Skip to content

Spark – Unity Catalog

Unity Catalog IoManagers work on both classic and serverless compute. UC handles authentication and access control, so no credential injection is needed.

Partitioning

All UC IoManagers support partition_by via the @task decorator. "backfill_key" is auto-injected on write and auto-filtered on read. Managed tables use partitionBy() with saveAsTable(); volume paths use partitionBy() with save().

io = SparkUCTableIoManager(
    catalog="main",
    schema="staging",
    mode="overwrite",
)

@task(io_manager=io, partition_by="backfill_key")
def extract(): ...

Managed Tables

databricks_bundle_decorators.io_managers.SparkUCTableIoManager(catalog, schema, write_options=None, read_options=None, mode='error', *, auto_filter=True)

Bases: IoManager

Persist PySpark DataFrames as Unity Catalog managed Delta tables.

Uses saveAsTable / spark.table() with the three-level namespace catalog.schema.task_key.

Unity Catalog manages access control and storage location, so no credential configuration is required. Works on both classic and serverless compute.

Parameters:

Name Type Description Default
catalog str

Unity Catalog catalog name (e.g. "main").

 required
schema str

Unity Catalog schema (database) name (e.g. "staging").

 required
write_options dict[str, str] | None

Extra Spark writer options applied via .option(k, v).

 None
read_options dict[str, str] | None

Extra Spark reader options applied via .option(k, v).

 None
mode str

Delta write mode ("error", "overwrite", "append", etc.). Defaults to "error" to prevent accidental data loss. For merge operations, return a DeltaMergeBuilder from your task instead.

 'error'
Example

::

from databricks_bundle_decorators.io_managers import SparkUCTableIoManager

io = SparkUCTableIoManager(catalog="main", schema="staging")

@task(io_manager=io)
def extract():
    spark = SparkSession.getActiveSession()
    return spark.range(10)

@task
def transform(df):   # spark.table("main.staging.extract")
    df.show()
Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def __init__(
    self,
    catalog: str,
    schema: str,
    write_options: dict[str, str] | None = None,
    read_options: dict[str, str] | None = None,
    mode: str = "error",
    *,
    auto_filter: bool = True,
) -> None:
    self.catalog = catalog
    self.schema = schema
    self._write_options = write_options or {}
    self._read_options = read_options or {}
    self._mode = mode
    self.auto_filter = auto_filter

setup()

Obtain the active SparkSession.

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
102
103
104
105
106
107
108
109
def setup(self) -> None:
    """Obtain the active SparkSession."""
    from pyspark.sql import SparkSession

    self._spark = SparkSession.getActiveSession()
    if self._spark is None:
        msg = "No active SparkSession found."
        raise RuntimeError(msg)

write(context, obj)

Write a PySpark DataFrame or execute a DeltaMergeBuilder.

  • If obj is a DeltaMergeBuilder, calls .execute().
  • Otherwise writes via saveAsTable with the configured mode, partition_by, and write_options.

When partition_by includes "backfill_key", the column is injected automatically from the context.

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def write(self, context: OutputContext, obj: Any) -> None:
    """Write a PySpark DataFrame or execute a DeltaMergeBuilder.

    - If *obj* is a ``DeltaMergeBuilder``, calls ``.execute()``.
    - Otherwise writes via ``saveAsTable`` with the configured
      ``mode``, ``partition_by``, and ``write_options``.

    When ``partition_by`` includes ``"backfill_key"``, the column
    is injected automatically from the context.
    """
    _merge_cls: type | None = None
    try:
        from delta.tables import DeltaMergeBuilder

        _merge_cls = DeltaMergeBuilder
    except ImportError:
        pass

    if _merge_cls is not None and isinstance(obj, _merge_cls):
        obj.execute()
        return

    partition_by = context.partition_by

    # Inject backfill_key column if it's a partition column
    if _needs_backfill_key_col(partition_by):
        from pyspark.sql import functions as F

        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.withColumn("backfill_key", F.lit(bk))

    table = self._table_name(context.task_key)
    writer = obj.write.format("delta").mode(self._mode)
    if partition_by:
        # Extract partition values from data before writing
        self._last_partition_values = _spark_extract_partition_values(
            obj, partition_by
        )
        writer = writer.partitionBy(*partition_by)
        # Scope overwrite to affected partitions only (skip on first write
        # because saveAsTable + replaceWhere requires an existing table)
        if self._mode == "overwrite" and self._spark.catalog.tableExists(table):
            writer = writer.option(
                "replaceWhere",
                _build_replace_where(self._last_partition_values),
            )
    for k, v in self._write_options.items():
        writer = writer.option(k, v)
    writer.saveAsTable(table)

read(context)

Read a Unity Catalog managed table as a PySpark DataFrame.

When partition_by includes "backfill_key", reads are filtered to the current partition unless the upstream dependency uses all_partitions() or the consuming task uses @task(all_partitions=True).

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def read(self, context: InputContext) -> Any:
    """Read a Unity Catalog managed table as a PySpark DataFrame.

    When ``partition_by`` includes ``"backfill_key"``, reads are
    filtered to the current partition unless the upstream
    dependency uses `all_partitions()` or the consuming
    task uses ``@task(all_partitions=True)``.
    """
    table = self._table_name(context.upstream_task_key)
    result = self._spark.table(table)

    if context.partition_filter and not context.all_partitions:
        result = _spark_apply_partition_filter(result, context.partition_filter)
    elif (
        self.auto_filter
        and _needs_backfill_key_col(context.partition_by)
        and not context.all_partitions
    ):
        from pyspark.sql import functions as F

        result = result.filter(
            F.col("backfill_key") == _resolve_backfill_key(context.backfill_key)
        )

    return result

Volume – Delta

databricks_bundle_decorators.io_managers.SparkUCVolumeDeltaIoManager(catalog, schema, volume, write_options=None, read_options=None, mode='error', *, auto_filter=True)

Bases: IoManager

Persist PySpark DataFrames as Delta tables in UC Volumes.

Writes to /Volumes/<catalog>/<schema>/<volume>/<task_key> using the standard Delta format.

Parameters:

Name Type Description Default
catalog str

Unity Catalog catalog name.

 required
schema str

Unity Catalog schema (database) name.

 required
volume str

Unity Catalog volume name.

 required
write_options dict[str, str] | None

Extra Spark writer options applied via .option(k, v).

 None
read_options dict[str, str] | None

Extra Spark reader options applied via .option(k, v).

 None
mode str

Delta write mode ("error", "overwrite", "append", etc.). Defaults to "error" to prevent accidental data loss. For merge operations, return a DeltaMergeBuilder from your task instead.

 'error'
Example

::

from databricks_bundle_decorators.io_managers import (
    SparkUCVolumeDeltaIoManager,
)

io = SparkUCVolumeDeltaIoManager(
    catalog="main", schema="staging", volume="raw_data",
)

@task(io_manager=io)
def extract():
    spark = SparkSession.getActiveSession()
    return spark.range(10)
Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def __init__(
    self,
    catalog: str,
    schema: str,
    volume: str,
    write_options: dict[str, str] | None = None,
    read_options: dict[str, str] | None = None,
    mode: str = "error",
    *,
    auto_filter: bool = True,
) -> None:
    self.catalog = catalog
    self.schema = schema
    self.volume = volume
    self._write_options = write_options or {}
    self._read_options = read_options or {}
    self._mode = mode
    self.auto_filter = auto_filter

setup()

Obtain the active SparkSession.

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
254
255
256
257
258
259
260
261
def setup(self) -> None:
    """Obtain the active SparkSession."""
    from pyspark.sql import SparkSession

    self._spark = SparkSession.getActiveSession()
    if self._spark is None:
        msg = "No active SparkSession found."
        raise RuntimeError(msg)

write(context, obj)

Write a PySpark DataFrame or execute a DeltaMergeBuilder.

  • If obj is a DeltaMergeBuilder, calls .execute().
  • Otherwise writes via save() with the configured mode, partition_by, and write_options.

When partition_by includes "backfill_key", the column is injected automatically from the context.

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
def write(self, context: OutputContext, obj: Any) -> None:
    """Write a PySpark DataFrame or execute a DeltaMergeBuilder.

    - If *obj* is a ``DeltaMergeBuilder``, calls ``.execute()``.
    - Otherwise writes via ``save()`` with the configured
      ``mode``, ``partition_by``, and ``write_options``.

    When ``partition_by`` includes ``"backfill_key"``, the column
    is injected automatically from the context.
    """
    _merge_cls: type | None = None
    try:
        from delta.tables import DeltaMergeBuilder

        _merge_cls = DeltaMergeBuilder
    except ImportError:
        pass

    if _merge_cls is not None and isinstance(obj, _merge_cls):
        obj.execute()
        return

    partition_by = context.partition_by

    # Inject backfill_key column if it's a partition column
    if _needs_backfill_key_col(partition_by):
        from pyspark.sql import functions as F

        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.withColumn("backfill_key", F.lit(bk))

    uri = self._uri(context.task_key)
    writer = obj.write.format("delta").mode(self._mode)
    if partition_by:
        # Extract partition values from data before writing
        self._last_partition_values = _spark_extract_partition_values(
            obj, partition_by
        )
        writer = writer.partitionBy(*partition_by)
        # Scope overwrite to affected partitions only
        if self._mode == "overwrite":
            writer = writer.option(
                "replaceWhere",
                _build_replace_where(self._last_partition_values),
            )
    for k, v in self._write_options.items():
        writer = writer.option(k, v)
    writer.save(uri)

read(context)

Read Delta from a UC Volume path as a PySpark DataFrame.

When partition_by includes "backfill_key", reads are filtered to the current partition unless the upstream dependency uses all_partitions() or the consuming task uses @task(all_partitions=True).

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def read(self, context: InputContext) -> Any:
    """Read Delta from a UC Volume path as a PySpark DataFrame.

    When ``partition_by`` includes ``"backfill_key"``, reads are
    filtered to the current partition unless the upstream
    dependency uses `all_partitions()` or the consuming
    task uses ``@task(all_partitions=True)``.
    """
    uri = self._uri(context.upstream_task_key)
    reader = self._spark.read.format("delta")
    for k, v in self._read_options.items():
        reader = reader.option(k, v)
    result = reader.load(uri)

    if context.partition_filter and not context.all_partitions:
        result = _spark_apply_partition_filter(result, context.partition_filter)
    elif (
        self.auto_filter
        and _needs_backfill_key_col(context.partition_by)
        and not context.all_partitions
    ):
        from pyspark.sql import functions as F

        result = result.filter(
            F.col("backfill_key") == _resolve_backfill_key(context.backfill_key)
        )

    return result

Volume – Parquet

databricks_bundle_decorators.io_managers.SparkUCVolumeParquetIoManager(catalog, schema, volume, write_options=None, read_options=None, *, auto_filter=True)

Bases: IoManager

Persist PySpark DataFrames as Parquet in UC Volumes.

Writes to /Volumes/<catalog>/<schema>/<volume>/<task_key>.parquet using the Parquet format.

Parameters:

Name Type Description Default
catalog str

Unity Catalog catalog name.

 required
schema str

Unity Catalog schema (database) name.

 required
volume str

Unity Catalog volume name.

 required
write_options dict[str, str] | None

Extra Spark writer options applied via .option(k, v).

 None
read_options dict[str, str] | None

Extra Spark reader options applied via .option(k, v).

 None
Example

::

from databricks_bundle_decorators.io_managers import (
    SparkUCVolumeParquetIoManager,
)

io = SparkUCVolumeParquetIoManager(
    catalog="main", schema="staging", volume="raw_data",
)

@task(io_manager=io)
def extract():
    spark = SparkSession.getActiveSession()
    return spark.range(10)
Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
def __init__(
    self,
    catalog: str,
    schema: str,
    volume: str,
    write_options: dict[str, str] | None = None,
    read_options: dict[str, str] | None = None,
    *,
    auto_filter: bool = True,
) -> None:
    self.catalog = catalog
    self.schema = schema
    self.volume = volume
    self._write_options = write_options or {}
    self._read_options = read_options or {}
    self.auto_filter = auto_filter

setup()

Obtain the active SparkSession.

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
401
402
403
404
405
406
407
408
def setup(self) -> None:
    """Obtain the active SparkSession."""
    from pyspark.sql import SparkSession

    self._spark = SparkSession.getActiveSession()
    if self._spark is None:
        msg = "No active SparkSession found."
        raise RuntimeError(msg)

write(context, obj)

Write a PySpark DataFrame as Parquet to a UC Volume path.

When partition_by includes "backfill_key", the column is injected automatically from the context.

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
def write(self, context: OutputContext, obj: Any) -> None:
    """Write a PySpark DataFrame as Parquet to a UC Volume path.

    When ``partition_by`` includes ``"backfill_key"``, the column
    is injected automatically from the context.
    """
    partition_by = context.partition_by

    # Inject backfill_key column if it's a partition column
    if _needs_backfill_key_col(partition_by):
        from pyspark.sql import functions as F

        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.withColumn("backfill_key", F.lit(bk))

    uri = self._uri(context.task_key)
    writer = obj.write.format("parquet").mode("overwrite")
    if partition_by:
        # Extract partition values from data before writing
        self._last_partition_values = _spark_extract_partition_values(
            obj, partition_by
        )
        writer = writer.partitionBy(*partition_by)
        # Only overwrite partitions present in the data
        writer = writer.option("partitionOverwriteMode", "dynamic")
    for k, v in self._write_options.items():
        writer = writer.option(k, v)
    writer.save(uri)

read(context)

Read Parquet from a UC Volume path as a PySpark DataFrame.

When partition_by includes "backfill_key", reads are filtered to the current partition unless the upstream dependency uses all_partitions() or the consuming task uses @task(all_partitions=True).

Source code in src/databricks_bundle_decorators/io_managers/spark_uc.py 
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
def read(self, context: InputContext) -> Any:
    """Read Parquet from a UC Volume path as a PySpark DataFrame.

    When ``partition_by`` includes ``"backfill_key"``, reads are
    filtered to the current partition unless the upstream
    dependency uses `all_partitions()` or the consuming
    task uses ``@task(all_partitions=True)``.
    """
    uri = self._uri(context.upstream_task_key)
    reader = self._spark.read.format("parquet")
    for k, v in self._read_options.items():
        reader = reader.option(k, v)
    result = reader.load(uri)

    if context.partition_filter and not context.all_partitions:
        result = _spark_apply_partition_filter(result, context.partition_filter)
    elif (
        self.auto_filter
        and _needs_backfill_key_col(context.partition_by)
        and not context.all_partitions
    ):
        from pyspark.sql import functions as F

        result = result.filter(
            F.col("backfill_key") == _resolve_backfill_key(context.backfill_key)
        )

    return result