Rewriting Kedro’s Spark dataset

If you run Kedro on Databricks, you run SparkDataset. It is the piece of the ecosystem with the least room to be wrong and the most accumulated history: years of path-handling special cases for DBFS, Unity Catalog, workspace files and object storage, each added for a good reason, none of them aware of the others.

Two things forced the issue. Spark Connect decouples the client from the cluster, and the old implementation assumed a local SparkSession it could reach into. And the path logic had reached the point where fixing one environment reliably broke another.

Approach

A rewrite of a critical dataset is only acceptable if nobody is forced to take it. So SparkDatasetV2 ships next to the original rather than replacing it. Users opt in, the old class keeps working, and the migration happens on their schedule instead of mine.

The work split into three pieces:

• The dataset itself. A clean implementation covering versioned reads and writes, schema loading from an external file, and the full set of Spark formats, at around 380 lines against a 767-line test suite.
• Shared Spark utilities. Session resolution that understands both a local session and a Spark Connect remote, extracted so other datasets stop reimplementing it.
• Shared Databricks utilities. One place that knows how to interpret a path, whether it arrives as DBFS, a workspace path, a Unity Catalog volume or a cloud URI.

Pulling the last two out into _utils was the part that mattered most for the long term. The path logic had been duplicated across several datasets with small divergences; now there is a single implementation with tests, and the follow-up work to make SparkHiveDataset.exists() Spark Connect-compatible was a small change rather than another copy of the same code.

Why it is more than a port

Under Spark Connect the client has no JVM to reach into. A surprising amount of Spark integration code quietly depends on that: anything touching sparkContext, Hadoop configuration or the JVM gateway. Making the dataset work in both worlds meant establishing which capabilities are actually available in each and routing around the gaps, rather than assuming a session and finding out at runtime.

Outcome

Thirteen files, close to two thousand lines, and no breaking change for existing users. I presented it on a Kedro Coffee Chat, at that point as a speaker rather than the host. I took on hosting afterwards.