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

spiral/client.py

@@ -1,24 +1,29 @@
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING
 
 import jwt
+import pyarrow as pa
 
 from spiral.api import SpiralAPI
 from spiral.api.projects import CreateProjectRequest, CreateProjectResponse
+from spiral.core.client import Operations
 from spiral.core.client import Spiral as CoreSpiral
+from spiral.datetime_ import timestamp_micros
+from spiral.expressions import ExprLike
+from spiral.scan import Scan
 from spiral.settings import Settings, settings
 
 if TYPE_CHECKING:
     from spiral.iceberg import Iceberg
+    from spiral.key_space_index import KeySpaceIndex
     from spiral.project import Project
+    from spiral.table import Table
+    from spiral.text_index import TextIndex
 
 
 class Spiral:
     def __init__(self, config: Settings | None = None):
         self._config = config or settings()
-        self._api = self._config.api
-        self._core = CoreSpiral(
-            api_url=self._config.spiraldb.uri, spfs_url=self._config.spfs.uri, authn=self._config.authn
-        )
         self._org = None
 
     @property
@@ -27,7 +32,11 @@ class Spiral:
 
     @property
     def api(self) -> SpiralAPI:
-        return self._api
+        return self._config.api
+
+    @property
+    def _core(self) -> CoreSpiral:
+        return self._config.core
 
     @property
     def organization(self) -> str:
@@ -45,7 +54,7 @@ class Spiral:
         """List project IDs."""
         from .project import Project
 
-        return [Project(self, id=p.id, name=p.name) for p in self.api.project.list()]
+        return [Project(self, project_id=p.id, name=p.name) for p in self.api.project.list()]
 
     def create_project(
         self,
@@ -64,7 +73,102 @@ class Spiral:
         from spiral.project import Project
 
         # We avoid an API call since we'd just be fetching a human-readable name. Seems a waste in most cases.
-        return Project(self, id=project_id, name=project_id)
+        return Project(self, project_id=project_id, name=project_id)
+
+    def table(self, table_id: str) -> "Table":
+        """Open a table using an ID."""
+        from spiral.table import Table
+
+        return Table(self, self._core.table(table_id))
+
+    def text_index(self, index_id: str) -> "TextIndex":
+        """Open a text index using an ID."""
+        from spiral.text_index import TextIndex
+
+        return TextIndex(self._core.text_index(index_id))
+
+    def key_space_index(self, index_id: str) -> "KeySpaceIndex":
+        """Open a key space index using an ID."""
+        from spiral.key_space_index import KeySpaceIndex
+
+        return KeySpaceIndex(self._core.key_space_index(index_id))
+
+    def scan(
+        self,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        asof: datetime | int | None = None,
+        exclude_keys: bool = False,
+    ) -> Scan:
+        """Starts a read transaction on the Spiral.
+
+        Args:
+            projections: a set of expressions that return struct arrays.
+            where: a query expression to apply to the data.
+            asof: only data written before the given timestamp will be returned, caveats around compaction.
+            exclude_keys: whether to exclude the key columns in the scan result, defaults to False.
+                Note that if a projection includes a key column, it will be included in the result.
+        """
+        from spiral import expressions as se
+
+        if isinstance(asof, datetime):
+            asof = timestamp_micros(asof)
+
+        # Combine all projections into a single struct.
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return Scan(
+            self._core.scan(
+                projection.__expr__,
+                filter=where.__expr__ if where else None,
+                asof=asof,
+                exclude_keys=exclude_keys,
+            ),
+        )
+
+    # TODO(marko): This should be query, and search should be query + scan.
+    def search(
+        self,
+        top_k: int,
+        *rank_by: ExprLike,
+        filters: ExprLike | None = None,
+        freshness_window: timedelta | None = None,
+    ) -> pa.RecordBatchReader:
+        """Queries the index with the given rank by and filters clauses. Returns a steam of scored keys.
+
+        Args:
+            top_k: The number of top results to return.
+            rank_by: Rank by expressions are combined for scoring.
+                See `se.text.find` and `se.text.boost` for scoring expressions.
+            filters: The `filters` expression is used to filter the results.
+                It must return a boolean value and use only conjunctions (ANDs). Expressions in filters
+                statement are considered either a `must` or `must_not` clause in search terminology.
+            freshness_window: If provided, the index will not be refreshed if its freshness does not exceed this window.
+        """
+        from spiral import expressions as se
+
+        if not rank_by:
+            raise ValueError("At least one rank by expression is required.")
+        rank_by = se.or_(*rank_by)
+        if filters is not None:
+            filters = se.lift(filters)
+
+        if freshness_window is None:
+            freshness_window = timedelta(seconds=0)
+        freshness_window_s = int(freshness_window.total_seconds())
+
+        return self._core.search(
+            top_k=top_k,
+            rank_by=rank_by.__expr__,
+            filters=filters.__expr__ if filters else None,
+            freshness_window_s=freshness_window_s,
+        )
+
+    def _ops(self) -> Operations:
+        """Access maintenance operations."""
+        return self._core._ops(format=settings().file_format)
 
     @property
     def iceberg(self) -> "Iceberg":
@@ -72,7 +176,6 @@ class Spiral:
         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.
-
         """
         from spiral.iceberg import Iceberg
 

spiral/core/authn/__init__.pyi

@@ -0,0 +1,27 @@
+from spiral.api.types import OrgId
+
+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."""
+        ...
+
+    def logout(self) -> None:
+        """Logout from the device authentication session."""
+        ...

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,25 +67,29 @@ 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_scan(self, column_group: ColumnGroup) -> ColumnGroupScan: ...
-    def key_space_scan(self, table_id: str) -> KeySpaceScan: ...
     def metrics(self) -> dict[str, Any]: ...
+    def _prepare_shard(
+        self,
+        output_path: str,
+        key_range: KeyRange,
+        expected_cardinality: int | None = None,
+    ) -> None: ...
 
-class KeySpaceScan:
+class KeySpaceState:
     manifest: FragmentManifest
 
     def key_schema(self) -> Schema: ...
 
-class ColumnGroupScan:
+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_scan(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_scan(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)}",