pyspiral 0.5.0__cp310-abi3-macosx_11_0_arm64.whl → 0.6.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/core/client/__init__.pyi

@@ -1,34 +1,11 @@
-from spiral.api.types import DatasetName, IndexName, OrgId, ProjectId, RootUri, TableName
-from spiral.core.index import SearchScan, TextIndex
-from spiral.core.table import Table, TableMaintenance, TableScan, TableSnapshot, TableTransaction
-from spiral.core.table.spec import Schema
-from spiral.expressions import Expr
-
-class Token:
-    def __init__(self, value: str): ...
-    def expose_secret(self) -> str: ...
-
-class Authn:
-    @staticmethod
-    def from_token(token: Token) -> Authn: ...
-    @staticmethod
-    def from_fallback() -> Authn: ...
-    @staticmethod
-    def from_device() -> Authn: ...
-    def token(self) -> Token | None: ...
-
-class DeviceCodeAuth:
-    @staticmethod
-    def default() -> DeviceCodeAuth:
-        """Return the static device code instance."""
-        ...
-    def authenticate(self, force: bool = False, org_id: OrgId | None = None) -> Token:
-        """Authenticate using device code flow."""
-        ...
+from typing import Any, Literal
 
-    def logout(self) -> None:
-        """Logout from the device authentication session."""
-        ...
+import pyarrow as pa
+from spiral.api.types import DatasetName, IndexName, ProjectId, RootUri, TableName
+from spiral.core.authn import Authn
+from spiral.core.table import ColumnGroupState, KeyRange, KeySpaceState, Scan, Snapshot, Table, Transaction
+from spiral.core.table.spec import ColumnGroup, Schema
+from spiral.expressions import Expr
 
 class Spiral:
     """A client for Spiral database"""
@@ -43,6 +20,39 @@ class Spiral:
     def authn(self) -> Authn:
         """Get the current authentication context."""
         ...
+
+    def scan(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        exclude_keys: bool = False,
+    ) -> Scan:
+        """Construct a table scan."""
+        ...
+
+    def transaction(self, table: Table, format: str | None = None) -> Transaction:
+        """Being a table transaction."""
+        ...
+
+    def search(
+        self,
+        top_k: int,
+        rank_by: Expr,
+        *,
+        filters: Expr | None = None,
+        freshness_window_s: int | None = None,
+    ) -> pa.RecordBatchReader:
+        """Search an index.
+
+        Searching an index returns a stream of record batches that match table's key schema + float score column.
+        """
+        ...
+
+    def table(self, table_id: str) -> Table:
+        """Get a table."""
+        ...
+
     def create_table(
         self,
         project_id: ProjectId,
@@ -56,63 +66,125 @@ class Spiral:
         """Create a new table in the specified project."""
         ...
 
-    def get_table(self, table_id: str) -> Table:
-        """Get and open table."""
-
-    def open_table(self, table_id: str, key_schema: Schema, root_uri: RootUri) -> Table:
-        """Open a table. This does not make any network calls."""
+    def text_index(self, index_id: str) -> TextIndex:
+        """Get a text index."""
         ...
 
-    def open_table_scan(
+    def create_text_index(
         self,
+        project_id: ProjectId,
+        name: IndexName,
         projection: Expr,
         filter: Expr | None = None,
-        asof: int | None = None,
-        exclude_keys: bool = False,
-    ) -> TableScan:
-        """Construct a table scan."""
+        *,
+        root_uri: RootUri | None = None,
+        exist_ok: bool = False,
+    ) -> TextIndex:
+        """Create a new index in the specified project."""
         ...
 
-    def open_transaction(self, table: Table, format: str | None = None) -> TableTransaction:
-        """Being transaction."""
+    def key_space_index(self, index_id: str) -> KeySpaceIndex:
+        """Get a key space index."""
         ...
 
-    def open_maintenance(self, table: Table, format: str | None = None) -> TableMaintenance:
-        """Access maintenance operations for a table."""
-        ...
-    def create_text_index(
+    def create_key_space_index(
         self,
         project_id: ProjectId,
         name: IndexName,
+        granularity: int,
         projection: Expr,
         filter: Expr | None = None,
         *,
         root_uri: RootUri | None = None,
         exist_ok: bool = False,
-    ) -> TextIndex:
-        """Create a new index in the specified project."""
+    ) -> KeySpaceIndex:
+        """Create a new key space index in the specified project."""
         ...
 
-    def get_text_index(self, index_id: str) -> TextIndex:
-        """Get a text-based index."""
+    def _ops(self, *, format: str | None = None) -> Operations:
+        """Access maintenance operations.
+
+        IMPORTANT: This API is internal and is currently exposed for development & testing.
+            Maintenance operations are run by SpiralDB.
+        """
         ...
 
-    def open_search_scan(
+class TextIndex:
+    id: str
+
+class KeySpaceIndex:
+    id: str
+    table_id: str
+    granularity: int
+    projection: Expr
+    filter: Expr
+    asof: int
+
+class Shard:
+    key_range: KeyRange
+    cardinality: int
+
+class Operations:
+    def flush_wal(self, table: Table) -> None:
+        """
+        Flush the write-ahead log of the table.
+        """
+        ...
+    def compact_key_space(
         self,
-        rank_by: Expr,
-        top_k: int,
-        # NOTE(marko): Required for now.
-        freshness_window_s: int,
         *,
-        filter: Expr | None = None,
-    ) -> SearchScan:
-        """Query an index."""
+        table: Table,
+        mode: Literal["plan", "read", "write"] | None = None,
+        partition_bytes_min: int | None = None,
+    ):
+        """
+        Compact the key space of the table.
+        """
         ...
-
-    def _sync_snapshot(self, index_id: str, snapshot: TableSnapshot) -> None:
+    def compact_column_group(
+        self,
+        table: Table,
+        column_group: ColumnGroup,
+        *,
+        mode: Literal["plan", "read", "write"] | None = None,
+        partition_bytes_min: int | None = None,
+    ):
         """
-        Synchronize an index with a table snapshot.
-
-        IMPORTANT: This is only exposed for testing purposes and should not be used.
+        Compact a column group in the table.
+        """
+        ...
+    def update_text_index(self, index: TextIndex, snapshot: Snapshot) -> None:
+        """
+        Index table changes up to the given snapshot.
+        """
+        ...
+    def update_key_space_index(self, index: KeySpaceIndex, snapshot: Snapshot) -> None:
+        """
+        Index table changes up to the given snapshot.
+        """
+        ...
+    def key_space_state(self, table: Table, *, asof: int | None = None) -> KeySpaceState:
+        """
+        The key space state for the table.
+        """
+        ...
+    def column_group_state(
+        self, table: Table, key_space_state: KeySpaceState, column_group: ColumnGroup, *, asof: int | None = None
+    ) -> ColumnGroupState:
+        """
+        The state the column group of the table.
+        """
+        ...
+    def column_groups_states(
+        self, table: Table, key_space_state: KeySpaceState, *, asof: int | None = None
+    ) -> list[ColumnGroupState]:
+        """
+        The state of each column group of the table.
+        """
+        ...
+    def compute_shards(self, index: KeySpaceIndex) -> list[Shard]:
+        """
+        Compute the scan shards from a key space index.
         """
         ...
+    def metrics(self) -> dict[str, Any]: ...

spiral/core/table/__init__.pyi

@@ -1,4 +1,4 @@
-from typing import Any, Literal
+from typing import Any
 
 import pyarrow as pa
 from spiral.expressions import Expr
@@ -42,22 +42,24 @@ class Table:
 
     def get_wal(self, *, asof: int | None) -> WriteAheadLog: ...
     def get_schema(self, *, asof: int | None) -> Schema: ...
-    def get_snapshot(self, *, asof: int | None) -> TableSnapshot: ...
+    def get_snapshot(self, *, asof: int | None) -> Snapshot: ...
 
-class TableSnapshot:
+class Snapshot:
     """A snapshot of a table at a specific point in time."""
 
     asof: int
     table: Table
     wal: WriteAheadLog
 
-class TableScan:
+class Scan:
     def key_schema(self) -> Schema: ...
     def schema(self) -> Schema: ...
     def is_empty(self) -> bool: ...
-    def split(self) -> list[KeyRange]: ...
+    def splits(self) -> list[KeyRange]: ...
     def table_ids(self) -> list[str]: ...
     def column_groups(self) -> list[ColumnGroup]: ...
+    def column_group_state(self, column_group: ColumnGroup) -> ColumnGroupState: ...
+    def key_space_state(self, table_id: str) -> KeySpaceState: ...
     def to_record_batches(
         self,
         key_table: pa.Table | pa.RecordBatch | None = None,
@@ -65,13 +67,16 @@ class TableScan:
     ) -> pa.RecordBatchReader: ...
     def to_shuffled_record_batches(
         self,
+        strategy: ShuffleStrategy | None = None,
         batch_readahead: int | None = None,
-        shuffle_buffer_size: int | None = None,
-        shuffle_pool_num_rows: int | None = None,
     ) -> pa.RecordBatchReader: ...
-    def column_group_state(self, column_group: ColumnGroup) -> ColumnGroupState: ...
-    def key_space_state(self, table_id: str) -> KeySpaceState: ...
     def metrics(self) -> dict[str, Any]: ...
+    def _prepare_shard(
+        self,
+        output_path: str,
+        key_range: KeyRange,
+        expected_cardinality: int | None = None,
+    ) -> None: ...
 
 class KeySpaceState:
     manifest: FragmentManifest
@@ -80,10 +85,11 @@ class KeySpaceState:
 
 class ColumnGroupState:
     manifest: FragmentManifest
+    column_group: ColumnGroup
 
     def schema(self) -> Schema: ...
 
-class TableTransaction:
+class Transaction:
     status: str
 
     def write(self, expr: Expr, *, partition_size_bytes: int | None = None): ...
@@ -91,19 +97,23 @@ class TableTransaction:
     def abort(self): ...
     def metrics(self) -> dict[str, Any]: ...
 
-class TableMaintenance:
-    def flush_wal(self): ...
-    def compact_key_space(
-        self,
-        *,
-        mode: Literal["plan", "read", "write"] | None = None,
-        partition_bytes_min: int | None = None,
-    ): ...
-    def compact_column_group(
+class ShuffleStrategy:
+    # Results are buffered in a pool of `buffer_size` rows and shuffled again.
+    shuffle_buffer_size: int
+
+    # All randomness is derived from this seed. If None, a random seed is generated from the OS.
+    seed: int | None
+
+    # `approximate_batch_size` controls the maximum approximate size of each shard. Shards that
+    # are larger than this size are further split assuming uniform distribution of keys. Note
+    # that this is a best-effort and can be widely off. The purpose of this is to improve
+    # shuffling, rather than to support sharding. If not present, splits derived from the table
+    # are used in the attempt to minimize wasted reads.
+    approximate_buffer_size: int | None
+
+    def __init__(
         self,
-        column_group: ColumnGroup,
-        *,
-        mode: Literal["plan", "read", "write"] | None = None,
-        partition_bytes_min: int | None = None,
+        seed: int | None = None,
+        shard_size: int | None = None,
+        buffer_size: int | None = None,
     ): ...
-    def metrics(self) -> dict[str, Any]: ...

spiral/core/table/metastore/__init__.pyi

@@ -56,7 +56,3 @@ class PyMetastore:
         authn: Authn,
     ) -> PyMetastore:
         """Construct a PyMetastore backed by an HTTP metastore service."""
-
-    @staticmethod
-    def test(table_id: str, root_uri: str, key_schema: Schema) -> PyMetastore:
-        """Construct a PyMetastore backed by an in-memory mock metastore service."""

spiral/core/table/spec/__init__.pyi

@@ -101,8 +101,6 @@ class FragmentLevel:
         ...
 
 class Key:
-    key: bytes
-
     def __init__(self, key: bytes): ...
     def __bytes__(self): ...
     def step(self) -> Key:

spiral/{tables/dataset.py → dataset.py}

@@ -4,14 +4,15 @@ import pyarrow as pa
 import pyarrow.compute as pc
 import pyarrow.dataset as ds
 
-from spiral.tables import Scan, Snapshot
+from spiral.scan import Scan
+from spiral.snapshot import Snapshot
 
 
-class TableDataset(ds.Dataset):
+class Dataset(ds.Dataset):
     def __init__(self, snapshot: Snapshot):
         self._snapshot = snapshot
         self._table = snapshot.table
-        self._schema: pa.Schema = self._snapshot._snapshot.table.get_schema(asof=self._snapshot.asof).to_arrow()
+        self._schema: pa.Schema = self._snapshot.schema().to_arrow()
 
         # We don't actually initialize a Dataset, we just implement enough of the API
         # to fool both DuckDB and Polars.
@@ -43,7 +44,7 @@ class TableDataset(ds.Dataset):
             memory_pool,
         ).count_rows()
 
-    def filter(self, expression: pc.Expression) -> "TableDataset":
+    def filter(self, expression: pc.Expression) -> "Dataset":
         raise NotImplementedError("filter not implemented")
 
     def get_fragments(self, filter: pc.Expression | None = None):
@@ -89,7 +90,7 @@ class TableDataset(ds.Dataset):
     def join_asof(self, right_dataset, on, by, tolerance, right_on=None, right_by=None):
         raise NotImplementedError("join_asof not implemented")
 
-    def replace_schema(self, schema: pa.Schema) -> "TableDataset":
+    def replace_schema(self, schema: pa.Schema) -> "Dataset":
         raise NotImplementedError("replace_schema not implemented")
 
     def scanner(
@@ -112,13 +113,18 @@ class TableDataset(ds.Dataset):
             )
 
         scan = (
-            self._snapshot.scan(
+            self._table.spiral.scan(
                 {c: self._table[c] for c in columns},
                 where=filter,
+                asof=self._snapshot.asof,
                 exclude_keys=True,
             )
             if columns
-            else self._snapshot.scan(where=filter)
+            else self._table.spiral.scan(
+                self._table,
+                where=filter,
+                asof=self._snapshot.asof,
+            )
         )
         self._last_scan = scan
 

spiral/{tables/debug → debug}/manifests.py

@@ -1,23 +1,34 @@
+from collections.abc import Iterable
+
 from spiral import datetime_
-from spiral.core.table import TableScan
+from spiral.core.table import Scan
 from spiral.core.table.manifests import FragmentManifest
-from spiral.tables.debug.metrics import _format_bytes
+from spiral.core.table.spec import ColumnGroup
+from spiral.debug.metrics import _format_bytes
 
 
-def display_manifests(scan: TableScan):
+def display_scan_manifests(scan: Scan):
     """Display all manifests in a scan."""
     if len(scan.table_ids()) != 1:
         raise NotImplementedError("Multiple table scans are not supported.")
     table_id = scan.table_ids()[0]
+    key_space_manifest = scan.key_space_state(table_id).manifest
+    column_group_manifests = (
+        (column_group, scan.column_group_state(column_group).manifest) for column_group in scan.column_groups()
+    )
+
+    display_manifests(key_space_manifest, column_group_manifests)
+
 
-    key_space_manifest: FragmentManifest = scan.key_space_state(table_id).manifest
+def display_manifests(
+    key_space_manifest: FragmentManifest, column_group_manifests: Iterable[tuple[ColumnGroup, FragmentManifest]]
+):
     _table_of_fragments(
         key_space_manifest,
         title="Key Space manifest",
     )
 
-    for column_group in scan.column_groups():
-        column_group_manifest: FragmentManifest = scan.column_group_state(column_group).manifest
+    for column_group, column_group_manifest in column_group_manifests:
         _table_of_fragments(
             column_group_manifest,
             title=f"Column Group manifest for {str(column_group)}",

spiral/{tables/debug → debug}/scan.py

@@ -1,12 +1,12 @@
 from datetime import datetime
 
-from spiral.core.table import TableScan
+from spiral.core.table import Scan
 from spiral.core.table.manifests import FragmentFile, FragmentManifest
 from spiral.core.table.spec import Key
 from spiral.types_ import Timestamp
 
 
-def show_scan(scan: TableScan):
+def show_scan(scan: Scan):
     """Displays a scan in a way that is useful for debugging."""
     table_ids = scan.table_ids()
     if len(table_ids) > 1:
@@ -14,7 +14,7 @@ def show_scan(scan: TableScan):
     table_id = table_ids[0]
     column_groups = scan.column_groups()
 
-    splits = scan.split()
+    splits = scan.splits()
     key_space_state = scan.key_space_state(table_id)
 
     # Collect all key bounds from all manifests. This makes sure all visualizations are aligned.

spiral/expressions/base.py

@@ -132,7 +132,7 @@ class Expr:
                 packed = packed.select(exclude=exclude)
             return packed
 
-        if not paths:
+        if not paths and not exclude:
             return self
 
         return se.select(self, names=list(paths), exclude=exclude)
@@ -145,5 +145,5 @@ class Expr:
         return Expr(_lib.expr.binary(op, self.__expr__, rhs.__expr__))
 
 
-ScalarLike: TypeAlias = bool | int | float | str | list | datetime.datetime | None
-ExprLike: TypeAlias = Expr | dict | ScalarLike
+ScalarLike: TypeAlias = bool | int | float | str | list["ScalarLike"] | datetime.datetime | None
+ExprLike: TypeAlias = Expr | dict["ExprLike", "ExprLike"] | ScalarLike

spiral/expressions/udf.py

@@ -38,7 +38,7 @@ class RefUDF(BaseUDF):
         super().__init__(_lib.expr.udf.create(name, return_type=self.return_type, invoke=self.invoke, scope="ref"))
 
     @abc.abstractmethod
-    def invoke(self, fp: _lib.FileObject, *input_args: pa.Array) -> pa.Array:
+    def invoke(self, fp, *input_args: pa.Array) -> pa.Array:
         """Invoke the UDF with the given arguments.
 
         NOTE: The first argument is always the ref cell. All array input args will be sliced to the appropriate row.

spiral/{iceberg/client.py → iceberg.py}

@@ -13,10 +13,8 @@ class Iceberg:
     and ACID transactions to your warehouse.
     """
 
-    def __init__(self, spiral: "Spiral", *, project_id: str | None = None):
+    def __init__(self, spiral: "Spiral"):
         self._spiral = spiral
-        self._project_id = project_id
-
         self._api = self._spiral.config.api
 
     def catalog(self) -> "Catalog":

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