pyspiral 0.7.18__cp312-abi3-manylinux_2_28_x86_64.whl

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: ...

spiral/core/table/__init__.pyi

@@ -0,0 +1,141 @@
+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, Operation, 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 ScanState:
+    def to_json(self) -> str: ...
+    @staticmethod
+    def from_json(json: str) -> ScanState: ...
+
+class MaterializablePlan:
+    pass
+
+class EvaluatedExecutablePlan:
+    pass
+
+class EvaluatedPlanStream:
+    def __next__(self) -> EvaluatedExecutablePlan: ...
+    def __iter__(self) -> EvaluatedPlanStream: ...
+
+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 plan_state(self) -> ScanState: ...
+    def materializable_plan(self) -> MaterializablePlan: ...
+    def to_record_batches(
+        self,
+        key_range: KeyRange | None = None,
+        key_table: pa.Table | pa.RecordBatch | None = None,
+        batch_readahead: int | None = None,
+        progress: bool = True,
+    ) -> pa.RecordBatchReader: ...
+    def to_unordered_record_batches(
+        self,
+        key_table: pa.Table | pa.RecordBatch | None = None,
+        batch_readahead: int | None = None,
+        progress: bool = True,
+    ) -> 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 evaluate_analyze(
+        self, key_table: pa.Table | pa.RecordBatch | None = None, batch_readahead: int | None = None
+    ) -> EvaluatedPlanStream: ...
+    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 writeback(
+        self,
+        scan: Scan,
+        *,
+        key_range: KeyRange | None = None,
+        partition_size_bytes: int | None = None,
+        batch_readahead: int | None = None,
+    ): ...
+    def drop_columns(self, column_paths: list[str]): ...
+    def ops(self) -> list[Operation]: ...
+    def take(self) -> list[Operation]: ...
+    def include(self, ops: list[Operation]): ...
+    def commit(self, *, compact: bool = False, manifest_rows: int | None = None): ...
+    def abort(self): ...
+    def is_empty(self) -> bool: ...

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."""

spiral/core/table/spec/__init__.pyi

@@ -0,0 +1,215 @@
+"""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 Operation:
+    # Base class for all operations in the WAL.
+    def to_json(self) -> str: ...
+    @staticmethod
+    def from_json(json: str) -> Operation: ...
+
+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:
+    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: ...
+    def __reduce__(self) -> tuple[type[Key], tuple[bytes]]: ...
+
+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
+    manifest_handle: ManifestHandle
+
+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."""