pyspiral 0.6.9__cp312-abi3-macosx_11_0_arm64.whl → 0.7.12__cp312-abi3-macosx_11_0_arm64.whl

spiral/table.py

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any
 
 from spiral.core.table import Table as CoreTable
 from spiral.core.table.spec import Schema
+from spiral.enrichment import Enrichment
 from spiral.expressions.base import Expr, ExprLike
 from spiral.settings import settings
 from spiral.snapshot import Snapshot
@@ -12,12 +13,11 @@ if TYPE_CHECKING:
     import duckdb
     import polars as pl
     import pyarrow.dataset as ds
-    import streaming
-    import torch.utils.data as torchdata  # noqa
 
     from spiral.client import Spiral
     from spiral.dataloader import SpiralDataLoader
     from spiral.key_space_index import KeySpaceIndex
+    from spiral.streaming_ import SpiralStream
 
 
 class Table(Expr):
@@ -50,6 +50,14 @@ class Table(Expr):
         """Returns the fully qualified identifier of the table."""
         return self._identifier or self.table_id
 
+    @property
+    def project(self) -> str | None:
+        """Returns the project of the table."""
+        if self._identifier is None:
+            return None
+        project, _, _ = self._identifier.split(".")
+        return project
+
     @property
     def dataset(self) -> str | None:
         """Returns the dataset of the table."""
@@ -75,7 +83,7 @@ class Table(Expr):
     def __repr__(self):
         return f'Table("{self.identifier}")'
 
-    def __getitem__(self, item: str) -> Expr:
+    def __getitem__(self, item: str | int | list[str]) -> Expr:
         return super().__getitem__(item)
 
     def select(self, *paths: str, exclude: list[str] = None) -> "Expr":
@@ -110,6 +118,28 @@ class Table(Expr):
                 partition_size_bytes=partition_size_bytes,
             )
 
+    def enrich(
+        self,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+    ) -> Enrichment:
+        """Returns an Enrichment object that, when applied, produces new columns.
+
+        Enrichment can be applied in different ways, e.g. distributed.
+
+        :param projections: Projection expressions deriving new columns to write back.
+            Expressions can be over multiple Spiral tables, but all tables including
+            this one must share the same key schema.
+        :param where: Optional filter expression to apply when reading the input tables.
+        """
+        from spiral import expressions as se
+
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return Enrichment(self, projection, where)
+
     def drop_columns(self, column_paths: list[str]) -> None:
         """
         Drops the specified columns from the table.
@@ -136,7 +166,7 @@ class Table(Expr):
              it is important that the primary key columns are unique within the transaction.
              The behavior is undefined if this is not the case.
         """
-        return Transaction(self.spiral._core.transaction(self.core, settings().file_format, retries=retries))
+        return Transaction(self.spiral.core.transaction(self.core, settings().file_format, retries=retries))
 
     def to_dataset(self) -> "ds.Dataset":
         """Returns a PyArrow Dataset representing the table."""
@@ -175,7 +205,7 @@ class Table(Expr):
         if index.asof == 0:
             raise ValueError("Index have to be synced before it can be used.")
 
-        shards = self.spiral._core._ops().compute_shards(index=index.core)
+        shards = self.spiral.internal.compute_shards(index=index.core)
 
         return self.spiral.scan(
             projection if projection is not None else index.projection,
@@ -208,7 +238,7 @@ class Table(Expr):
         if index.asof == 0:
             raise ValueError("Index have to be synced before it can be used.")
 
-        shards = self.spiral._core._ops().compute_shards(index=index.core)
+        shards = self.spiral.core.internal.compute_shards(index=index.core)
 
         return self.spiral.scan(
             projection if projection is not None else index.projection,
@@ -240,7 +270,7 @@ class Table(Expr):
         if index.asof == 0:
             raise ValueError("Index have to be synced before it can be used.")
 
-        shards = self.spiral._core._ops().compute_shards(index=index.core)
+        shards = self.spiral.core.internal.compute_shards(index=index.core)
 
         return self.spiral.scan(
             index.projection,
@@ -255,7 +285,7 @@ class Table(Expr):
         projection: Expr | None = None,
         cache_dir: str | None = None,
         shard_row_block_size: int | None = None,
-    ) -> "streaming.Stream":
+    ) -> "SpiralStream":
         """Returns a stream to be used with MosaicML's StreamingDataset.
 
         Requires `streaming` package to be installed.
@@ -282,7 +312,7 @@ class Table(Expr):
             where=index.filter,
             asof=index.asof,
         )
-        shards = self.spiral._core._ops().compute_shards(index=index.core)
+        shards = self.spiral.internal.compute_shards(index=index.core)
 
         return SpiralStream(
             sp=self.spiral,
@@ -290,4 +320,4 @@ class Table(Expr):
             shards=shards,
             cache_dir=cache_dir,
             shard_row_block_size=shard_row_block_size,
-        )  # type: ignore[return-value]
+        )

spiral/transaction.py

@@ -1,5 +1,13 @@
+import logging
+from pathlib import Path
+
+from spiral.core.table import KeyRange
 from spiral.core.table import Transaction as CoreTransaction
+from spiral.core.table.spec import Operation
 from spiral.expressions.base import ExprLike
+from spiral.scan import Scan
+
+logger = logging.getLogger(__name__)
 
 
 class Transaction:
@@ -17,6 +25,10 @@ class Transaction:
         """The status of the transaction."""
         return self._core.status
 
+    def is_empty(self) -> bool:
+        """Check if the transaction has no operations."""
+        return self._core.is_empty()
+
     def __enter__(self):
         return self
 
@@ -39,6 +51,27 @@ class Transaction:
 
         self._core.write(record_batches, partition_size_bytes=partition_size_bytes)
 
+    def writeback(
+        self,
+        scan: Scan,
+        *,
+        key_range: KeyRange | None = None,
+        partition_size_bytes: int | None = None,
+        batch_readahead: int | None = None,
+    ):
+        """Write back the results of a scan to the table.
+
+        :param scan: The scan to write back.
+            The scan does NOT need to be over the same table as transaction,
+            but it does need to have the same key schema.
+        :param key_range: Optional key range to limit the writeback to.
+        :param partition_size_bytes: The maximum partition size in bytes.
+        :param batch_readahead: The number of batches to read ahead when evaluating the scan.
+        """
+        self._core.writeback(
+            scan.core, key_range=key_range, partition_size_bytes=partition_size_bytes, batch_readahead=batch_readahead
+        )
+
     def drop_columns(self, column_paths: list[str]):
         """
         Drops the specified columns from the table.
@@ -49,9 +82,73 @@ class Transaction:
         """
         self._core.drop_columns(column_paths)
 
-    def commit(self):
+    def take(self) -> list[Operation]:
+        """Take the operations from the transaction
+
+        Transaction can no longer be committed or aborted after calling this method.
+        ."""
+        return self._core.take()
+
+    def include(self, ops: list[Operation]):
+        """Include the given operations in the transaction.
+
+        Checks for conflicts between the included operations and any existing operations.
+        """
+        self._core.include(ops)
+
+    def commit(self, *, compact: bool = False, tx_dump: str | None = None):
         """Commit the transaction."""
-        self._core.commit()
+        if tx_dump is not None:
+            try:
+                # Create parent directories if they don't exist
+                dump_path = Path(tx_dump)
+                dump_path.parent.mkdir(parents=True, exist_ok=True)
+
+                # Write operations to a JSONL file
+                with open(dump_path, "w") as f:
+                    for op in self._core.ops():
+                        f.write(op.to_json() + "\n")
+
+                logger.info(f"Transaction dumped to {tx_dump}")
+            except Exception as e:
+                logger.error(f"Failed to dump transaction to {tx_dump}: {e}")
+
+        self._core.commit(compact=compact)
+
+    @staticmethod
+    def load_dumps(*tx_dump: str) -> list[Operation]:
+        """Load a transaction from a dump file."""
+        import json
+
+        dumps = list(tx_dump)
+        ops: list[Operation] = []
+
+        for dump in dumps:
+            with open(dump) as f:
+                lines = f.readlines()
+
+            for line in lines:
+                line = line.strip()
+                if not line:
+                    continue
+
+                # Each line may contain multiple JSON objects concatenated together
+                # This is due to a bug in the dump writing code.
+                # Use JSONDecoder to parse them one by one
+                decoder = json.JSONDecoder()
+                idx = 0
+                while idx < len(line):
+                    try:
+                        obj, end_idx = decoder.raw_decode(line, idx)
+                        ops.append(Operation.from_json(json.dumps(obj)))
+                        idx = end_idx
+                        # Skip whitespace between JSON objects
+                        while idx < len(line) and line[idx].isspace():
+                            idx += 1
+                    except json.JSONDecodeError as e:
+                        raise ValueError(f"Failed to parse JSON at position {idx} in line: {line}") from e
+
+        return ops
 
     def abort(self):
         """Abort the transaction."""

spiral/expressions/io.py

@@ -1,100 +0,0 @@
-import tarfile
-from io import BytesIO
-
-import pyarrow as pa
-
-from spiral.expressions.base import Expr, ExprLike
-from spiral.expressions.struct import pack
-from spiral.expressions.udf import UDF
-
-
-def read_file(path: ExprLike) -> Expr:
-    """
-    Read file path(s) from disk into a struct with a single field "bytes" containing the file contents.
-
-    Args:
-        path: Expression evaluating to an array of strings representing local disk paths.
-    """
-    to_pack = {"path": path}
-    return FileRead()(pack(to_pack))
-
-
-class FileRead(UDF):
-    RES_DTYPE: pa.DataType = pa.struct(
-        [
-            pa.field("bytes", pa.large_binary()),
-        ]
-    )
-
-    def __init__(self):
-        super().__init__("file.read")
-
-    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
-        return FileRead.RES_DTYPE
-
-    def invoke(self, *input_args: pa.Array) -> pa.Array:
-        if len(input_args) != 1:
-            raise ValueError(f"Expected 1 argument, got {len(input_args)}")
-        arg = input_args[0]
-
-        res = []
-        for req in arg:
-            with open(req["path"].as_py(), "rb") as f:
-                res.append({"bytes": f.read()})
-
-        return pa.array(res, type=FileRead.RES_DTYPE)
-
-
-def read_tar(path: ExprLike = None, bytes_: ExprLike = None) -> "Expr":
-    # Untar a vector of paths / byte arrays representing tarballs.
-    if path is None and bytes_ is None:
-        raise ValueError("Expected either path or bytes_ to be provided")
-    to_pack = {}
-    if path is not None:
-        to_pack["path"] = path
-    if bytes_ is not None:
-        to_pack["bytes"] = bytes_
-    return TarRead()(pack(to_pack))
-
-
-class TarRead(UDF):
-    RES_DTYPE = pa.list_(
-        pa.struct(
-            [
-                pa.field("name", pa.string()),
-                pa.field("bytes", pa.large_binary()),
-            ]
-        )
-    )
-
-    def __init__(self):
-        super().__init__("tar.read")
-
-    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
-        return TarRead.RES_DTYPE
-
-    def invoke(self, *input_args: pa.Array) -> pa.Array:
-        if len(input_args) != 1:
-            raise ValueError(f"Expected 1 argument, got {len(input_args)}")
-        arg = input_args[0]
-
-        res = []
-        for req in arg:
-            if "path" in req:
-                kwargs = {"name": req["path"].as_py()}
-            elif "bytes" in req:
-                kwargs = {"fileobj": BytesIO(req["bytes"].as_py())}
-            else:
-                raise ValueError("Expected path or bytes_ to be provided")
-
-            files = []
-            with tarfile.open(**kwargs) as f:
-                for m in f.getmembers():
-                    m: tarfile.TarInfo
-                    if m.type == tarfile.DIRTYPE:
-                        continue
-                    # TODO(ngates): skip other types too maybe? Why are we even skipping directories?
-                    files.append({"name": m.name, "bytes": f.extractfile(m).read()})
-            res.append(files)
-
-        return pa.array(res, type=TarRead.RES_DTYPE)

spiral/expressions/mp4.py

@@ -1,62 +0,0 @@
-import pyarrow as pa
-
-from spiral.expressions.base import Expr, ExprLike
-
-_MP4_RES_DTYPE: pa.DataType = pa.struct(
-    [
-        pa.field("pixels", pa.large_binary()),
-        pa.field("height", pa.uint32()),
-        pa.field("width", pa.uint32()),
-        pa.field("frames", pa.uint32()),
-    ]
-)
-
-
-# TODO(marko): Support optional range and crop.
-#   IMPORTANT: Frames is currently broken and defaults to full.
-def read(expr: ExprLike | str, frames: ExprLike | str, crop: ExprLike | str):
-    """
-    Read referenced cell in a `MP4` format. Requires `ffmpeg`.
-
-    Args:
-        expr: The referenced `Mp4` bytes.
-            A str is assumed to be the `se.aux` expression.
-        frames: The range of frames to read. Each element must be a list of two uint32,
-            frame start and frame end, or null / empty list to read all frames.
-            A str is assumed to be the `se.aux` expression.
-        crop: The crop of the frames to read. Each element must be a list of four uint32,
-            x, y, width, height or null / empty list to read full frames.
-            A str is assumed to be the `se.aux` expression.
-
-    Returns:
-        An array where each element is a decoded cropped video with fields:
-            pixels: RGB8 bytes, frames * width * height * 3.
-            width: Width of the image with type `pa.uint32()`.
-            height: Height of the image with type `pa.uint32()`.
-            frames: Number of frames with type `pa.uint32()`.
-    """
-    from spiral import _lib
-    from spiral.expressions import aux, lift
-
-    if isinstance(expr, str):
-        expr = aux(
-            expr,
-            pa.struct([("__ref__", pa.struct([("id", pa.string()), ("begin", pa.uint64()), ("end", pa.uint64())]))]),
-        )
-    if isinstance(frames, str):
-        frames = aux(frames, pa.list_(pa.uint32()))
-    if isinstance(crop, str):
-        crop = aux(crop, pa.list_(pa.uint32()))
-
-    expr = lift(expr)
-    frames = lift(frames)
-    crop = lift(crop)
-
-    return Expr(
-        _lib.expr.video.read(
-            expr.__expr__,
-            frames.__expr__,
-            crop.__expr__,
-            format="mp4",
-        )
-    )

spiral/expressions/png.py

@@ -1,18 +0,0 @@
-from spiral.expressions.base import Expr, ExprLike
-
-
-def encode(expr: ExprLike) -> Expr:
-    """Encode the given expression as a PNG image.
-
-    Args:
-        expr: The expression to encode.
-            Expects a struct with `pixels`, `width`, `height`, `channels`, `channel_bit_depth` fields.
-
-    Returns:
-        The encoded PNG images.
-    """
-    from spiral import _lib
-    from spiral.expressions import lift
-
-    expr = lift(expr)
-    return Expr(_lib.expr.img.encode(expr.__expr__, format="png"))

spiral/expressions/qoi.py

@@ -1,18 +0,0 @@
-from spiral.expressions.base import Expr, ExprLike
-
-
-def encode(expr: ExprLike) -> Expr:
-    """Encode the given expression as a QOI image.
-
-    Args:
-        expr: The expression to encode.
-            Expects a struct with `pixels`, `width`, `height`, `channels`, `channel_bit_depth` fields.
-
-    Returns:
-        The encoded QOI images.
-    """
-    from spiral import _lib
-    from spiral.expressions import lift
-
-    expr = lift(expr)
-    return Expr(_lib.expr.img.encode(expr.__expr__, format="qoi"))

spiral/expressions/refs.py

@@ -1,58 +0,0 @@
-import pyarrow as pa
-
-from spiral.expressions.base import Expr, ExprLike
-
-
-def ref(expr: ExprLike, field: str | None = None) -> Expr:
-    """Store binary values as references. This expression can only be used on write.
-
-    It is often better to store large cell values, such as bytes columns, that aren't used in filter expressions as
-    references. This enables more efficient scan pruning. Many of the Spiral's cell pushdown expressions work
-    over references.
-
-    Args:
-        expr: The expression to store as a reference.
-        field: If the expr evaluates into struct, the field name of that struct that should be referenced.
-            If `None`, the expr must evaluate into a type that supports referencing.
-    """
-    from spiral import _lib
-    from spiral.expressions import lift
-
-    expr = lift(expr)
-    return Expr(_lib.expr.refs.ref(expr.__expr__, field))
-
-
-def deref(expr: ExprLike | str, field: str | None = None) -> Expr:
-    """De-reference referenced values.
-
-    See `ref` for more information on Spiral's reference values. This expression is used to de-reference referenced
-    column back into their original form, e.g. binary.
-
-    Args:
-        expr: The expression to de-reference. A str is assumed to be the `se.aux` expression.
-        field: If the expr evaluates into struct, the field name of that struct that should be de-referenced.
-            If `None`, the expr must evaluate into a reference type.
-    """
-    from spiral import _lib
-    from spiral.expressions import aux, lift
-
-    if isinstance(expr, str):
-        expr = aux(
-            expr,
-            pa.struct([("__ref__", pa.struct([("id", pa.string()), ("begin", pa.uint64()), ("end", pa.uint64())]))]),
-        )
-
-    expr = lift(expr)
-    return Expr(_lib.expr.refs.deref(expr.__expr__, field=field))
-
-
-def nbytes(expr: ExprLike) -> Expr:
-    """Return the number of bytes in a reference.
-
-    Args:
-        expr: The ref expression to get the number of bytes from.
-    """
-    from spiral.expressions import lift
-
-    expr = lift(expr)
-    return expr["__ref__"]["end"] - expr["__ref__"]["begin"]