pyspiral 0.6.11__cp312-abi3-manylinux_2_28_aarch64.whl → 0.6.13__cp312-abi3-manylinux_2_28_aarch64.whl

spiral/scan.py

@@ -1,8 +1,10 @@
+from functools import partial
 from typing import TYPE_CHECKING, Any, Optional
 
 import pyarrow as pa
 
 from spiral.core.client import Shard, ShuffleConfig
+from spiral.core.table import KeyRange
 from spiral.core.table import Scan as CoreScan
 from spiral.core.table.spec import Schema
 from spiral.settings import CI, DEV
@@ -15,13 +17,15 @@ if TYPE_CHECKING:
     import streaming  # noqa
     import torch.utils.data as torchdata  # noqa
 
+    from spiral.client import Spiral
     from spiral.dataloader import SpiralDataLoader, World  # noqa
 
 
 class Scan:
     """Scan object."""
 
-    def __init__(self, core: CoreScan):
+    def __init__(self, spiral: "Spiral", core: CoreScan):
+        self.spiral = spiral
         self.core = core
 
     @property
@@ -34,6 +38,11 @@ class Scan:
         """Returns the schema of the scan."""
         return self.core.schema()
 
+    @property
+    def key_schema(self) -> Schema:
+        """Returns the key schema of the scan."""
+        return self.core.key_schema()
+
     def is_empty(self) -> bool:
         """Check if the Spiral is empty for the given key range.
 
@@ -44,6 +53,8 @@ class Scan:
 
     def to_record_batches(
         self,
+        *,
+        key_range: KeyRange | None = None,
         key_table: pa.Table | pa.RecordBatchReader | None = None,
         batch_size: int | None = None,
         batch_readahead: int | None = None,
@@ -51,6 +62,9 @@ class Scan:
         """Read as a stream of RecordBatches.
 
         Args:
+            key_range: Optional key range to filter the scan.
+                If provided, the scan will only return rows within the key range.
+                Only one of key_range or key_table can be provided.
             key_table: a table of keys to "take" (including aux columns for cell-push-down).
                 If None, the scan will be executed without a key table.
             batch_size: the maximum number of rows per returned batch.
@@ -58,6 +72,9 @@ class Scan:
                     RecordBatchReader, the batch_size argument must be None, and the existing batching is respected.
             batch_readahead: the number of batches to prefetch in the background.
         """
+        if key_range is not None and key_table is not None:
+            raise ValueError("Only one of key_range or key_table can be provided.")
+
         if isinstance(key_table, pa.RecordBatchReader):
             if batch_size is not None:
                 raise ValueError(
@@ -66,46 +83,54 @@ class Scan:
         elif isinstance(key_table, pa.Table):
             key_table = key_table.to_reader(max_chunksize=batch_size)
 
-        return self.core.to_record_batches(key_table=key_table, batch_readahead=batch_readahead)
+        return self.core.to_record_batches(key_range=key_range, key_table=key_table, batch_readahead=batch_readahead)
 
     def to_table(
         self,
+        *,
+        key_range: KeyRange | None = None,
         key_table: pa.Table | pa.RecordBatchReader | None = None,
     ) -> pa.Table:
         """Read into a single PyArrow Table.
 
         Args:
+            key_range: Optional key range to filter the scan.
+                If provided, the scan will only return rows within the key range.
+                Only one of key_range or key_table can be provided.
             key_table: a table of keys to "take" (including aux columns for cell-push-down).
                 If None, the scan will be executed without a key table.
         """
         # NOTE: Evaluates fully on Rust side which improved debuggability.
-        if DEV and not CI and key_table is None:
+        if DEV and not CI and key_table is None and key_range is None:
             rb = self.core.to_record_batch()
             return pa.Table.from_batches([rb])
 
-        return self.to_record_batches(key_table=key_table).read_all()
+        return self.to_record_batches(key_range=key_range, key_table=key_table).read_all()
 
     def to_dask(self) -> "dd.DataFrame":
         """Read into a Dask DataFrame.
 
         Requires the `dask` package to be installed.
+
+        IMPORTANT: Dask execution has some limitations, e.g. UDFs are not currently supported. These limitations
+        usually manifest as serialization errors when Dask workers attempt to serialize the state. If you are
+        encountering such issues, please reach out to the support for assistance.
         """
         import dask.dataframe as dd
-        import pandas as pd
 
-        def _read_shard(shard: Shard) -> pd.DataFrame:
-            # TODO(ngates): we need a way to preserve the existing asofs?
-            raise NotImplementedError()
-
-        # Fetch a set of partition ranges
+        _read_shard = partial(
+            _read_shard_task,
+            settings_dict=self.spiral.config.model_dump(),
+            state_json=self.core.scan_state().to_json(),
+        )
         return dd.from_map(_read_shard, self.shards())
 
-    def to_pandas(self) -> "pd.DataFrame":
+    def to_pandas(self, *, key_range: KeyRange | None = None) -> "pd.DataFrame":
         """Read into a Pandas DataFrame.
 
         Requires the `pandas` package to be installed.
         """
-        return self.to_table().to_pandas()
+        return self.to_table(key_range=key_range).to_pandas()
 
     def to_polars(self) -> "pl.DataFrame":
         """Read into a Polars DataFrame.
@@ -160,16 +185,18 @@ class Scan:
 
         Returns:
             SpiralDataLoader with shards partitioned for this rank.
-        """
-        # Example usage:
-        #
-        # Auto-detect from PyTorch distributed:
-        #   loader: SpiralDataLoader = scan.to_distributed_data_loader(batch_size=32)
-        #
-        # Explicit world configuration:
-        #   world = World(rank=0, world_size=4)
-        #   loader: SpiralDataLoader = scan.to_distributed_data_loader(world=world, batch_size=32)
 
+        Auto-detect from PyTorch distributed:
+        ```python
+        loader: SpiralDataLoader = scan.to_distributed_data_loader(batch_size=32)
+        ```
+
+        Explicit world configuration:
+        ```python
+        world = World(rank=0, world_size=4)
+        loader: SpiralDataLoader = scan.to_distributed_data_loader(world=world, batch_size=32)
+        ```
+        """
         from spiral.dataloader import SpiralDataLoader, World
 
         if world is None:
@@ -203,19 +230,21 @@ class Scan:
 
         Returns:
             New SpiralDataLoader instance configured to resume from the checkpoint.
+
+        Save checkpoint during training:
+        ```python
+        loader = scan.to_distributed_data_loader(batch_size=32, seed=42)
+        checkpoint = loader.state_dict()
+        ```
+
+        Resume later - uses same shards from checkpoint:
+        ```python
+        resumed_loader = scan.resume_data_loader(
+            checkpoint,
+            batch_size=32,
+            transform_fn=my_transform,
+        )
         """
-        # Example usage:
-        #
-        # Save checkpoint during training:
-        #   loader = scan.to_distributed_data_loader(batch_size=32, seed=42)
-        #   checkpoint = loader.state_dict()
-        #
-        # Resume later - uses same shards from checkpoint:
-        #   resumed_loader = scan.resume_data_loader(
-        #       checkpoint,
-        #       batch_size=32,
-        #       transform_fn=my_transform,
-        #   )
         from spiral.dataloader import SpiralDataLoader
 
         return SpiralDataLoader.from_state_dict(self, state, **kwargs)
@@ -283,3 +312,17 @@ class Scan:
         from spiral.debug.metrics import display_metrics
 
         display_metrics(self.metrics)
+
+
+# NOTE(marko): This function must be picklable!
+def _read_shard_task(shard: Shard, *, settings_dict, state_json) -> "pd.DataFrame":
+    from spiral import Spiral
+    from spiral.core.table import ScanState
+    from spiral.settings import Settings
+
+    settings: Settings = Settings.model_validate(settings_dict)
+    sp = Spiral(config=settings)
+    state = ScanState.from_json(state_json)
+    task_scan = Scan(sp, sp.core.load_scan(state))
+
+    return task_scan.to_pandas(key_range=shard.key_range)

spiral/settings.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Annotated
 
 import typer
-from pydantic import Field, ValidatorFunctionWrapHandler, WrapValidator
+from pydantic import Field, PlainSerializer, ValidatorFunctionWrapHandler, WrapValidator
 from pydantic_settings import (
     BaseSettings,
     InitSettingsSource,
@@ -28,13 +28,16 @@ PACKAGE_NAME = "pyspiral"
 
 
 def validate_token(v, handler: ValidatorFunctionWrapHandler):
-    if isinstance(v, str):
-        return Token(v)
-    else:
-        raise ValueError("Token value must be a string")
+    if not isinstance(v, str):
+        raise ValueError("Token value (SPIRAL__SPIRALDB__TOKEN) must be a string")
+    return Token(v)
 
 
-TokenType = Annotated[Token, WrapValidator(validate_token)]
+TokenType = Annotated[
+    Token,
+    WrapValidator(validate_token),
+    PlainSerializer(lambda token: token.expose_secret(), return_type=str),
+]
 
 
 class SpiralDBSettings(BaseSettings):

spiral/streaming_/stream.py

@@ -101,7 +101,7 @@ class SpiralStream:
             return 0
 
         # Prepare the shard, writing it to disk.
-        self._sp._ops().prepare_shard(
+        self._sp.internal.prepare_shard(
             shard_path, self._scan.core, shard.shard, row_block_size=self._shard_row_block_size
         )
 

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."""
@@ -110,6 +118,30 @@ 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
+
+        # Combine table with all projections into a single struct.
+        # The table is included to ensure key columns are present in the scan output.
+        projection = se.merge(self, *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 +168,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 +207,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 +240,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 +272,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 +287,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 +314,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 +322,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,8 @@
+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
 
 
 class Transaction:
@@ -17,6 +20,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 +46,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,6 +77,20 @@ class Transaction:
         """
         self._core.drop_columns(column_paths)
 
+    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):
         """Commit the transaction."""
         self._core.commit()

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"))