Polars¶

def write(self, context: OutputContext, obj: Any) -> None:
    """Write a Polars DataFrame or LazyFrame to Parquet.

    - `polars.DataFrame` → ``write_parquet`` (native ``partition_by``)
    - `polars.LazyFrame` → ``sink_parquet`` (``pl.PartitionBy``)

    When ``partition_by`` is set on the ``@task`` decorator, writes
    to Hive-style partitioned directories.
    """
    import polars as pl  # noqa: PLC0415

    base_uri = self._uri(context.asset_name)
    partition_by = context.partition_by
    _logger.info("Writing to %s (partition_by=%s)", base_uri, partition_by)

    # Inject backfill_key column if it's a partition column
    has_bk_col = isinstance(
        obj, (pl.DataFrame, pl.LazyFrame)
    ) and "backfill_key" in (
        obj.collect_schema().names()
        if isinstance(obj, pl.LazyFrame)
        else obj.columns
    )
    if _should_inject_backfill_key(partition_by, has_backfill_key_col=has_bk_col):
        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.with_columns(pl.lit(bk).alias("backfill_key"))

    if partition_by:
        # Extract partition values from data before writing
        self._last_partition_values = _polars_extract_partition_values(
            obj, partition_by
        )
        if isinstance(obj, pl.LazyFrame):
            obj.sink_parquet(
                pl.PartitionBy(base_uri, key=partition_by),
                mkdir=True,
                storage_options=self.storage_options,
                **self._write_options,
            )
        elif isinstance(obj, pl.DataFrame):
            obj.write_parquet(
                base_uri,
                partition_by=partition_by,
                mkdir=True,
                storage_options=self.storage_options,
                **self._write_options,
            )
        else:
            msg = (
                f"PolarsParquetIoManager.write() expects a polars.DataFrame or "
                f"polars.LazyFrame, got {type(obj).__name__}"
            )
            raise TypeError(msg)
    else:
        uri = f"{base_uri}.parquet"
        if isinstance(obj, pl.LazyFrame):
            obj.sink_parquet(
                uri, storage_options=self.storage_options, **self._write_options
            )
        elif isinstance(obj, pl.DataFrame):
            obj.write_parquet(
                uri, storage_options=self.storage_options, **self._write_options
            )
        else:
            msg = (
                f"PolarsParquetIoManager.write() expects a polars.DataFrame or "
                f"polars.LazyFrame, got {type(obj).__name__}"
            )
            raise TypeError(msg)

def read(self, context: InputContext) -> Any:
    """Read Parquet as a LazyFrame or DataFrame.

    If the downstream parameter is annotated as `polars.DataFrame`,
    returns ``read_parquet`` (eager).  Otherwise returns ``scan_parquet``
    (lazy `polars.LazyFrame`) — this is the default for
    unannotated parameters.

    When ``partition_by`` is set on the producing ``@task``, reads
    from the Hive-partitioned directory.  By default only the
    current ``backfill_key`` partition is returned; use
    `all_partitions()` on the upstream dependency or
    ``@task(all_partitions=True)`` on the consuming task to read
    all partitions.
    """
    import polars as pl  # noqa: PLC0415

    base_uri = self._uri(context.upstream_asset_name)
    partition_by = context.partition_by
    _logger.info(
        "Reading from %s (partition_filter=%s)", base_uri, context.partition_filter
    )

    if partition_by:
        glob_uri = f"{base_uri}/**/*.parquet"
        if context.expected_type is pl.DataFrame:
            result = pl.read_parquet(
                glob_uri,
                hive_partitioning=True,
                storage_options=self.storage_options,
                **self._read_options,
            )
        else:
            result = pl.scan_parquet(
                glob_uri,
                hive_partitioning=True,
                storage_options=self.storage_options,
                **self._read_options,
            )
        if context.partition_filter and not context.all_partitions:
            result = _polars_apply_partition_filter(
                result, context.partition_filter
            )
        elif (
            self.auto_filter
            and _needs_backfill_key_col(partition_by)
            and not context.all_partitions
        ):
            result = result.filter(
                pl.col("backfill_key")
                == _resolve_backfill_key(context.backfill_key)
            )
        return result

    uri = f"{base_uri}.parquet"
    if context.expected_type is pl.DataFrame:
        return pl.read_parquet(
            uri, storage_options=self.storage_options, **self._read_options
        )
    return pl.scan_parquet(
        uri, storage_options=self.storage_options, **self._read_options
    )

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

    - `polars.DataFrame` → ``write_delta``
    - `polars.LazyFrame` → ``sink_delta``
    - `DeltaMerge` → merge/upsert operation

    When ``partition_by`` is set on the ``@task`` decorator, writes
    with ``delta_write_options={"partition_by": ...}``.
    """
    # Handle DeltaMerge definitions — build a fresh merger and execute.
    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(),
        )
        merger = obj._build_merger(uri, storage_options=self.storage_options)
        if merger is None:
            # Target table doesn't exist yet — write source data directly.
            obj._initial_write(
                uri,
                storage_options=self.storage_options,
                partition_by=context.partition_by,
                write_options=dict(self._write_options),
            )
        else:
            merger.execute()
        if context.partition_by:
            self._last_partition_values = _polars_extract_partition_values(
                obj.source, context.partition_by
            )
        else:
            self._last_partition_values = {}
        return

    import polars as pl  # noqa: PLC0415

    uri = self._uri(context.asset_name)
    partition_by = context.partition_by

    # Inject backfill_key column if it's a partition column
    has_bk_col = isinstance(
        obj, (pl.DataFrame, pl.LazyFrame)
    ) and "backfill_key" in (
        obj.collect_schema().names()
        if isinstance(obj, pl.LazyFrame)
        else obj.columns
    )
    if _should_inject_backfill_key(partition_by, has_backfill_key_col=has_bk_col):
        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.with_columns(pl.lit(bk).alias("backfill_key"))

    # Merge partition_by into write_options for Delta
    write_opts = dict(self._write_options)
    if partition_by:
        delta_opts = write_opts.setdefault("delta_write_options", {})
        delta_opts.setdefault("partition_by", partition_by)

    # Extract partition values from data before writing
    if partition_by:
        self._last_partition_values = _polars_extract_partition_values(
            obj, partition_by
        )

    # Scope overwrite to affected partitions only
    if self._mode == "overwrite" and partition_by and self._last_partition_values:
        delta_opts = write_opts.setdefault("delta_write_options", {})
        delta_opts.setdefault(
            "predicate", _build_replace_where(self._last_partition_values)
        )

    _logger.info(
        "Writing to %s (mode=%s, partition_by=%s)", uri, self._mode, partition_by
    )

    if isinstance(obj, pl.LazyFrame):
        obj.sink_delta(
            uri,
            mode=self._mode,
            storage_options=self.storage_options,
            **write_opts,
        )
    elif isinstance(obj, pl.DataFrame):
        obj.write_delta(
            uri,
            mode=self._mode,
            storage_options=self.storage_options,
            **write_opts,
        )
    else:
        msg = (
            f"PolarsDeltaIoManager.write() expects a polars.DataFrame, "
            f"polars.LazyFrame, or DeltaMerge, "
            f"got {type(obj).__name__}"
        )
        raise TypeError(msg)

def write_with_retry(self, context: OutputContext, obj: Any) -> None:
    """Write with retry logic.

    `DeltaMerge` and DataFrame/LazyFrame writes are all retried
    when `RetryConfig` is configured.
    """
    # DeltaMerge — retry-safe path: rebuild merger on each attempt.
    if isinstance(obj, DeltaMerge):
        if self.retry is None:
            self.write(context, obj)
            return

        uri = self._uri(context.asset_name)

        def _execute_merge() -> None:
            merger = obj._build_merger(uri, storage_options=self.storage_options)
            if merger is None:
                obj._initial_write(
                    uri,
                    storage_options=self.storage_options,
                    partition_by=context.partition_by,
                    write_options=dict(self._write_options),
                )
            else:
                merger.execute()

        wait_kwargs: dict[str, Any] = {
            "multiplier": self.retry.delay,
            "exp_base": self.retry.backoff_factor,
        }
        if self.retry.max_delay is not None:
            wait_kwargs["max"] = self.retry.max_delay

        retryer = retry(
            stop=stop_after_attempt(self.retry.max_attempts),
            wait=wait_exponential(**wait_kwargs),
            reraise=True,
            before_sleep=before_sleep_log(_logger, logging.WARNING),
        )
        retryer(_execute_merge)()
        if context.partition_by:
            self._last_partition_values = _polars_extract_partition_values(
                obj.source, context.partition_by
            )
        else:
            self._last_partition_values = {}
        return

    super().write_with_retry(context, obj)

def read(self, context: InputContext) -> Any:
    """Read a Delta table as a LazyFrame or DataFrame.

    If the downstream parameter is annotated as `polars.DataFrame`,
    returns ``read_delta`` (eager).  Otherwise returns ``scan_delta``
    (lazy `polars.LazyFrame`) — this is the default for
    unannotated parameters.

    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)``.
    """
    import polars as pl  # noqa: PLC0415

    uri = self._uri(context.upstream_asset_name)
    _logger.info(
        "Reading from %s (partition_filter=%s)", uri, context.partition_filter
    )

    if context.expected_type is pl.DataFrame:
        result = pl.read_delta(
            uri, storage_options=self.storage_options, **self._read_options
        )
    else:
        result = pl.scan_delta(
            uri, storage_options=self.storage_options, **self._read_options
        )

    if context.partition_filter and not context.all_partitions:
        result = _polars_apply_partition_filter(result, context.partition_filter)
    elif (
        self.auto_filter
        and _needs_backfill_key_col(context.partition_by)
        and not context.all_partitions
    ):
        result = result.filter(
            pl.col("backfill_key") == _resolve_backfill_key(context.backfill_key)
        )

    return result

def write(self, context: OutputContext, obj: Any) -> None:
    """Write a Polars DataFrame or LazyFrame as NDJSON.

    - `polars.LazyFrame` → ``sink_ndjson``
    - `polars.DataFrame` → ``.lazy()`` then ``sink_ndjson``

    When ``partition_by`` is set on the ``@task`` decorator, writes
    to Hive-style partitioned directories using
    ``pl.PartitionBy``.
    """
    import polars as pl  # noqa: PLC0415

    base_uri = self._uri(context.asset_name)
    partition_by = context.partition_by
    _logger.info("Writing to %s (partition_by=%s)", base_uri, partition_by)

    # Inject backfill_key column if it's a partition column
    has_bk_col = isinstance(
        obj, (pl.DataFrame, pl.LazyFrame)
    ) and "backfill_key" in (
        obj.collect_schema().names()
        if isinstance(obj, pl.LazyFrame)
        else obj.columns
    )
    if _should_inject_backfill_key(partition_by, has_backfill_key_col=has_bk_col):
        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.with_columns(pl.lit(bk).alias("backfill_key"))

    if partition_by:
        # Extract partition values from data before writing
        self._last_partition_values = _polars_extract_partition_values(
            obj, partition_by
        )
        if isinstance(obj, pl.LazyFrame):
            obj.sink_ndjson(
                pl.PartitionBy(base_uri, key=partition_by),
                mkdir=True,
                storage_options=self.storage_options,
                **self._write_options,
            )
        elif isinstance(obj, pl.DataFrame):
            obj.lazy().sink_ndjson(
                pl.PartitionBy(base_uri, key=partition_by),
                mkdir=True,
                storage_options=self.storage_options,
                **self._write_options,
            )
        else:
            msg = (
                f"PolarsJsonIoManager.write() expects a polars.DataFrame or "
                f"polars.LazyFrame, got {type(obj).__name__}"
            )
            raise TypeError(msg)
    else:
        uri = f"{base_uri}.ndjson"
        if isinstance(obj, pl.LazyFrame):
            obj.sink_ndjson(
                uri, storage_options=self.storage_options, **self._write_options
            )
        elif isinstance(obj, pl.DataFrame):
            obj.lazy().sink_ndjson(
                uri, storage_options=self.storage_options, **self._write_options
            )
        else:
            msg = (
                f"PolarsJsonIoManager.write() expects a polars.DataFrame or "
                f"polars.LazyFrame, got {type(obj).__name__}"
            )
            raise TypeError(msg)

def read(self, context: InputContext) -> Any:
    """Read NDJSON as a LazyFrame or DataFrame.

    If the downstream parameter is annotated as `polars.DataFrame`,
    returns ``read_ndjson`` (eager).  Otherwise returns ``scan_ndjson``
    (lazy `polars.LazyFrame`) — this is the default for
    unannotated parameters.

    When ``partition_by`` is set, reads from the Hive-partitioned
    directory.  By default only the current ``backfill_key``
    partition is returned; use `all_partitions()` on the
    upstream dependency or ``@task(all_partitions=True)`` on
    the consuming task to read all partitions.
    """
    import polars as pl  # noqa: PLC0415

    base_uri = self._uri(context.upstream_asset_name)
    partition_by = context.partition_by
    _logger.info(
        "Reading from %s (partition_filter=%s)", base_uri, context.partition_filter
    )

    if partition_by:
        glob_uri = f"{base_uri}/**/*.jsonl"
        if context.expected_type is pl.DataFrame:
            result = pl.read_ndjson(
                glob_uri,
                storage_options=self.storage_options,
                **self._read_options,
            )
        else:
            result = pl.scan_ndjson(
                glob_uri,
                storage_options=self.storage_options,
                **self._read_options,
            )
        if context.partition_filter and not context.all_partitions:
            result = _polars_apply_partition_filter(
                result, context.partition_filter
            )
        elif (
            self.auto_filter
            and _needs_backfill_key_col(partition_by)
            and not context.all_partitions
        ):
            result = result.filter(
                pl.col("backfill_key")
                == _resolve_backfill_key(context.backfill_key)
            )
        return result

    uri = f"{base_uri}.ndjson"
    if context.expected_type is pl.DataFrame:
        return pl.read_ndjson(
            uri, storage_options=self.storage_options, **self._read_options
        )
    return pl.scan_ndjson(
        uri, storage_options=self.storage_options, **self._read_options
    )

def write(self, context: OutputContext, obj: Any) -> None:
    """Write a Polars DataFrame or LazyFrame as CSV.

    - `polars.LazyFrame` → ``sink_csv``
    - `polars.DataFrame` → ``write_csv`` (single file) or
      ``.lazy().sink_csv(PartitionBy)`` (partitioned)

    When ``partition_by`` is set on the ``@task`` decorator, writes
    to Hive-style partitioned directories using
    ``pl.PartitionBy``.
    """
    import polars as pl  # noqa: PLC0415

    base_uri = self._uri(context.asset_name)
    partition_by = context.partition_by
    _logger.info("Writing to %s (partition_by=%s)", base_uri, partition_by)

    # Inject backfill_key column if it's a partition column
    has_bk_col = isinstance(
        obj, (pl.DataFrame, pl.LazyFrame)
    ) and "backfill_key" in (
        obj.collect_schema().names()
        if isinstance(obj, pl.LazyFrame)
        else obj.columns
    )
    if _should_inject_backfill_key(partition_by, has_backfill_key_col=has_bk_col):
        bk = _resolve_backfill_key(context.backfill_key)
        obj = obj.with_columns(pl.lit(bk).alias("backfill_key"))

    if partition_by:
        # Extract partition values from data before writing
        self._last_partition_values = _polars_extract_partition_values(
            obj, partition_by
        )
        if isinstance(obj, pl.LazyFrame):
            obj.sink_csv(
                pl.PartitionBy(base_uri, key=partition_by),
                mkdir=True,
                storage_options=self.storage_options,
                **self._write_options,
            )
        elif isinstance(obj, pl.DataFrame):
            obj.lazy().sink_csv(
                pl.PartitionBy(base_uri, key=partition_by),
                mkdir=True,
                storage_options=self.storage_options,
                **self._write_options,
            )
        else:
            msg = (
                f"PolarsCsvIoManager.write() expects a polars.DataFrame or "
                f"polars.LazyFrame, got {type(obj).__name__}"
            )
            raise TypeError(msg)
    else:
        uri = f"{base_uri}.csv"
        if isinstance(obj, pl.LazyFrame):
            obj.sink_csv(
                uri, storage_options=self.storage_options, **self._write_options
            )
        elif isinstance(obj, pl.DataFrame):
            obj.write_csv(
                uri, storage_options=self.storage_options, **self._write_options
            )
        else:
            msg = (
                f"PolarsCsvIoManager.write() expects a polars.DataFrame or "
                f"polars.LazyFrame, got {type(obj).__name__}"
            )
            raise TypeError(msg)

def read(self, context: InputContext) -> Any:
    """Read CSV as a LazyFrame or DataFrame.

    If the downstream parameter is annotated as `polars.DataFrame`,
    returns ``read_csv`` (eager).  Otherwise returns ``scan_csv``
    (lazy `polars.LazyFrame`) — this is the default for
    unannotated parameters.

    When ``partition_by`` is set, reads from the Hive-partitioned
    directory.  By default only the current ``backfill_key``
    partition is returned; use `all_partitions()` on the
    upstream dependency or ``@task(all_partitions=True)`` on
    the consuming task to read all partitions.
    """
    import polars as pl  # noqa: PLC0415

    base_uri = self._uri(context.upstream_asset_name)
    partition_by = context.partition_by
    _logger.info(
        "Reading from %s (partition_filter=%s)", base_uri, context.partition_filter
    )

    if partition_by:
        glob_uri = f"{base_uri}/**/*.csv"
        if context.expected_type is pl.DataFrame:
            result = pl.read_csv(
                glob_uri,
                storage_options=self.storage_options,
                **self._read_options,
            )
        else:
            result = pl.scan_csv(
                glob_uri,
                storage_options=self.storage_options,
                **self._read_options,
            )
        if context.partition_filter and not context.all_partitions:
            result = _polars_apply_partition_filter(
                result, context.partition_filter
            )
        elif (
            self.auto_filter
            and _needs_backfill_key_col(partition_by)
            and not context.all_partitions
        ):
            result = result.filter(
                pl.col("backfill_key")
                == _resolve_backfill_key(context.backfill_key)
            )
        return result

    uri = f"{base_uri}.csv"
    if context.expected_type is pl.DataFrame:
        return pl.read_csv(
            uri, storage_options=self.storage_options, **self._read_options
        )
    return pl.scan_csv(
        uri, storage_options=self.storage_options, **self._read_options
    )