pyspiral 0.8.2__cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

spiral/client.py

@@ -0,0 +1,269 @@
+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.authn import Authn
+from spiral.core.client import Internal, Shard
+from spiral.core.client import Spiral as CoreSpiral
+from spiral.core.config import ClientSettings
+from spiral.datetime_ import timestamp_micros
+from spiral.expressions import ExprLike
+from spiral.scan import Scan
+
+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:
+    """Main client for interacting with the Spiral data platform.
+
+    Configuration is loaded with the following priority (highest to lowest):
+    1. Explicit parameters.
+    2. Environment variables (`SPIRAL__*`)
+    3. Config file (`~/.spiral.toml`)
+    4. Default values (production URLs)
+
+    Examples:
+
+    ```python
+    import spiral
+    # Default configuration
+    sp = spiral.Spiral()
+
+    # With config overrides
+    sp = spiral.Spiral(overrides={"limits.concurrency": "16"})
+    ```
+
+    Args:
+        config: Custom ClientSettings object. Defaults to global settings.
+        overrides: Configuration overrides using dot notation,
+            see the [Client Configuration](https://docs.spiraldb.com/config) page for a full list.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: ClientSettings | None = None,
+        overrides: dict[str, str] | None = None,
+    ):
+        self._overrides = overrides
+        self._config = config
+        self._org = None
+        self._core = None
+        self._api = None
+
+    @property
+    def config(self) -> ClientSettings:
+        """Returns the client's configuration"""
+        return self.core.config()
+
+    @property
+    def authn(self) -> Authn:
+        """Get the authentication handler for this client."""
+        return self.core.authn()
+
+    @property
+    def api(self) -> SpiralAPI:
+        if self._api is None:
+            self._api = SpiralAPI(self.authn, base_url=self.config.server_url)
+        return self._api
+
+    @property
+    def core(self) -> CoreSpiral:
+        if self._core is None:
+            self._core = CoreSpiral(
+                config=self._config,
+                overrides=self._overrides,
+            )
+
+        return self._core
+
+    @property
+    def internal(self) -> Internal:
+        return self.core.internal(format=self.config.file_format)
+
+    @property
+    def organization(self) -> str:
+        if self._org is None:
+            token = self.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: execute the scan on the version of the table as of the given timestamp.
+        """
+        from spiral import expressions as se
+
+        if isinstance(asof, datetime):
+            asof = timestamp_micros(asof)
+
+        # Combine all projections into a single struct.
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return Scan(
+            self,
+            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 resume_scan(self, state_json: str) -> Scan:
+        """Resumes a previously started scan using its scan state.
+
+        Args:
+            state_json: The scan state returned by a previous scan.
+        """
+        from spiral.core.table import ScanState
+
+        state = ScanState.from_json(state_json)
+        return Scan(self, self.core.load_scan(state))
+
+    def compute_shards(
+        self,
+        max_batch_size: int,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        asof: datetime | int | None = None,
+        stream: bool = False,
+    ) -> list[Shard]:
+        """Computes shards over the given projections and filter.
+
+        Args:
+            max_batch_size: The maximum number of rows per shard.
+            projections: a set of expressions that return struct arrays.
+            where: a query expression to apply to the data.
+            asof: execute the scan on the version of the table as of the given timestamp.
+            stream: if true, builds shards in a streaming fashion, suitable for very large tables.
+        """
+        from spiral import expressions as se
+
+        if isinstance(asof, datetime):
+            asof = timestamp_micros(asof)
+
+        # Combine all projections into a single struct.
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return self.core.compute_shards(
+            max_batch_size, projection.__expr__, where.__expr__ if where else None, asof=asof, stream=stream
+        )
+
+    @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,21 @@
+from spiral.api.types import OrgId
+
+class Token:
+    def __init__(self, value: str): ...
+    def expose_secret(self) -> str: ...
+
+class 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,258 @@
+from enum import Enum
+from typing import Any
+
+import pyarrow as pa
+from spiral.api.types import DatasetName, IndexName, ProjectId, RootUri, TableId, TableName
+from spiral.core.authn import Authn
+from spiral.core.config import ClientSettings
+from spiral.core.table import ColumnGroupState, KeyRange, KeySpaceState, Scan, ScanState, 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,
+        config: ClientSettings | None = None,
+        overrides: dict[str, str] | None = None,
+    ):
+        """Initialize the Spiral client.
+
+        Args:
+            config: Client configuration, defaults to the global config.
+            overrides: Configuration overrides using dot notation,
+                see the [Client Configuration](/python-client) page for a full list.
+        """
+        ...
+
+    def authn(self) -> Authn:
+        """Get the current authentication context."""
+        ...
+
+    def scan(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        key_columns: KeyColumns | None = None,
+    ) -> Scan:
+        """Construct a table scan."""
+        ...
+
+    def load_scan(self, plan_state: ScanState) -> Scan:
+        """Load a scan from a serialized scan state."""
+        ...
+
+    def transaction(self, table: Table, *, partition_max_bytes: int | 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,
+        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 move_table(
+        self,
+        table_id: TableId,
+        new_dataset: DatasetName,
+    ):
+        """Move a table to a dataset in the same project."""
+        ...
+
+    def rename_table(
+        self,
+        table_id: TableId,
+        new_table: TableName,
+    ):
+        """Rename a table."""
+        ...
+
+    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 compute_shards(
+        self,
+        max_batch_size: int,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        stream: bool = False,
+    ) -> list[Shard]:
+        """Constructs shards for a given projection (and filter).
+
+        Useful for distributing work.
+        """
+        ...
+
+    def internal(self, *, format: str | None = None) -> Internal:
+        """Internal client APIs. It can change without notice."""
+        ...
+
+    def config(self) -> ClientSettings:
+        """Client-side configuration."""
+        ...
+
+class KeyColumns(Enum):
+    IfProjected = 0
+    Included = 1
+    Only = 2
+
+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.
+    """
+
+    buffer_size: int
+    seed: int | None
+
+    def __init__(
+        self,
+        buffer_size: int,
+        *,
+        seed: int | None = None,
+    ): ...
+
+class Internal:
+    def flush_wal(self, table: Table) -> None:
+        """
+        Flush the write-ahead log of 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 key_space_index_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/config/__init__.pyi

@@ -0,0 +1,35 @@
+class ClientSettings:
+    """Client configuration loaded from ~/.spiral.toml and environment variables."""
+
+    @staticmethod
+    def load() -> ClientSettings:
+        """Load ClientSettings from ~/.spiral.toml and environment variables.
+
+        Configuration priority (highest to lowest):
+        1. Environment variables (SPIRAL__*)
+        2. Config file (~/.spiral.toml)
+        3. Default values
+        """
+        ...
+
+    @property
+    def server_url(self) -> str:
+        """The Spiral API endpoint URL."""
+        ...
+
+    @property
+    def spfs_url(self) -> str:
+        """The SpFS endpoint URL."""
+        ...
+
+    @property
+    def file_format(self) -> str:
+        """File format for table storage (vortex or parquet)."""
+        ...
+
+    def to_json(self) -> str:
+        """Serialize to a JSON string"""
+        ...
+    @staticmethod
+    def from_json(json: str) -> ClientSettings:
+        """Deserialize from a JSON-formatted string"""

spiral/core/expr/__init__.pyi

@@ -0,0 +1,15 @@
+from pyarrow import Array, DataType, Scalar
+
+class Expr:
+    """Low level expression class."""
+
+def aux(name: str, data_type: DataType) -> Expr: ...
+
+# Array is correct (there is no ArrayData), see the table here:
+# https://arrow.apache.org/rust/arrow_pyarrow/index.html
+def scalar(array: Array[Scalar[DataType]]) -> Expr: ...
+def not_(expr: Expr) -> Expr: ...
+def is_null(expr: Expr) -> Expr: ...
+def binary(op: str, expr: Expr, Expr: Expr) -> Expr: ...
+def cast(_expr: Expr, _data_type: DataType) -> Expr: ...
+def array_lit(array: Array[Scalar[DataType]]) -> Expr: ...

spiral/core/expr/images/__init__.pyi

@@ -0,0 +1,3 @@
+from .. import Expr
+
+def encode_(images: Expr, format: str) -> Expr: ...

spiral/core/expr/list_/__init__.pyi

@@ -0,0 +1,4 @@
+from .. import Expr
+
+def contains(list: Expr, expr: Expr) -> Expr: ...
+def element_at(list: Expr, element: Expr) -> Expr: ...

spiral/core/expr/pushdown/__init__.pyi

@@ -0,0 +1,3 @@
+from .. import Expr
+
+def expensive(expr: Expr) -> Expr: ...

spiral/core/expr/refs/__init__.pyi

@@ -0,0 +1,4 @@
+from .. import Expr
+
+def ref(expr: Expr, field: str | None) -> Expr: ...
+def deref(expr: Expr, field: str | None) -> Expr: ...

spiral/core/expr/s3/__init__.pyi

@@ -0,0 +1,3 @@
+from .. import Expr
+
+def get(expr: Expr, abort_on_error: bool) -> Expr: ...

spiral/core/expr/str_/__init__.pyi

@@ -0,0 +1,3 @@
+from .. import Expr
+
+def substr(expr: Expr, begin: int, end: int | None) -> Expr: ...

spiral/core/expr/struct_/__init__.pyi

@@ -0,0 +1,6 @@
+from .. import Expr
+
+def getitem(expr: Expr, item: str) -> Expr: ...
+def select(expr: Expr, including: list[str] | None = None, excluding: list[str] | None = None) -> Expr: ...
+def pack(names: list[str], children: list[str], nullable: bool) -> Expr: ...
+def merge(names: list[Expr]) -> Expr: ...

spiral/core/expr/text/__init__.pyi

@@ -0,0 +1,5 @@
+from .. import Expr
+
+def field(expr: Expr, tokeneizer: str | None) -> Expr: ...
+def find(expr: Expr, term: str) -> Expr: ...
+def boost(expr: Expr, factor: float) -> Expr: ...

spiral/core/expr/udf/__init__.pyi

@@ -0,0 +1,14 @@
+from collections.abc import Callable
+
+from pyarrow import Array, DataType, Scalar
+
+from .. import Expr
+
+class UDF:
+    def __call__(self, args: list[Expr]) -> Expr: ...
+
+def create(
+    name: str,
+    return_type: Callable[[tuple[DataType, ...]], DataType],
+    invoke: Callable[[tuple[Array[Scalar[DataType]], ...]], Array[Scalar[DataType]]],
+) -> UDF: ...

spiral/core/expr/video/__init__.pyi

@@ -0,0 +1,3 @@
+from .. import Expr
+
+def read(expr: Expr, ranges: Expr, crops: Expr, format: str) -> Expr: ...