Spark – Unity Catalog¶

def setup(self) -> None:
    """Obtain the active SparkSession."""
    from pyspark.sql import SparkSession  # noqa: PLC0415

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

def write(self, context: OutputContext, obj: Any) -> None:
    """Write a PySpark DataFrame or `DeltaMerge`.

    - If *obj* is a `DeltaMerge`, builds and executes a merge.
    - 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.
    """
    from databricks_bundle_decorators.merge import DeltaMerge  # noqa: PLC0415

    loc = self._location_path(context.asset_name)

    if isinstance(obj, DeltaMerge):
        table = self._table_name(context.asset_name)
        _logger.info(
            "Merging into %s (predicate=%r, actions=%s)",
            table,
            obj.predicate,
            obj._describe_actions(),
        )
        builder = obj._build_spark_merger(table)
        if builder is None:
            write_opts = dict(self._write_options)
            if loc:
                write_opts["path"] = loc
            obj._initial_spark_write(
                table,
                partition_by=context.partition_by,
                write_options=write_opts,
            )
        else:
            builder.execute()
        if context.partition_by:
            self._last_partition_values = _spark_extract_partition_values(
                obj.source, context.partition_by
            )
        else:
            self._last_partition_values = {}
        return

    partition_by = context.partition_by

    # Inject backfill_key column if it's a partition column
    if _should_inject_backfill_key(
        partition_by, has_backfill_key_col="backfill_key" in obj.columns
    ):
        from pyspark.sql import functions as F  # noqa: N812, PLC0415

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

    table = self._table_name(context.asset_name)
    _logger.info(
        "Writing to %s (mode=%s, partition_by=%s, location=%s)",
        table,
        self._mode,
        partition_by,
        loc,
    )
    writer = obj.write.format("delta").mode(self._mode)
    if loc:
        writer = writer.option("path", loc)
    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)

def read(self, context: InputContext) -> Any:
    """Read a Unity Catalog table as a PySpark DataFrame.

    Works for both managed and external tables — the read path
    uses ``spark.table()`` which resolves via the UC catalog
    regardless of storage location.

    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_asset_name)
    _logger.info(
        "Reading from %s (partition_filter=%s)", table, context.partition_filter
    )
    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  # noqa: N812, PLC0415

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

    return result

def setup(self) -> None:
    """Obtain the active SparkSession."""
    from pyspark.sql import SparkSession  # noqa: PLC0415

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

def write(self, context: OutputContext, obj: Any) -> None:
    """Write a PySpark DataFrame or `DeltaMerge`.

    - If *obj* is a `DeltaMerge`, builds and executes a merge.
    - 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.
    """
    from databricks_bundle_decorators.merge import DeltaMerge  # noqa: PLC0415

    if isinstance(obj, DeltaMerge):
        uri = self._uri(context.asset_name)
        _logger.info(
            "Merging into %s (predicate=%r, actions=%s)",
            uri,
            obj.predicate,
            obj._describe_actions(),
        )
        builder = obj._build_spark_merger(uri)
        if builder is None:
            obj._initial_spark_write(
                uri,
                partition_by=context.partition_by,
                write_options=dict(self._write_options),
            )
        else:
            builder.execute()
        if context.partition_by:
            self._last_partition_values = _spark_extract_partition_values(
                obj.source, context.partition_by
            )
        else:
            self._last_partition_values = {}
        return

    partition_by = context.partition_by

    # Inject backfill_key column if it's a partition column
    if _should_inject_backfill_key(
        partition_by, has_backfill_key_col="backfill_key" in obj.columns
    ):
        from pyspark.sql import functions as F  # noqa: N812, PLC0415

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

    uri = self._uri(context.asset_name)
    _logger.info(
        "Writing to %s (mode=%s, partition_by=%s)", uri, self._mode, partition_by
    )
    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)

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_asset_name)
    _logger.info(
        "Reading from %s (partition_filter=%s)", uri, context.partition_filter
    )
    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  # noqa: N812, PLC0415

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

    return result

def setup(self) -> None:
    """Obtain the active SparkSession."""
    from pyspark.sql import SparkSession  # noqa: PLC0415

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

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 _should_inject_backfill_key(
        partition_by, has_backfill_key_col="backfill_key" in obj.columns
    ):
        from pyspark.sql import functions as F  # noqa: N812, PLC0415

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

    uri = self._uri(context.asset_name)
    _logger.info("Writing to %s (partition_by=%s)", uri, partition_by)
    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)

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_asset_name)
    _logger.info(
        "Reading from %s (partition_filter=%s)", uri, context.partition_filter
    )
    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  # noqa: N812, PLC0415

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

    return result