pyspiral 0.6.8__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/client.py

@@ -0,0 +1,178 @@
+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._org = None
+
+    @property
+    def config(self) -> Settings:
+        return self._config
+
+    @property
+    def api(self) -> SpiralAPI:
+        return self._config.api
+
+    @property
+    def _core(self) -> CoreSpiral:
+        return self._config.core
+
+    @property
+    def organization(self) -> str:
+        if self._org is None:
+            token = self._config.authn.token()
+            if token is None:
+                raise ValueError("Authentication failed.")
+            token_payload = jwt.decode(token.expose_secret(), options={"verify_signature": False})
+            if "org_id" not in token_payload:
+                raise ValueError("Please create an organization.")
+            self._org = token_payload["org_id"]
+        return self._org
+
+    def list_projects(self) -> list["Project"]:
+        """List project IDs."""
+        from .project import Project
+
+        return [Project(self, project_id=p.id, name=p.name) for p in self.api.project.list()]
+
+    def create_project(
+        self,
+        id_prefix: str | None = None,
+        *,
+        name: str | None = None,
+    ) -> "Project":
+        """Create a project in the current, or given, organization."""
+        from .project import Project
+
+        res: CreateProjectResponse = self.api.project.create(CreateProjectRequest(id_prefix=id_prefix, name=name))
+        return Project(self, res.project.id, name=res.project.name)
+
+    def project(self, project_id: str) -> "Project":
+        """Open an existing project."""
+        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, 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,
+    ) -> 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.
+        """
+        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,
+            ),
+        )
+
+    # 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 stream 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":
+        """
+        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
+
+        return Iceberg(self)

spiral/core/_tools/__init__.pyi

@@ -0,0 +1,5 @@
+from ..table.spec import Schema
+
+def pretty_key(key: bytes, schema: Schema) -> str:
+    """Represent a key in a human-readable way."""
+    ...

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(api_url: str) -> 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

@@ -0,0 +1,237 @@
+from typing import Any, Literal
+
+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"""
+    def __init__(
+        self,
+        api_url: str | None = None,
+        spfs_url: str | None = None,
+        authn: Authn | None = None,
+    ):
+        """Initialize the Spiral client."""
+        ...
+    def authn(self) -> Authn:
+        """Get the current authentication context."""
+        ...
+
+    def scan(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+    ) -> Scan:
+        """Construct a table scan."""
+        ...
+
+    def transaction(self, table: Table, format: str | None = None, retries: int | None = 3) -> 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,
+        dataset: DatasetName,
+        table: TableName,
+        key_schema: Schema,
+        *,
+        root_uri: RootUri | None = None,
+        exist_ok: bool = False,
+    ) -> Table:
+        """Create a new table in the specified project."""
+        ...
+
+    def text_index(self, index_id: str) -> TextIndex:
+        """Get a text index."""
+        ...
+
+    def create_text_index(
+        self,
+        project_id: ProjectId,
+        name: IndexName,
+        projection: Expr,
+        filter: Expr | None = None,
+        *,
+        root_uri: RootUri | None = None,
+        exist_ok: bool = False,
+    ) -> TextIndex:
+        """Create a new index in the specified project."""
+        ...
+
+    def key_space_index(self, index_id: str) -> KeySpaceIndex:
+        """Get a key space 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,
+    ) -> KeySpaceIndex:
+        """Create a new key space index in the specified project."""
+        ...
+
+    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.
+        """
+        ...
+
+class TextIndex:
+    id: str
+
+class KeySpaceIndex:
+    id: str
+    table_id: str
+    granularity: int
+    projection: Expr
+    filter: Expr
+    asof: int
+
+class Shard:
+    """A shard representing a partition of data.
+
+    Attributes:
+        key_range: The key range for this shard.
+        cardinality: The number of rows in this shard, if known.
+    """
+
+    key_range: KeyRange
+    cardinality: int | None
+
+    def __init__(self, key_range: KeyRange, cardinality: int | None): ...
+    def __getnewargs__(self) -> tuple[KeyRange, int | None]: ...
+
+class ShuffleConfig:
+    """Configuration for within-shard sample shuffling.
+
+    This controls how samples are shuffled within a buffer, separate from
+    which shards to read (which is specified as a parameter to the scan).
+
+    Attributes:
+        buffer_size: Size of the buffer pool for shuffling samples.
+        seed: Random seed for reproducibility. If None, uses OS randomness.
+        max_batch_size: Maximum batch size for output chunks. If None,
+            defaults to max(1, buffer_size / 16).
+    """
+
+    buffer_size: int
+    seed: int | None
+    max_batch_size: int | None
+
+    def __init__(
+        self,
+        buffer_size: int,
+        *,
+        seed: int | None = None,
+        max_batch_size: int | None = None,
+    ): ...
+
+class Operations:
+    def flush_wal(self, table: Table, *, keep_latest_s: int | None = None) -> None:
+        """
+        Flush the write-ahead log of the table.
+        """
+        ...
+    def compact_key_space(
+        self,
+        *,
+        table: Table,
+        mode: Literal["plan", "read", "write"] | None = None,
+        partition_bytes_min: int | None = None,
+    ):
+        """
+        Compact the key space of the table.
+        """
+        ...
+    def compact_column_group(
+        self,
+        table: Table,
+        column_group: ColumnGroup,
+        *,
+        mode: Literal["plan", "read", "write"] | None = None,
+        partition_bytes_min: int | None = None,
+    ):
+        """
+        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, snapshot: Snapshot) -> KeySpaceState:
+        """
+        The key space state for the table.
+        """
+        ...
+    def column_group_state(
+        self, snapshot: Snapshot, key_space_state: KeySpaceState, column_group: ColumnGroup
+    ) -> ColumnGroupState:
+        """
+        The state the column group of the table.
+        """
+        ...
+    def column_groups_states(self, snapshot: Snapshot, key_space_state: KeySpaceState) -> 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 prepare_shard(
+        self,
+        output_path: str,
+        scan: Scan,
+        shard: Shard,
+        row_block_size: int = 8192,
+    ) -> None:
+        """
+        Prepare a shard locally. Used for `SpiralStream` integration with `streaming` which requires on-disk shards.
+        """
+        ...
+    def metrics(self) -> dict[str, Any]: ...
+
+def flush_telemetry() -> None:
+    """Flush telemetry data to the configured exporter."""
+    ...

spiral/core/table/__init__.pyi

@@ -0,0 +1,101 @@
+from typing import Any
+
+import pyarrow as pa
+from spiral.core.client import Shard, ShuffleConfig
+
+from .manifests import FragmentManifest
+from .metastore import PyMetastore
+from .spec import ColumnGroup, Key, Schema, WriteAheadLog
+
+class KeyRange:
+    """A right-exclusive range of keys."""
+
+    def __init__(self, *, begin: Key, end: Key): ...
+
+    begin: Key
+    end: Key
+
+    def union(self, other: KeyRange) -> KeyRange: ...
+    def __or__(self, other: KeyRange) -> KeyRange: ...
+    def intersection(self, key_extent: KeyRange) -> KeyRange | None: ...
+    def __and__(self, other: KeyRange) -> KeyRange | None: ...
+    def contains(self, item: Key) -> bool: ...
+    def __contains__(self, item: Key) -> bool: ...
+    def is_disjoint(self, key_range: KeyRange) -> bool:
+        return self.end <= key_range.begin or self.begin >= key_range.end
+
+    @staticmethod
+    def beginning_with(begin: Key) -> KeyRange: ...
+    @staticmethod
+    def ending_with(end: Key) -> KeyRange: ...
+    @staticmethod
+    def full() -> KeyRange: ...
+    def __reduce__(self) -> tuple[type[KeyRange], tuple[Key, Key]]: ...
+
+class Table:
+    def __init__(self, metastore: PyMetastore): ...
+
+    id: str
+    root_uri: str
+    mount_id: str | None
+    key_schema: Schema
+    metastore: PyMetastore
+
+    def get_wal(self, *, asof: int | None) -> WriteAheadLog: ...
+    def get_schema(self, *, asof: int | None) -> Schema: ...
+    def get_snapshot(self, *, asof: int | None) -> Snapshot: ...
+
+class Snapshot:
+    """A snapshot of a table at a specific point in time."""
+
+    asof: int
+    table: Table
+    wal: WriteAheadLog
+
+class Scan:
+    def key_schema(self) -> Schema: ...
+    def schema(self) -> Schema: ...
+    def is_empty(self) -> bool: ...
+    def splits(self) -> list[KeyRange]: ...
+    def shards(self) -> list[Shard]: ...
+    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,
+        batch_readahead: int | None = None,
+    ) -> pa.RecordBatchReader: ...
+    def to_shuffled_record_batches(
+        self,
+        shards: list[Shard] | None = None,
+        shuffle: ShuffleConfig | None = None,
+        max_batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        infinite: bool = False,
+    ) -> pa.RecordBatchReader:
+        # If `infinite` is True, shards are shuffled after exhausted but not before the first pass.
+        # Otherwise, shards are not shuffle and shuffle config is only used for shuffle buffer.
+        ...
+    def metrics(self) -> dict[str, Any]: ...
+
+class KeySpaceState:
+    manifest: FragmentManifest
+
+    def key_schema(self) -> Schema: ...
+
+class ColumnGroupState:
+    manifest: FragmentManifest
+    column_group: ColumnGroup
+
+    def schema(self) -> Schema: ...
+
+class Transaction:
+    status: str
+
+    def write(self, table: pa.RecordBatchReader, *, partition_size_bytes: int | None = None): ...
+    def drop_columns(self, column_paths: list[str]): ...
+    def commit(self): ...
+    def abort(self): ...
+    def metrics(self) -> dict[str, Any]: ...

spiral/core/table/manifests/__init__.pyi

@@ -0,0 +1,35 @@
+import pyarrow as pa
+from spiral.core.table import KeyRange
+from spiral.core.table.spec import FileFormat, FragmentLevel, KeyExtent, KeySpan
+from spiral.types_ import Timestamp
+
+class FragmentManifest:
+    def __len__(self): ...
+    def __getitem__(self, idx: int): ...
+    def to_arrow(self) -> pa.RecordBatchReader: ...
+    @staticmethod
+    def compute_schema() -> pa.Schema: ...
+    @staticmethod
+    def from_fragment(fragment_file: FragmentFile) -> FragmentManifest: ...
+    @staticmethod
+    def from_arrow(reader: pa.RecordBatchReader) -> FragmentManifest: ...
+    @staticmethod
+    def empty() -> FragmentManifest: ...
+
+class FragmentFile:
+    id: str
+    committed_at: Timestamp | None
+    compacted_at: Timestamp | None
+    format: FileFormat
+    format_metadata: bytes | None
+    size_bytes: int
+    column_ids: list[str]
+    level: FragmentLevel
+    # NOTE: Empty for key space files.
+    column_ids: list[str]
+    ks_id: str
+    key_span: KeySpan
+    key_extent: KeyExtent
+
+    @property
+    def key_range(self) -> KeyRange: ...

spiral/core/table/metastore/__init__.pyi

@@ -0,0 +1,58 @@
+"""The SpiralDB metastore API."""
+
+from spiral.core.client import Authn
+from spiral.core.table.spec import ColumnGroup, ColumnGroupMetadata, FileFormat, KeySpaceMetadata, Schema, WriteAheadLog
+from spiral.types_ import Uri
+
+class FileHandle:
+    def __init__(self, *, uri: str, format: FileFormat, spfs_token: str | None): ...
+
+    uri: str
+    format: FileFormat
+    spfs_token: str | None
+
+class FileRef:
+    def __init__(self, *, id: str, file_type: FileType, file_format: FileFormat): ...
+
+    id: str
+    file_type: FileType
+    file_format: FileFormat
+
+    def resolve_uri(self, root_uri: str) -> str:
+        """Resolves the file reference URI given the root URI."""
+
+class FileType:
+    FragmentFile: FileType
+    FragmentManifest: FileType
+    ReferenceFile: FileType
+
+    def __int__(self) -> int:
+        """Returns the protobuf enum int value."""
+
+class PyMetastore:
+    """Rust implementation of the metastore API."""
+
+    @property
+    def table_id(self) -> str: ...
+    @property
+    def root_uri(self) -> Uri: ...
+    @property
+    def key_schema(self) -> Schema: ...
+    def get_wal(self) -> WriteAheadLog:
+        """Return the log for the table."""
+    def get_key_space_metadata(self) -> KeySpaceMetadata:
+        """Return the metadata for the key space."""
+        ...
+    def get_column_group_metadata(self, column_group: ColumnGroup) -> ColumnGroupMetadata:
+        """Return the metadata for a column group."""
+        ...
+
+    @staticmethod
+    def http(
+        table_id: str,
+        root_uri: str,
+        key_schema: Schema,
+        base_url: str,
+        authn: Authn,
+    ) -> PyMetastore:
+        """Construct a PyMetastore backed by an HTTP metastore service."""