pyspiral 0.8.9__cp311-abi3-macosx_11_0_arm64.whl

spiral/expressions/tiff.py

@@ -0,0 +1,225 @@
+import pyarrow as pa
+
+from spiral.expressions.base import Expr, ExprLike
+
+_TIFF_RES_DTYPE: pa.DataType = pa.struct(
+    [
+        pa.field("pixels", pa.large_binary()),
+        pa.field("height", pa.uint32()),
+        pa.field("width", pa.uint32()),
+        pa.field("channels", pa.uint8()),
+        pa.field("channel_bit_depth", pa.uint8()),
+    ]
+)
+
+
+def read(
+    expr: ExprLike,
+    indexes: ExprLike | int | None = None,
+    window: ExprLike | tuple[tuple[int, int], tuple[int, int]] | None = None,
+    boundless: ExprLike | bool | None = None,
+) -> Expr:
+    """
+    Read referenced cell in a `TIFF` format. Requires `rasterio` to be installed.
+
+    Args:
+        expr: The referenced `TIFF` bytes.
+        indexes: The band indexes to read. Defaults to all.
+        window: The window to read. In format (row_range_tuple, col_range_tuple). Defaults to full window.
+        boundless: If `True`, windows that extend beyond the dataset's extent
+            are permitted and partially or completely filled arrays will be returned as appropriate.
+
+    Returns:
+        An array where each element is a decoded image with fields:
+            pixels: bytes of shape (channels, width, height).
+            width: Width of the image with type `pa.uint32()`.
+            height: Height of the image with type `pa.uint32()`.
+            channels: Number of channels of the image with type `pa.uint8()`.
+                If `indexes` is not None, this is the length of `indexes` or 1 if `indexes` is an int.
+            channel_bit_depth: Bit depth of the channel with type `pa.uint8()`.
+    """
+    try:
+        import rasterio  # noqa: F401
+    except ImportError:
+        raise ImportError("`rasterio` is required for tiff.read")
+
+    return TiffReadUDF()(expr, indexes, window, boundless)
+
+
+def select(
+    expr: ExprLike,
+    shape: ExprLike | dict,
+    indexes: ExprLike | int | None = None,
+) -> Expr:
+    """
+    Select the shape out of the referenced cell in a `TIFF` format. Requires `rasterio` to be installed.
+
+    Args:
+        expr: The referenced `TIFF` bytes.
+        shape: [GeoJSON-like](https://geojson.org/) shape.
+        indexes: The band indexes to read. Defaults to all.
+
+    Returns:
+        An array where each element is a decoded image with fields:
+            pixels: bytes of shape (len(indexes) or 1, width, height).
+            width: Width of the image with type `pa.uint32()`.
+            height: Height of the image with type `pa.uint32()`.
+            channels: Number of channels of the image with type `pa.uint8()`.
+                If `indexes` is not None, this is the length of `indexes` or 1 if `indexes` is an int.
+            channel_bit_depth: Bit depth of the channel with type `pa.uint8()`.
+    """
+    try:
+        import rasterio  # noqa: F401
+    except ImportError:
+        raise ImportError("`rasterio` is required for tiff.select")
+
+    return TiffSelectUDF()(expr, shape, indexes)
+
+
+class TiffReadUDF:
+    def __init__(self):
+        super().__init__("tiff.read")
+
+    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
+        return _TIFF_RES_DTYPE
+
+    def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
+        try:
+            import rasterio
+        except ImportError:
+            raise ImportError("`rasterio` is required for tiff.read")
+
+        from rasterio.windows import Window
+
+        if len(input_args) != 4:
+            raise ValueError("tiff.read expects exactly 4 arguments: expr, indexes, window, boundless")
+
+        _, indexes, window, boundless = input_args
+
+        indexes = indexes[0].as_py()
+        if indexes is not None and not isinstance(indexes, int) and not isinstance(indexes, list):
+            raise ValueError(f"tiff.read expects indexes to be None or an int or a list, got {indexes}")
+
+        boundless = boundless[0].as_py()
+        if boundless is not None and not isinstance(boundless, bool):
+            raise ValueError(f"tiff.read expects boundless to be None or a bool, got {boundless}")
+
+        window = window[0].as_py()
+        if window is not None:
+            if len(window) != 2:
+                raise ValueError(f"tiff.read window invalid, got {window}")
+            window = Window.from_slices(slice(*window[0]), slice(*window[1]), boundless=boundless or False)
+
+        opener = _VsiOpener(fp)
+        with rasterio.open("ref", opener=opener) as src:
+            src: rasterio.DatasetReader
+            # TODO(marko): We know the size and dtype so we should be able to preallocate the result and read into it.
+            #   This matters more if we want to rewrite this function to work with multiple inputs at once, in which
+            #   case we should first consider using Rust GDAL bindings - I believe rasterio uses GDAL under the hood.
+            result = src.read(indexes=indexes, window=window)
+            return _return_result(result, indexes)
+
+
+class TiffSelectUDF:
+    def __init__(self):
+        super().__init__("tiff.select")
+
+    def return_type(self, *input_types: pa.DataType) -> pa.DataType:
+        return _TIFF_RES_DTYPE
+
+    def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
+        try:
+            import rasterio
+        except ImportError:
+            raise ImportError("`rasterio` is required for tiff.select")
+
+        from rasterio.mask import raster_geometry_mask
+
+        if len(input_args) != 3:
+            raise ValueError("tiff.select expects exactly 3 arguments: expr, shape, indexes")
+
+        _, shape, indexes = input_args
+
+        shape = shape[0].as_py()
+        if shape is None:
+            raise ValueError("tiff.select expects shape to be a GeoJSON-like shape")
+
+        indexes = indexes[0].as_py()
+        if indexes is not None and not isinstance(indexes, int) and not isinstance(indexes, list):
+            raise ValueError(f"tiff.select expects indexes to be None or an int or a list, got {indexes}")
+
+        opener = _VsiOpener(fp)
+        with rasterio.open("ref", opener=opener) as src:
+            src: rasterio.DatasetReader
+
+            shape_mask, _, window = raster_geometry_mask(src, [shape], crop=True)
+            out_shape = (src.count,) + shape_mask.shape
+
+            result = src.read(window=window, indexes=indexes, out_shape=out_shape, masked=True)
+            return _return_result(result, indexes)
+
+
+def _return_result(result, indexes) -> pa.Array:
+    import numpy as np
+
+    result: np.ndarray
+
+    channels = result.shape[0]
+    if indexes is None:
+        pass
+    elif isinstance(indexes, int):
+        assert channels == 1, f"Expected 1 channel, got {channels}"
+    else:
+        assert channels == len(indexes), f"Expected {len(indexes)} channels, got {channels}"
+
+    if result.dtype == np.uint8:
+        channel_bit_depth = 8
+    elif result.dtype == np.uint16:
+        channel_bit_depth = 16
+    else:
+        raise ValueError(f"Unsupported bit width: {result.dtype}")
+
+    return pa.array(
+        [
+            {
+                "pixels": result.tobytes(),
+                "height": result.shape[1],
+                "width": result.shape[2],
+                "channels": channels,
+                "channel_bit_depth": channel_bit_depth,
+            }
+        ],
+        type=_TIFF_RES_DTYPE,
+    )
+
+
+class _VsiOpener:
+    """
+    VSI file opener which returns a constant file-like on open.
+
+    Must match https://rasterio.readthedocs.io/en/stable/topics/vsi.html#python-file-and-filesystem-openers spec but
+    only `open` is needed when going through rasterio.
+    """
+
+    def __init__(self, file_like):
+        self._file_like = file_like
+
+    def open(self, _path, mode):
+        if mode not in {"r", "rb"}:
+            raise ValueError(f"Unsupported mode: {mode}")
+        return self._file_like
+
+    def isdir(self, _):
+        return False
+
+    def isfile(self, _):
+        return False
+
+    def mtime(self, _):
+        return 0
+
+    def size(self, _):
+        return self._file_like.size()
+
+    def modified(self, _):
+        raise NotImplementedError

spiral/expressions/udf.py

@@ -0,0 +1,66 @@
+import abc
+
+import pyarrow as pa
+
+from spiral import _lib
+from spiral.expressions.base import Expr, ExprLike
+
+
+class UDF(abc.ABC):
+    """A User-Defined Function (UDF). This class should be subclassed to define custom UDFs.
+
+    Example:
+
+    ```python
+    import spiral
+    from spiral.demo import fineweb
+
+    sp = spiral.Spiral()
+    fineweb_table = fineweb(sp)
+
+    from spiral import expressions as se
+    import pyarrow as pa
+
+    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
+
+        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))
+
+    my_add = MyAdd()
+
+    expr = my_add(fineweb_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)
+
+    def __call__(self, scope: ExprLike) -> Expr:
+        """Create an expression that calls this UDF with the given arguments."""
+        from spiral import expressions as se
+
+        return Expr(self._udf(se.lift(scope).__expr__))
+
+    @abc.abstractmethod
+    def return_type(self, scope: pa.DataType) -> pa.DataType:
+        """Must return the return type of the UDF given the input scope type.
+
+        All expressions in Spiral must return nullable (Arrow default) types,
+        including nested structs, meaning that all fields in structs must also be nullable,
+        and if those fields are structs, their fields must also be nullable, and so on.
+        """
+        ...
+
+    @abc.abstractmethod
+    def invoke(self, scope: pa.Array) -> pa.Array:
+        """Must implement the UDF logic given the input scope array."""
+        ...

spiral/grpc_.py

@@ -0,0 +1,32 @@
+from collections.abc import AsyncIterator, Awaitable, Callable
+from typing import TypeVar
+
+R = TypeVar("R")
+T = TypeVar("T")
+
+
+async def paged(stub_fn: Callable[[R], Awaitable[T]], request: R, page_size: int = None) -> AsyncIterator[T]:
+    """Page through a gRPC paged API.
+
+    Assumes fields exist as per https://cloud.google.com/apis/design/design_patterns#list_pagination
+    """
+    next_page_token: str | None = None
+    while True:
+        request.page_size = page_size
+        request.page_token = next_page_token
+        res = await stub_fn(request)
+        if not res.next_page_token:
+            # No more items
+            yield res
+            break
+
+        next_page_token = res.next_page_token
+        yield res
+
+
+async def paged_items(
+    stub_fn: Callable[[R], Awaitable[T]], request: R, collection_name: str, page_size: int = None
+) -> AsyncIterator:
+    async for page in paged(stub_fn, request, page_size=page_size):
+        for item in getattr(page, collection_name):
+            yield item

spiral/iceberg.py

@@ -0,0 +1,31 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+
+    from spiral.client import Spiral
+
+
+class Iceberg:
+    """
+    Apache Iceberg is a powerful open-source table format designed for high-performance data lakes.
+    Iceberg brings reliability, scalability, and advanced features like time travel, schema evolution,
+    and ACID transactions to your warehouse.
+    """
+
+    def __init__(self, spiral: "Spiral"):
+        self._spiral = spiral
+        self._api = self._spiral.api
+
+    def catalog(self) -> "Catalog":
+        """Open the Iceberg catalog."""
+        from pyiceberg.catalog import load_catalog
+
+        return load_catalog(
+            "default",
+            **{
+                "type": "rest",
+                "uri": self._spiral.config.server_url + "/iceberg",
+                "token": self._spiral.authn.token().expose_secret(),
+            },
+        )

spiral/iterable_dataset.py

@@ -0,0 +1,106 @@
+from collections.abc import Callable, Iterator
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+
+if TYPE_CHECKING:
+    import datasets.iterable_dataset as hf  # noqa
+    import streaming  # noqa
+    import torch.utils.data as torchdata  # noqa
+
+
+def _hf_compatible_schema(schema: pa.Schema) -> pa.Schema:
+    """
+    Replace string-view and binary-view columns in the schema with strings/binary.
+    Recursively handles nested types (struct, list, etc).
+    We use this converted schema as Features in the returned Dataset.
+    Remove this method once we have https://github.com/huggingface/datasets/pull/7718
+    """
+
+    def _convert_type(dtype: pa.DataType) -> pa.DataType:
+        if dtype == pa.string_view():
+            return pa.string()
+        elif dtype == pa.binary_view():
+            return pa.binary()
+        elif pa.types.is_struct(dtype):
+            new_fields = [
+                pa.field(field.name, _convert_type(field.type), nullable=field.nullable, metadata=field.metadata)
+                for field in dtype
+            ]
+            return pa.struct(new_fields)
+        elif pa.types.is_list(dtype):
+            return pa.list_(_convert_type(dtype.value_type))
+        elif pa.types.is_large_list(dtype):
+            return pa.large_list(_convert_type(dtype.value_type))
+        elif pa.types.is_fixed_size_list(dtype):
+            return pa.list_(_convert_type(dtype.value_type), dtype.list_size)
+        elif pa.types.is_map(dtype):
+            return pa.map_(_convert_type(dtype.key_type), _convert_type(dtype.item_type))
+        else:
+            return dtype
+
+    new_fields = []
+    for field in schema:
+        new_type = _convert_type(field.type)
+        new_fields.append(pa.field(field.name, new_type, nullable=field.nullable, metadata=field.metadata))
+
+    return pa.schema(new_fields)
+
+
+def to_iterable_dataset(stream: pa.RecordBatchReader) -> "hf.IterableDataset":
+    from datasets import DatasetInfo, Features
+    from datasets.builder import ArrowExamplesIterable
+    from datasets.iterable_dataset import IterableDataset
+
+    def _generate_tables(**kwargs) -> Iterator[tuple[int, pa.Table]]:
+        # This key is unused when training with IterableDataset.
+        # Default implementation returns shard id, e.g. parquet row group id.
+        for i, rb in enumerate(stream):
+            yield i, pa.Table.from_batches([rb], stream.schema)
+
+    # TODO(marko): This is temporary until we stop returning IterableDataset from this function.
+    class _IterableDataset(IterableDataset):
+        # Diff with datasets.iterable_dataset.IterableDataset:
+        # - Removes torch handling which attempts to handle worker processes.
+        # - Assumes arrow iterator.
+        def __iter__(self):
+            from datasets.formatting import get_formatter
+
+            prepared_ex_iterable = self._prepare_ex_iterable_for_iteration()
+            if self._formatting and (prepared_ex_iterable.iter_arrow or self._formatting.is_table):
+                formatter = get_formatter(self._formatting.format_type, features=self.features)
+                iterator = prepared_ex_iterable.iter_arrow()
+                for key, pa_table in iterator:
+                    yield formatter.format_row(pa_table)
+                return
+
+            for key, example in prepared_ex_iterable:
+                # no need to format thanks to FormattedExamplesIterable
+                yield example
+
+        def map(self, *args, **kwargs):
+            # Map constructs a new IterableDataset, so we need to "patch" it
+            base = super().map(*args, **kwargs)
+            if isinstance(base, IterableDataset):
+                # Patch __iter__ to avoid torch handling
+                base.__class__ = _IterableDataset  # type: ignore
+            return base
+
+    class _ArrowExamplesIterable(ArrowExamplesIterable):
+        def __init__(self, generate_tables_fn: Callable[..., Iterator[tuple[int, pa.Table]]], features: Features):
+            # NOTE: generate_tables_fn type annotations are wrong, return type must be an iterable of tuples.
+            super().__init__(generate_tables_fn, kwargs={})  # type: ignore
+            self._features = features
+
+        @property
+        def is_typed(self) -> bool:
+            return True
+
+        @property
+        def features(self) -> Features:
+            return self._features
+
+    target_features = Features.from_arrow_schema(_hf_compatible_schema(stream.schema))
+    ex_iterable = _ArrowExamplesIterable(_generate_tables, target_features)
+    info = DatasetInfo(features=target_features)
+    return _IterableDataset(ex_iterable=ex_iterable, info=info)

spiral/key_space_index.py

@@ -0,0 +1,44 @@
+from spiral.core.client import KeySpaceIndex as CoreKeySpaceIndex
+from spiral.expressions import Expr
+from spiral.types_ import Timestamp
+
+
+class KeySpaceIndex:
+    """
+    KeysIndex represents an optionally materialized key space, defined by a projection and a filter over a table.
+    It can be used to efficiently and precisely shard the table for parallel processing or distributed training.
+
+    An index is defined by:
+    - A granularity that defines the target size of key ranges in the index.
+        IMPORTANT: Actual key ranges may be smaller, but will not exceed twice the granularity.
+    - A projection expression that defines which columns are included in the resulting key space.
+    - An optional filter expression that defines which rows are included in the index.
+    """
+
+    def __init__(self, core: CoreKeySpaceIndex, *, name: str | None = None):
+        self.core = core
+        self._name = name
+
+    @property
+    def index_id(self) -> str:
+        return self.core.id
+
+    @property
+    def table_id(self) -> str:
+        return self.core.table_id
+
+    @property
+    def name(self) -> str:
+        return self._name or self.index_id
+
+    @property
+    def asof(self) -> Timestamp:
+        return self.core.asof
+
+    @property
+    def projection(self) -> Expr:
+        return Expr(self.core.projection)
+
+    @property
+    def filter(self) -> Expr | None:
+        return Expr(self.core.filter) if self.core.filter is not None else None