pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/client.py

@@ -0,0 +1,79 @@
+from typing import TYPE_CHECKING
+
+import jwt
+
+from spiral.api import SpiralAPI
+from spiral.api.projects import CreateProjectRequest, CreateProjectResponse
+from spiral.core.client import Spiral as CoreSpiral
+from spiral.settings import Settings, settings
+
+if TYPE_CHECKING:
+    from spiral.iceberg import Iceberg
+    from spiral.project import Project
+
+
+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
+    def config(self) -> Settings:
+        return self._config
+
+    @property
+    def api(self) -> SpiralAPI:
+        return self._api
+
+    @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, 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, id=project_id, name=project_id)
+
+    @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/client/__init__.pyi

@@ -0,0 +1,117 @@
+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."""
+        ...
+
+    def logout(self) -> None:
+        """Logout from the device authentication session."""
+        ...
+
+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 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 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 open_table_scan(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        exclude_keys: bool = False,
+    ) -> TableScan:
+        """Construct a table scan."""
+        ...
+
+    def open_transaction(self, table: Table, format: str | None = None) -> TableTransaction:
+        """Being transaction."""
+        ...
+
+    def open_maintenance(self, table: Table, format: str | None = None) -> TableMaintenance:
+        """Access maintenance operations for a table."""
+        ...
+    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 get_text_index(self, index_id: str) -> TextIndex:
+        """Get a text-based index."""
+        ...
+
+    def open_search_scan(
+        self,
+        rank_by: Expr,
+        top_k: int,
+        # NOTE(marko): Required for now.
+        freshness_window_s: int,
+        *,
+        filter: Expr | None = None,
+    ) -> SearchScan:
+        """Query an index."""
+        ...
+
+    def _sync_snapshot(self, index_id: str, snapshot: TableSnapshot) -> None:
+        """Synchronize an index with a table snapshot.
+
+        IMPORTANT: This is only exposed for testing purposes and should not be used.
+        """
+        ...

spiral/core/index/__init__.pyi

@@ -0,0 +1,15 @@
+import pyarrow as pa
+
+class IndexStatus:
+    status: str
+    staleness_s: int | None
+    # An extent of keys that are indexed.
+    # key_extent: KeyExtent | None
+
+class TextIndex:
+    id: str
+
+    def status(self) -> IndexStatus: ...
+
+class SearchScan:
+    def to_record_batches(self) -> pa.RecordBatchReader: ...

spiral/core/table/__init__.pyi

@@ -0,0 +1,108 @@
+from typing import Any, Literal
+
+import pyarrow as pa
+from spiral.expressions import Expr
+
+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: ...
+
+class Table:
+    def __init__(self, metastore: PyMetastore): ...
+
+    id: str
+    root_uri: str
+    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) -> TableSnapshot: ...
+
+class TableSnapshot:
+    """A snapshot of a table at a specific point in time."""
+
+    asof: int
+    table: Table
+    wal: WriteAheadLog
+
+class TableScan:
+    def key_schema(self) -> Schema: ...
+    def schema(self) -> Schema: ...
+    def is_empty(self) -> bool: ...
+    def split(self) -> list[KeyRange]: ...
+    def table_ids(self) -> list[str]: ...
+    def column_groups(self) -> list[ColumnGroup]: ...
+    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,
+        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]: ...
+
+class KeySpaceScan:
+    manifest: FragmentManifest
+
+    def key_schema(self) -> Schema: ...
+
+class ColumnGroupScan:
+    manifest: FragmentManifest
+
+    def schema(self) -> Schema: ...
+
+class TableTransaction:
+    status: str
+
+    def write(self, expr: Expr, *, partition_size_bytes: int | None = None): ...
+    def commit(self): ...
+    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(
+        self,
+        column_group: ColumnGroup,
+        *,
+        mode: Literal["plan", "read", "write"] | None = None,
+        partition_bytes_min: int | None = None,
+    ): ...
+    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(data_schema: pa.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,62 @@
+"""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."""
+
+    @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

@@ -0,0 +1,214 @@
+"""Type definitions for the spiral.core.spec module shipped as part of the native library."""
+
+import pyarrow as pa
+
+class ColumnGroup:
+    def __init__(self, path: list[str]): ...
+    @property
+    def table_id(self) -> str: ...
+    @property
+    def path(self) -> list[str]: ...
+    def identifier(self, salt: int) -> str:
+        """Return the column group identifier based on the given salt."""
+
+    @staticmethod
+    def from_str(path: str) -> ColumnGroup: ...
+
+class KeySpaceMetadata:
+    def __init__(
+        self,
+        *,
+        manifest_handle: ManifestHandle | None,
+        last_modified_at: int,
+    ): ...
+
+    manifest_handle: ManifestHandle | None
+    last_modified_at: int
+
+    def asof(self, asof: int) -> KeySpaceMetadata:
+        """Returns the metadata as of a given timestamp. Currently just filtering versioned schemas."""
+        ...
+
+    def apply_wal(self, wal: WriteAheadLog) -> KeySpaceMetadata:
+        """Applies the given WAL to the metadata."""
+
+class ColumnGroupMetadata:
+    def __init__(
+        self,
+        *,
+        column_group: ColumnGroup,
+        manifest_handle: ManifestHandle | None,
+        last_modified_at: int,
+        schema_versions: list[VersionedSchema] | None,
+        immutable_schema: bool,
+        schema_salt: int,
+    ): ...
+
+    column_group: ColumnGroup
+    manifest_handle: ManifestHandle | None
+    last_modified_at: int
+    schema_versions: list[VersionedSchema]
+    immutable_schema: bool
+    schema_salt: int
+
+    def latest_schema(self) -> VersionedSchema:
+        """Returns the latest schema of the column group."""
+        ...
+
+    def asof(self, asof: int) -> ColumnGroupMetadata:
+        """Returns the metadata as of a given timestamp. Currently just filtering versioned schemas."""
+        ...
+
+    def apply_wal(self, wal: WriteAheadLog) -> ColumnGroupMetadata:
+        """Applies the given WAL to the metadata."""
+
+class LogEntry:
+    ts: int
+    operation: (
+        KeySpaceWriteOp
+        | ColumnGroupWriteOp
+        | SchemaEvolutionOp
+        | SchemaBreakOp
+        | KeySpaceCompactOp
+        | ColumnGroupCompactOp
+    )
+
+    def column_group(self) -> ColumnGroup | None:
+        """Returns the column group of the entry if it is associated with one."""
+
+class FileFormat:
+    def __init__(self, value: int): ...
+
+    Parquet: FileFormat
+    Protobuf: FileFormat
+    BinaryArray: FileFormat
+    Vortex: FileFormat
+
+    def __int__(self) -> int:
+        """Returns the protobuf enum int value."""
+        ...
+
+    def __str__(self) -> str:
+        """Returns the string representation of the file format."""
+        ...
+
+class FragmentLevel:
+    L0: FragmentLevel
+    L1: FragmentLevel
+
+    def __int__(self) -> int:
+        """Returns the protobuf enum int value."""
+        ...
+
+class Key:
+    key: bytes
+
+    def __init__(self, key: bytes): ...
+    def __bytes__(self): ...
+    def step(self) -> Key:
+        """Returns the next key in the key space."""
+
+    @staticmethod
+    def min() -> Key: ...
+    @staticmethod
+    def max() -> Key: ...
+
+class KeyExtent:
+    """An inclusive range of keys."""
+
+    def __init__(self, *, min: Key, max: Key): ...
+
+    min: Key
+    max: Key
+
+    def union(self, key_extent: KeyExtent) -> KeyExtent: ...
+    def __or__(self, other: KeyExtent) -> KeyExtent: ...
+    def intersection(self, key_extent: KeyExtent) -> KeyExtent | None: ...
+    def __and__(self, other: KeyExtent) -> KeyExtent | None: ...
+    def contains(self, item: Key) -> bool: ...
+    def __contains__(self, item: Key) -> bool: ...
+
+class KeySpan:
+    """An exclusive range of keys as indexed by their position in a key space."""
+
+    def __init__(self, *, begin: int, end: int): ...
+
+    begin: int
+    end: int
+
+    def __len__(self) -> int: ...
+    def shift(self, offset: int) -> KeySpan: ...
+    def union(self, other: KeySpan) -> KeySpan: ...
+    def __or__(self, other: KeySpan) -> KeySpan: ...
+
+class ManifestHandle:
+    id: str
+    format: FileFormat
+    file_size: int
+
+class Schema:
+    def to_arrow(self) -> pa.Schema:
+        """Returns the Arrow schema."""
+        ...
+    @staticmethod
+    def from_arrow(arrow: pa.Schema) -> Schema:
+        """Creates a Schema from an Arrow schema."""
+        ...
+    def __len__(self):
+        """Returns the number of columns in the schema."""
+        ...
+    @property
+    def names(self) -> list[str]:
+        """Returns the names of the columns in the schema."""
+        ...
+
+class VersionedSchema:
+    ts: int
+    schema: Schema
+    column_ids: list[str]
+
+class KeySpaceWriteOp:
+    ks_id: str
+    manifest_handle: ManifestHandle
+
+class ColumnGroupWriteOp:
+    column_group: ColumnGroup
+    level: FragmentLevel
+    manifest_handle: ManifestHandle
+    key_span: KeySpan
+    key_extent: KeyExtent
+    column_ids: list[str]
+
+class SchemaEvolutionOp:
+    column_group: ColumnGroup
+
+class SchemaBreakOp:
+    column_group: ColumnGroup
+
+class KeySpaceCompactOp:
+    ks_ids: list[str]
+    moved_ks_ids: list[str]
+
+class ColumnGroupCompactOp:
+    column_group: ColumnGroup
+    fragment_ids: list[int]
+
+class WriteAheadLog:
+    def __init__(
+        self,
+        *,
+        entries: list[LogEntry] | None = None,
+        truncated_up_to: int = 0,
+    ): ...
+
+    entries: list[LogEntry]
+    truncated_up_to: int
+
+    @property
+    def last_modified_at(self) -> int:
+        """Returns the timestamp of the last modification of the log."""
+
+    def filter(
+        self, asof: int | None = None, since: int | None = None, column_group: ColumnGroup | None = None
+    ) -> WriteAheadLog:
+        """Filters the WAL to entries by the given parameters."""

spiral/datetime_.py

@@ -0,0 +1,27 @@
+import warnings
+from datetime import datetime, timedelta, timezone, tzinfo
+
+_THE_EPOCH = datetime.fromtimestamp(0, tz=timezone.utc)
+
+
+def local_tz() -> tzinfo:
+    """Determine this machine's local timezone."""
+    tz = datetime.now().astimezone().tzinfo
+    if tz is None:
+        raise ValueError("Could not determine this machine's local timezone.")
+    return tz
+
+
+def timestamp_micros(instant: datetime) -> int:
+    """The number of microseconds between the epoch and the given instant."""
+    if instant.tzinfo is None:
+        warnings.warn("assuming timezone-naive datetime is local time", stacklevel=2)
+        instant = instant.replace(tzinfo=local_tz())
+    return (instant - _THE_EPOCH) // timedelta(microseconds=1)
+
+
+def from_timestamp_micros(ts: int) -> datetime:
+    """Convert a timestamp in microseconds to a datetime."""
+    if ts < 0:
+        raise ValueError("Timestamp must be non-negative")
+    return _THE_EPOCH + timedelta(microseconds=ts)