pyspiral 0.6.6__cp312-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

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

spiral/project.py

@@ -0,0 +1,199 @@
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+
+from spiral.api.projects import KeySpaceIndexResource, TableResource, TextIndexResource
+from spiral.core.table.spec import Schema
+from spiral.expressions import ExprLike
+from spiral.key_space_index import KeySpaceIndex
+from spiral.table import Table
+from spiral.text_index import TextIndex
+from spiral.types_ import Uri
+
+if TYPE_CHECKING:
+    from spiral.client import Spiral
+
+
+class Project:
+    def __init__(self, spiral: "Spiral", project_id: str, name: str | None = None):
+        self._spiral = spiral
+        self._id = project_id
+        self._name = name
+
+    def __str__(self):
+        return self._id
+
+    def __repr__(self):
+        return f"Project(id={self._id}{', name=' + self._name if self._name else ''})"
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def name(self) -> str:
+        return self._name or self._id
+
+    def list_tables(self) -> list[TableResource]:
+        return list(self._spiral.api.project.list_tables(self._id))
+
+    def list_text_indexes(self) -> list[TextIndexResource]:
+        return list(self._spiral.api.project.list_text_indexes(self._id))
+
+    def list_key_space_indexes(self) -> list[KeySpaceIndexResource]:
+        return list(self._spiral.api.project.list_key_space_indexes(self._id))
+
+    def table(self, identifier: str) -> Table:
+        """Open a table with a `dataset.table` identifier, or `table` name using the `default` dataset."""
+        dataset, table = self._parse_table_identifier(identifier)
+
+        res = list(self._spiral.api.project.list_tables(project_id=self._id, dataset=dataset, table=table))
+        if len(res) == 0:
+            raise ValueError(f"Table not found: {self._id}.{dataset}.{table}")
+        res = res[0]
+
+        return Table(
+            self._spiral, self._spiral._core.table(res.id), identifier=f"{res.project_id}.{res.dataset}.{res.table}"
+        )
+
+    def create_table(
+        self,
+        identifier: str,
+        *,
+        key_schema: pa.Schema | Any,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> Table:
+        """Create a new table in the project.
+
+        Args:
+            identifier: The table identifier, in the form `dataset.table` or `table`.
+            key_schema: The schema of the table's keys.
+            root_uri: The root URI for the table.
+            exist_ok: If True, do not raise an error if the table already exists.
+        """
+        dataset, table = self._parse_table_identifier(identifier)
+
+        if not isinstance(key_schema, pa.Schema):
+            key_schema = pa.schema(key_schema)
+        key_schema = Schema.from_arrow(key_schema)
+
+        core_table = self._spiral._core.create_table(
+            project_id=self._id,
+            dataset=dataset,
+            table=table,
+            key_schema=key_schema,
+            root_uri=root_uri,
+            exist_ok=exist_ok,
+        )
+
+        return Table(self._spiral, core_table, identifier=f"{self._id}.{dataset}.{table}")
+
+    def _parse_table_identifier(self, identifier: str) -> tuple[str, str]:
+        parts = identifier.split(".")
+        if len(parts) == 1:
+            return "default", parts[0]
+        elif len(parts) == 2:
+            return parts[0], parts[1]
+        else:
+            raise ValueError(f"Invalid table identifier: {self._id}.{identifier}")
+
+    def text_index(self, name: str) -> TextIndex:
+        """Returns the index with the given name."""
+        res = list(self._spiral.api.project.list_text_indexes(project_id=self._id, name=name))
+        if len(res) == 0:
+            raise ValueError(f"Index not found: {name}")
+        res = res[0]
+
+        return TextIndex(self._spiral._core.text_index(res.id), name=name)
+
+    def create_text_index(
+        self,
+        name: str,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> TextIndex:
+        """Creates a text index over the table projection.
+
+        See `se.text.field` for how to create and configure indexable fields.
+
+        Args:
+            name: The index name. Must be unique within the project.
+            projections: At least one projection expression is required.
+                All projections must reference the same table.
+            where: An optional filter expression to apply to the index.
+            root_uri: The root URI for the index.
+            exist_ok: If True, do not raise an error if the index already exists.
+        """
+        from spiral import expressions as se
+
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        core_index = self._spiral._core.create_text_index(
+            project_id=self._id,
+            name=name,
+            projection=projection.__expr__,
+            filter=where.__expr__ if where else None,
+            root_uri=root_uri,
+            # TODO(marko): Validate that if an index exists, it's the same?
+            exist_ok=exist_ok,
+        )
+
+        return TextIndex(core_index, name=name)
+
+    def key_space_index(self, name: str) -> KeySpaceIndex:
+        """Returns the index with the given name."""
+        res = list(self._spiral.api.project.list_key_space_indexes(project_id=self._id, name=name))
+        if len(res) == 0:
+            raise ValueError(f"Index not found: {name}")
+        res = res[0]
+
+        return KeySpaceIndex(self._spiral._core.key_space_index(res.id), name=name)
+
+    def create_key_space_index(
+        self,
+        name: str,
+        granularity: int,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> KeySpaceIndex:
+        """Creates a key space index over the table projection.
+
+        Args:
+            name: The index name. Must be unique within the project.
+            granularity: The granularity at which to store keys, i.e. the size of desired key ranges.
+                The key ranges will not be greater than 2x the granularity, but may be smaller.
+            projections: At least one projection expression is required.
+                All projections must reference the same table.
+            where: An optional filter expression to apply to the index.
+            root_uri: The root URI for the index.
+            exist_ok: If True, do not raise an error if the index already exists.
+        """
+        from spiral import expressions as se
+
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        core_index = self._spiral._core.create_key_space_index(
+            project_id=self._id,
+            name=name,
+            granularity=granularity,
+            projection=projection.__expr__,
+            filter=where.__expr__ if where else None,
+            root_uri=root_uri,
+            # TODO(marko): Validate that if an index exists, it's the same?
+            exist_ok=exist_ok,
+        )
+
+        return KeySpaceIndex(core_index, name=name)