pyspiral 0.6.10__cp312-abi3-macosx_11_0_arm64.whl → 0.6.12__cp312-abi3-macosx_11_0_arm64.whl

spiral/expressions/tiff.py

@@ -2,7 +2,6 @@ import numpy as np
 import pyarrow as pa
 
 from spiral.expressions.base import Expr, ExprLike
-from spiral.expressions.udf import RefUDF
 
 _TIFF_RES_DTYPE: pa.DataType = pa.struct(
     [
@@ -78,7 +77,7 @@ def select(
     return TiffSelectUDF()(expr, shape, indexes)
 
 
-class TiffReadUDF(RefUDF):
+class TiffReadUDF:
     def __init__(self):
         super().__init__("tiff.read")
 
@@ -122,7 +121,7 @@ class TiffReadUDF(RefUDF):
             return _return_result(result, indexes)
 
 
-class TiffSelectUDF(RefUDF):
+class TiffSelectUDF:
     def __init__(self):
         super().__init__("tiff.select")
 

spiral/expressions/udf.py

@@ -3,44 +3,53 @@ import abc
 import pyarrow as pa
 
 from spiral import _lib
-from spiral.expressions.base import Expr
+from spiral.expressions.base import Expr, ExprLike
 
 
-class BaseUDF:
-    def __init__(self, udf):
-        self._udf = udf
+class UDF(abc.ABC):
+    """A User-Defined Function (UDF). This class should be subclassed to define custom UDFs.
 
-    def __call__(self, *args) -> Expr:
-        """Create an expression that calls this UDF with the given arguments."""
-        from spiral import expressions as se
+    Example:
 
-        args = [se.lift(arg).__expr__ for arg in args]
-        return Expr(self._udf(args))
+    ```python
+    from spiral import expressions as se
+    import pyarrow as pa
 
-    @abc.abstractmethod
-    def return_type(self, *input_types: pa.DataType) -> pa.DataType: ...
+    class MyAdd(se.UDF):
+        def __init__(self):
+            super().__init__("my_add")
 
+        def return_type(self, scope: pa.DataType):
+            if not isinstance(scope, pa.StructType):
+                raise ValueError("Expected struct type as input")
+            return scope.field(0).type
 
-class UDF(BaseUDF):
-    """A User-Defined Function (UDF)."""
+        def invoke(self, scope: pa.Array):
+            if not isinstance(scope, pa.StructArray):
+                raise ValueError("Expected struct array as input")
+            return pa.compute.add(scope.field(0), scope.field(1))
 
-    def __init__(self, name: str):
-        super().__init__(_lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke))
+    my_add = MyAdd()
 
-    @abc.abstractmethod
-    def invoke(self, *input_args: pa.Array) -> pa.Array: ...
+    expr = my_add(table.select("first_arg", "second_arg"))
+    ```
+    """
 
+    def __init__(self, name: str):
+        self._udf = _lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke)
 
-class RefUDF(BaseUDF):
-    """A UDF over a single ref cell, and therefore can access the file object."""
+    def __call__(self, scope: ExprLike) -> Expr:
+        """Create an expression that calls this UDF with the given arguments."""
+        from spiral import expressions as se
 
-    def __init__(self, name: str):
-        super().__init__(_lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke, scope="ref"))
+        return Expr(self._udf(se.lift(scope).__expr__))
 
     @abc.abstractmethod
-    def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
-        """Invoke the UDF with the given arguments.
+    def return_type(self, scope: pa.DataType) -> pa.DataType:
+        """Must return the return type of the UDF given the input scope type."""
+        ...
 
-        NOTE: The first argument is always the ref cell. All array input args will be sliced to the appropriate row.
-        """
+    @abc.abstractmethod
+    def invoke(self, scope: pa.Array) -> pa.Array:
+        """Must implement the UDF logic given the input scope array."""
         ...

spiral/project.py

@@ -53,7 +53,7 @@ class Project:
         res = res[0]
 
         return Table(
-            self._spiral, self._spiral._core.table(res.id), identifier=f"{res.project_id}.{res.dataset}.{res.table}"
+            self._spiral, self._spiral.core.table(res.id), identifier=f"{res.project_id}.{res.dataset}.{res.table}"
         )
 
     def create_table(
@@ -78,7 +78,7 @@ class Project:
             key_schema = pa.schema(key_schema)
         key_schema = Schema.from_arrow(key_schema)
 
-        core_table = self._spiral._core.create_table(
+        core_table = self._spiral.core.create_table(
             project_id=self._id,
             dataset=dataset,
             table=table,
@@ -105,7 +105,7 @@ class Project:
             raise ValueError(f"Index not found: {name}")
         res = res[0]
 
-        return TextIndex(self._spiral._core.text_index(res.id), name=name)
+        return TextIndex(self._spiral.core.text_index(res.id), name=name)
 
     def create_text_index(
         self,
@@ -135,7 +135,7 @@ class Project:
         if where is not None:
             where = se.lift(where)
 
-        core_index = self._spiral._core.create_text_index(
+        core_index = self._spiral.core.create_text_index(
             project_id=self._id,
             name=name,
             projection=projection.__expr__,
@@ -154,7 +154,7 @@ class Project:
             raise ValueError(f"Index not found: {name}")
         res = res[0]
 
-        return KeySpaceIndex(self._spiral._core.key_space_index(res.id), name=name)
+        return KeySpaceIndex(self._spiral.core.key_space_index(res.id), name=name)
 
     def create_key_space_index(
         self,
@@ -185,7 +185,7 @@ class Project:
         if where is not None:
             where = se.lift(where)
 
-        core_index = self._spiral._core.create_key_space_index(
+        core_index = self._spiral.core.create_key_space_index(
             project_id=self._id,
             name=name,
             granularity=granularity,

spiral/scan.py

@@ -4,6 +4,7 @@ import pyarrow as pa
 
 from spiral.core.client import Shard, ShuffleConfig
 from spiral.core.table import Scan as CoreScan
+from spiral.core.table import ScanState as CoreScanState
 from spiral.core.table.spec import Schema
 from spiral.settings import CI, DEV
 
@@ -18,12 +19,34 @@ if TYPE_CHECKING:
     from spiral.dataloader import SpiralDataLoader, World  # noqa
 
 
+class ScanState:
+    """
+    Evaluated properties of the scan
+    """
+
+    __slots__ = ("core",)
+
+    def __init__(self, core: CoreScanState):
+        self.core = core
+
+    def __getstate__(self):
+        return self.core.to_json()
+
+    def __setstate__(self, state):
+        self.core = CoreScanState.from_json(state)
+
+
 class Scan:
     """Scan object."""
 
     def __init__(self, core: CoreScan):
         self.core = core
 
+    @property
+    def scan_state(self) -> ScanState:
+        """Returns evaluated properties of the scan."""
+        return ScanState(self.core.scan_state())
+
     @property
     def metrics(self) -> dict[str, Any]:
         """Returns metrics about the scan."""
@@ -34,6 +57,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.
 

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

@@ -18,6 +18,7 @@ if TYPE_CHECKING:
     from spiral.client import Spiral
     from spiral.dataloader import SpiralDataLoader
     from spiral.key_space_index import KeySpaceIndex
+    from spiral.scan import Scan
 
 
 class Table(Expr):
@@ -110,6 +111,25 @@ class Table(Expr):
                 partition_size_bytes=partition_size_bytes,
             )
 
+    def writeback(
+        self,
+        scan: "Scan",
+        *,
+        partition_size_bytes: int | None = 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 partition_size_bytes: The maximum partition size in bytes.
+        """
+        with self.txn() as txn:
+            txn.writeback(
+                scan,
+                partition_size_bytes=partition_size_bytes,
+            )
+
     def drop_columns(self, column_paths: list[str]) -> None:
         """
         Drops the specified columns from the table.
@@ -136,7 +156,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 +195,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 +228,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 +260,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,
@@ -282,7 +302,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,

spiral/transaction.py

@@ -1,5 +1,7 @@
 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:
@@ -39,6 +41,17 @@ class Transaction:
 
         self._core.write(record_batches, partition_size_bytes=partition_size_bytes)
 
+    def writeback(self, scan: Scan, *, partition_size_bytes: 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 partition_size_bytes: The maximum partition size in bytes.
+            If not provided, the default partition size is used.
+        """
+        self._core.writeback(scan.core, partition_size_bytes=partition_size_bytes)
+
     def drop_columns(self, column_paths: list[str]):
         """
         Drops the specified columns from the table.
@@ -49,6 +62,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/http.py

@@ -1,86 +0,0 @@
-import hishel
-import httpx
-import pyarrow as pa
-
-from spiral.expressions.base import Expr, ExprLike
-from spiral.expressions.struct import pack
-from spiral.expressions.udf import UDF
-from spiral.settings import APP_DIR
-
-
-def get(url: ExprLike, headers: ExprLike = None, force_cache: bool = False) -> Expr:
-    """Submit a GET request to either a scalar of vector of URLs."""
-    to_pack = {"url": url}
-    if headers is not None:
-        to_pack["headers"] = headers
-    return HttpGet(force_cache)(pack(to_pack))
-
-
-class HttpGet(UDF):
-    RES_DTYPE: pa.DataType = pa.struct(
-        [
-            pa.field("bytes", pa.large_binary()),
-            pa.field("status", pa.int32()),
-            pa.field("headers", pa.map_(pa.string(), pa.string())),
-        ]
-    )
-
-    def __init__(self, force_cache: bool = False):
-        super().__init__("http.get")
-        self._force_cache = force_cache
-
-    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
-        return HttpGet.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)}")
-        result = _http_request(input_args[0], self._force_cache)
-        if isinstance(result, pa.ChunkedArray):
-            result = result.combine_chunks()
-        return result
-
-
-def _http_request(arg: pa.Array, force_cache: bool) -> pa.Array:
-    client = _HttpClient()
-
-    if isinstance(arg, pa.StructArray):
-        # We assume a vector of requests, but with potentially many arguments
-        return pa.array(
-            [
-                _response_dict(
-                    client.request(
-                        req.get("method", "GET").upper(),
-                        req["url"],
-                        headers=req.get("headers", {}),
-                        extensions={"force_cache": force_cache},
-                    )
-                )
-                for req in arg.to_pylist()
-            ],
-            type=HttpGet.RES_DTYPE,
-        )
-
-    raise TypeError(f"Unsupported argument: {arg} ({type(arg)})")
-
-
-def _response_dict(response: httpx.Response) -> dict:
-    if response.status_code != 200:
-        raise ValueError(f"Request failed with status {response.status_code}")
-    return {
-        "bytes": response.read(),
-        "status": response.status_code,
-        "headers": dict(response.headers),
-    }
-
-
-class _HttpClient(hishel.CacheClient):
-    _instance: "_HttpClient" = None
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super().__new__(cls)
-        return cls._instance
-
-    def __init__(self):
-        super().__init__(storage=hishel.FileStorage(base_path=APP_DIR / "http.cache", ttl=3600))

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