loom-kernel 0.10.0__tar.gz → 0.10.2__tar.gz

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/CHANGELOG.md

@@ -1,3 +1,37 @@
+# 🚀 Release 0.10.2 ([#46](https://github.com/the-reacher-data/loom-py/pull/46)) ([`0fa17f5`](https://github.com/the-reacher-data/loom-py/commit/0fa17f5735be843e68ba446993bb41f2d107c33b))
+
+
+## ✨ Features
+### polars
+- **polars:** streaming Arrow writes for replace_partitions
+
+
+
+
+
+
+
+
+
+
+
+# 🚀 Release 0.10.1 ([#44](https://github.com/the-reacher-data/loom-py/pull/44)) ([`a55cd27`](https://github.com/the-reacher-data/loom-py/commit/a55cd27020dee60aea6849b510500e7e6f5a83b2))
+
+
+
+## 🐛 Fixes
+### testing
+- **testing:** implement read_streaming on _PolarsStubReader (#44)
+
+
+
+
+
+
+
+
+
+
 # 🚀 Release 0.10.0 ([#41](https://github.com/the-reacher-data/loom-py/pull/41)) ([`099de38`](https://github.com/the-reacher-data/loom-py/commit/099de38b47cd7002a9c36fc56c9c85f673431022))
 
 

loom_kernel-0.10.2/CHANGELOG_RELEASE.md

@@ -0,0 +1,15 @@
+# 🚀 Release 0.10.2 ([#46](https://github.com/the-reacher-data/loom-py/pull/46)) ([`0fa17f5`](https://github.com/the-reacher-data/loom-py/commit/0fa17f5735be843e68ba446993bb41f2d107c33b))
+
+
+## ✨ Features
+### polars
+- **polars:** streaming Arrow writes for replace_partitions
+
+
+
+
+
+
+
+
+

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: loom-kernel
-Version: 0.10.0
+Version: 0.10.2
 Summary: Loom Python project
 License-File: LICENSE
 Classifier: Typing :: Typed

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "loom-kernel"
-version = "0.10.0"
+version = "0.10.2"
 description = "Loom Python project"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -268,7 +268,7 @@ markers = [
 
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.10.0"
+version = "0.10.2"
 tag_format = "v$version"
 version_files = [
     "pyproject.toml:project.version",

loom_kernel-0.10.2/src/loom/etl/backends/polars/_streaming.py

@@ -0,0 +1,82 @@
+"""Streaming-write helpers for the Polars Delta backend.
+
+Bridges a ``pl.LazyFrame`` to an Arrow ``RecordBatchReader`` via an
+intermediate IPC spool on disk.  Memory is bounded by Polars' streaming
+engine plus one batch in flight.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import NamedTuple
+
+import polars as pl
+import pyarrow as pa
+import pyarrow.ipc as ipc
+
+_SPOOL_DIR_ENV = "LOOM_SPOOL_DIR"
+
+
+class StreamingSpool(NamedTuple):
+    reader: pa.RecordBatchReader
+    spool_path: Path
+
+
+def _resolve_spool_dir(spool_dir: str | os.PathLike[str] | None) -> str | None:
+    if spool_dir is not None:
+        return str(spool_dir)
+    env_dir = os.environ.get(_SPOOL_DIR_ENV)
+    return env_dir if env_dir else None
+
+
+@contextmanager
+def spool_lazy_to_arrow_reader(
+    frame: pl.LazyFrame,
+    *,
+    spool_dir: str | os.PathLike[str] | None = None,
+) -> Iterator[StreamingSpool]:
+    """Sink *frame* to an IPC file and yield an Arrow ``RecordBatchReader``.
+
+    The spool file lives in a temporary directory removed on context exit
+    even when the consumer raises.  The yielded reader is single-pass;
+    callers needing a separate scan (e.g. for partition pruning) should
+    ``pl.scan_ipc(spool_path)`` while the context is active.
+
+    Args:
+        frame:     Lazy frame to materialise to a streaming Arrow source.
+        spool_dir: Base directory for the temporary spool.  When ``None``
+                   defaults to ``$LOOM_SPOOL_DIR`` if set, else the
+                   system temp dir.  Set this on container workloads
+                   where ``/tmp`` is a ``tmpfs`` mount (e.g. Fargate)
+                   to avoid spilling back to RAM.
+
+    Yields:
+        A :class:`StreamingSpool` with the reader and the spool path.
+    """
+    base_dir = _resolve_spool_dir(spool_dir)
+    with tempfile.TemporaryDirectory(prefix="loom-spool-", dir=base_dir) as tmp_dir:
+        spool_path = Path(tmp_dir) / "spool.arrow"
+        frame.sink_ipc(spool_path, compression="lz4")
+        source = ipc.open_file(spool_path)
+
+        def _iter_batches() -> Iterator[pa.RecordBatch]:
+            for index in range(source.num_record_batches):
+                yield source.get_batch(index)
+
+        reader = pa.RecordBatchReader.from_batches(source.schema, _iter_batches())
+        yield StreamingSpool(reader=reader, spool_path=spool_path)
+
+
+def partition_rows_from_spool(
+    spool_path: Path,
+    partition_cols: tuple[str, ...],
+) -> list[dict[str, object]]:
+    """Return distinct partition combinations from the spool file."""
+    frame = (
+        pl.scan_ipc(spool_path).select(list(partition_cols)).unique().collect(engine="streaming")
+    )
+    return list(frame.iter_rows(named=True))

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/src/loom/etl/backends/polars/_writer.py

@@ -29,7 +29,11 @@ from loom.etl.backends.polars._schema import (
 from loom.etl.declarative.target import SchemaMode
 from loom.etl.declarative.target._file import FileSpec
 from loom.etl.declarative.target._history import HistorifyRepairReport, HistorifySpec
-from loom.etl.declarative.target._table import AppendSpec, UpsertSpec
+from loom.etl.declarative.target._table import (
+    AppendSpec,
+    ReplacePartitionsSpec,
+    UpsertSpec,
+)
 from loom.etl.lineage._records import LineageRecord, WriteContext, get_lineage_schema
 from loom.etl.schema._schema import SchemaError
 from loom.etl.storage import (
@@ -46,6 +50,7 @@ from loom.etl.storage._locator import TableLocator, _as_locator
 
 from ._dtype import loom_type_to_polars
 from ._file_writer import PolarsFileWriter
+from ._streaming import partition_rows_from_spool, spool_lazy_to_arrow_reader
 
 _log = get_logger(__name__)
 
@@ -88,6 +93,27 @@ def _log_write_commit(mode: str, location: TableLocation) -> None:
         _log.info("delta write commit", mode=mode, uri=location.uri, metrics="unavailable")
 
 
+def _check_null_dtype_columns_lazy(frame: pl.LazyFrame) -> None:
+    """Lazy variant of :func:`_check_null_dtype_columns`.
+
+    Raises :exc:`~loom.etl.schema.SchemaError` if any column in the lazy
+    schema has dtype ``Null``.  Operates on ``collect_schema()`` (metadata
+    only), so it is safe to call on a frame about to be streamed.
+    """
+    schema = frame.collect_schema()
+    null_cols = [name for name, dtype in schema.items() if isinstance(dtype, pl.Null)]
+    if not null_cols:
+        return
+    schema_repr = "\n".join(f"  {n}: {d}" for n, d in schema.items())
+    raise SchemaError(
+        f"Delta Lake does not support the Null dtype. "
+        f"Column(s) with Null dtype: {null_cols}.\n"
+        f"Cast each column to its intended type before writing "
+        f"(use an explicit schema declaration or .cast() in the pipeline step).\n"
+        f"Full frame schema:\n{schema_repr}"
+    )
+
+
 def _check_null_dtype_columns(frame: pl.DataFrame) -> None:
     """Raise :exc:`~loom.etl.schema.SchemaError` if any column has dtype ``Null``.
 
@@ -189,6 +215,74 @@ class PolarsTargetWriter(_WritePolicy[pl.LazyFrame, pl.DataFrame, PolarsPhysical
         _check_null_dtype_columns(df)
         return df
 
+    def _do_replace_partitions(
+        self,
+        frame: pl.LazyFrame,
+        target: ResolvedTarget,
+        spec: ReplacePartitionsSpec,
+        streaming: bool,
+    ) -> None:
+        if not spec.streaming:
+            super()._do_replace_partitions(frame, target, spec, streaming)
+            return
+
+        existing = self._physical_schema(target)
+        if existing is None:
+            super()._do_replace_partitions(frame, target, spec, streaming)
+            return
+
+        aligned = self._align(frame, existing, spec.schema_mode)
+        _check_null_dtype_columns_lazy(aligned)
+        self._streaming_replace_partitions(aligned, target, spec)
+
+    def _streaming_replace_partitions(
+        self,
+        frame: pl.LazyFrame,
+        target: ResolvedTarget,
+        spec: ReplacePartitionsSpec,
+    ) -> None:
+        path_target = self._as_path_target(target)
+        location = path_target.location
+        with spool_lazy_to_arrow_reader(frame) as spool:
+            partition_rows = partition_rows_from_spool(spool.spool_path, spec.partition_cols)
+            if not partition_rows:
+                _log.warning("replace_partitions empty frame", table=location.uri)
+                return
+            null_rows = [r for r in partition_rows if any(v is None for v in r.values())]
+            if null_rows:
+                sample = null_rows[:5]
+                more = (
+                    f" (+{len(null_rows) - len(sample)} more)"
+                    if len(null_rows) > len(sample)
+                    else ""
+                )
+                raise ValueError(
+                    f"replace_partitions: null values in partition columns "
+                    f"{spec.partition_cols} for table {path_target}. Null partition "
+                    f"values cannot be used as a replace predicate. Ensure all "
+                    f"partition columns are non-null before writing. "
+                    f"Offending rows: {sample}{more}"
+                )
+            predicate = _build_partition_predicate(iter(partition_rows), spec.partition_cols)
+            _log.info(
+                "delta write start",
+                mode="replace_partitions",
+                uri=location.uri,
+                rows="streaming",
+                cols=None,
+                partition_cols=list(spec.partition_cols),
+                predicate=predicate,
+            )
+            write_deltalake(
+                location.uri,
+                spool.reader,
+                mode="overwrite",
+                predicate=predicate,
+                schema_mode=self._delta_schema_mode(spec.schema_mode),
+                **self._write_kwargs(location),
+            )
+        _log_write_commit("replace_partitions", location)
+
     def _predicate_to_sql(self, predicate: Any, params: Any) -> str:
         """Convert predicate to SQL (Polars uses internal SQL representation)."""
         return predicate_to_sql(predicate, params)

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/src/loom/etl/declarative/target/_into.py

@@ -82,38 +82,53 @@ class IntoTable:
         # Default write mode is REPLACE.  Call a write-mode method to change it.
         self._spec: Any = ReplaceSpec(table_ref=table_ref)
 
-    def append(self, *, schema: SchemaMode = SchemaMode.STRICT) -> IntoTable:
+    def append(
+        self,
+        *,
+        schema: SchemaMode = SchemaMode.STRICT,
+        streaming: bool = False,
+    ) -> IntoTable:
         """Write mode: append rows to the target table.
 
         Args:
-            schema: Schema evolution strategy.  Defaults to
-                    :attr:`~SchemaMode.STRICT`.
+            schema:    Schema evolution strategy.  Defaults to
+                       :attr:`~SchemaMode.STRICT`.
+            streaming: When ``True`` the Polars backend writes via an Arrow
+                       ``RecordBatchReader`` instead of materialising.
 
         Returns:
             New ``IntoTable`` with APPEND mode.
         """
-        return self._with(AppendSpec(table_ref=self._ref, schema_mode=schema))
+        return self._with(AppendSpec(table_ref=self._ref, schema_mode=schema, streaming=streaming))
 
-    def replace(self, *, schema: SchemaMode = SchemaMode.STRICT) -> IntoTable:
+    def replace(
+        self,
+        *,
+        schema: SchemaMode = SchemaMode.STRICT,
+        streaming: bool = False,
+    ) -> IntoTable:
         """Write mode: full replace of the target table.
 
         Overwrites all data in the table.  Use :meth:`replace_partitions` for
         partition-scoped overwrite or :meth:`replace_where` for predicate-scoped.
 
         Args:
-            schema: Schema evolution strategy.  Use
-                    :attr:`~SchemaMode.OVERWRITE` to replace the table schema
-                    entirely alongside the data.
+            schema:    Schema evolution strategy.  Use
+                       :attr:`~SchemaMode.OVERWRITE` to replace the table schema
+                       entirely alongside the data.
+            streaming: When ``True`` the Polars backend writes via an Arrow
+                       ``RecordBatchReader`` instead of materialising.
 
         Returns:
             New ``IntoTable`` with REPLACE mode.
         """
-        return self._with(ReplaceSpec(table_ref=self._ref, schema_mode=schema))
+        return self._with(ReplaceSpec(table_ref=self._ref, schema_mode=schema, streaming=streaming))
 
     def replace_partitions(
         self,
         *cols: str,
         schema: SchemaMode = SchemaMode.STRICT,
+        streaming: bool = False,
     ) -> IntoTable:
         """Write mode: replace the partitions present in the batch frame.
 
@@ -144,6 +159,7 @@ class IntoTable:
                 table_ref=self._ref,
                 partition_cols=cols,
                 schema_mode=schema,
+                streaming=streaming,
             )
         )
 
@@ -200,6 +216,7 @@ class IntoTable:
         predicate: PredicateNode,
         *,
         schema: SchemaMode = SchemaMode.STRICT,
+        streaming: bool = False,
     ) -> IntoTable:
         """Write mode: replace rows matching an explicit predicate.
 
@@ -228,6 +245,7 @@ class IntoTable:
                 table_ref=self._ref,
                 replace_predicate=predicate,
                 schema_mode=schema,
+                streaming=streaming,
             )
         )
 

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/src/loom/etl/declarative/target/_table.py

@@ -22,10 +22,16 @@ class AppendSpec:
     Args:
         table_ref:   Logical table reference.
         schema_mode: Schema evolution strategy.  Defaults to ``STRICT``.
+        streaming:   When ``True`` requests a streaming Arrow write.  In
+                     this release the Polars backend honours the flag
+                     only for :class:`ReplacePartitionsSpec`; other modes
+                     accept the flag for forward-compatibility but still
+                     materialise the frame.
     """
 
     table_ref: TableRef
     schema_mode: SchemaMode = SchemaMode.STRICT
+    streaming: bool = False
 
 
 @dataclass(frozen=True)
@@ -36,10 +42,16 @@ class ReplaceSpec:
         table_ref:   Logical table reference.
         schema_mode: Schema evolution strategy.  ``OVERWRITE`` replaces the
                      table schema alongside the data.
+        streaming:   When ``True`` requests a streaming Arrow write.  In
+                     this release the Polars backend honours the flag
+                     only for :class:`ReplacePartitionsSpec`; other modes
+                     accept the flag for forward-compatibility but still
+                     materialise the frame.
     """
 
     table_ref: TableRef
     schema_mode: SchemaMode = SchemaMode.STRICT
+    streaming: bool = False
 
 
 @dataclass(frozen=True)
@@ -53,11 +65,16 @@ class ReplacePartitionsSpec:
         table_ref:      Logical table reference.
         partition_cols: Partition columns used to build the predicate.
         schema_mode:    Schema evolution strategy.
+        streaming:      When ``True`` the Polars backend writes via an
+                        Arrow ``RecordBatchReader`` and computes the
+                        partition predicate via a cheap projection over
+                        the spool, bounding peak memory.
     """
 
     table_ref: TableRef
     partition_cols: tuple[str, ...]
     schema_mode: SchemaMode = SchemaMode.STRICT
+    streaming: bool = False
 
 
 @dataclass(frozen=True)
@@ -71,11 +88,18 @@ class ReplaceWhereSpec:
         table_ref:         Logical table reference.
         replace_predicate: Predicate node built with the col/params DSL.
         schema_mode:       Schema evolution strategy.
+        streaming:         When ``True`` requests a streaming Arrow write.
+                           In this release the Polars backend honours
+                           the flag only for
+                           :class:`ReplacePartitionsSpec`; other modes
+                           accept the flag for forward-compatibility but
+                           still materialise the frame.
     """
 
     table_ref: TableRef
     replace_predicate: PredicateNode
     schema_mode: SchemaMode = SchemaMode.STRICT
+    streaming: bool = False
 
 
 @dataclass(frozen=True)

{loom_kernel-0.10.0 → loom_kernel-0.10.2}/src/loom/etl/testing/_runners.py

@@ -50,6 +50,14 @@ class _PolarsStubReader:
         key = table_ref.ref if table_ref is not None else spec.alias
         return self._frames[key]
 
+    def read_streaming(
+        self,
+        spec: SourceSpec,
+        params_instance: Any,
+        /,
+    ) -> pl.LazyFrame:
+        return self.read(spec, params_instance)
+
     def execute_sql(self, frames: dict[str, Any], query: str, /) -> pl.LazyFrame:
         ctx = pl.SQLContext()
         for name, frame in frames.items():