pyspiral 0.6.6__cp312-abi3-macosx_11_0_arm64.whl

spiral/protogen/_/substrait/extensions/__init__.py

@@ -0,0 +1,169 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# sources: substrait/extensions/extensions.proto
+# plugin: python-betterproto2
+# This file has been @generated
+
+__all__ = (
+    "AdvancedExtension",
+    "SimpleExtensionDeclaration",
+    "SimpleExtensionDeclarationExtensionFunction",
+    "SimpleExtensionDeclarationExtensionType",
+    "SimpleExtensionDeclarationExtensionTypeVariation",
+    "SimpleExtensionUri",
+)
+
+from dataclasses import dataclass
+
+import betterproto2
+
+from ...message_pool import default_message_pool
+
+_COMPILER_VERSION = "0.9.0"
+betterproto2.check_compiler_version(_COMPILER_VERSION)
+
+
+@dataclass(eq=False, repr=False)
+class AdvancedExtension(betterproto2.Message):
+    """
+    A generic object that can be used to embed additional extension information
+    into the serialized substrait plan.
+    """
+
+    optimization: "list[__google__protobuf__.Any]" = betterproto2.field(1, betterproto2.TYPE_MESSAGE, repeated=True)
+    """
+    An optimization is helpful information that don't influence semantics. May
+    be ignored by a consumer.
+    """
+
+    enhancement: "__google__protobuf__.Any | None" = betterproto2.field(2, betterproto2.TYPE_MESSAGE, optional=True)
+    """
+    An enhancement alter semantics. Cannot be ignored by a consumer.
+    """
+
+
+default_message_pool.register_message("substrait.extensions", "AdvancedExtension", AdvancedExtension)
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclaration(betterproto2.Message):
+    """
+    Describes a mapping between a specific extension entity and the uri where
+    that extension can be found.
+
+    Oneofs:
+        - mapping_type:
+    """
+
+    extension_type: "SimpleExtensionDeclarationExtensionType | None" = betterproto2.field(
+        1, betterproto2.TYPE_MESSAGE, optional=True, group="mapping_type"
+    )
+
+    extension_type_variation: "SimpleExtensionDeclarationExtensionTypeVariation | None" = betterproto2.field(
+        2, betterproto2.TYPE_MESSAGE, optional=True, group="mapping_type"
+    )
+
+    extension_function: "SimpleExtensionDeclarationExtensionFunction | None" = betterproto2.field(
+        3, betterproto2.TYPE_MESSAGE, optional=True, group="mapping_type"
+    )
+
+
+default_message_pool.register_message("substrait.extensions", "SimpleExtensionDeclaration", SimpleExtensionDeclaration)
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclarationExtensionFunction(betterproto2.Message):
+    extension_uri_reference: "int" = betterproto2.field(1, betterproto2.TYPE_UINT32)
+    """
+    references the extension_uri_anchor defined for a specific extension URI.
+    """
+
+    function_anchor: "int" = betterproto2.field(2, betterproto2.TYPE_UINT32)
+    """
+    A surrogate key used in the context of a single plan to reference a
+    specific function
+    """
+
+    name: "str" = betterproto2.field(3, betterproto2.TYPE_STRING)
+    """
+    A function signature compound name
+    """
+
+
+default_message_pool.register_message(
+    "substrait.extensions", "SimpleExtensionDeclaration.ExtensionFunction", SimpleExtensionDeclarationExtensionFunction
+)
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclarationExtensionType(betterproto2.Message):
+    """
+    Describes a Type
+    """
+
+    extension_uri_reference: "int" = betterproto2.field(1, betterproto2.TYPE_UINT32)
+    """
+    references the extension_uri_anchor defined for a specific extension URI.
+    """
+
+    type_anchor: "int" = betterproto2.field(2, betterproto2.TYPE_UINT32)
+    """
+    A surrogate key used in the context of a single plan to reference a
+    specific extension type
+    """
+
+    name: "str" = betterproto2.field(3, betterproto2.TYPE_STRING)
+    """
+    the name of the type in the defined extension YAML.
+    """
+
+
+default_message_pool.register_message(
+    "substrait.extensions", "SimpleExtensionDeclaration.ExtensionType", SimpleExtensionDeclarationExtensionType
+)
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclarationExtensionTypeVariation(betterproto2.Message):
+    extension_uri_reference: "int" = betterproto2.field(1, betterproto2.TYPE_UINT32)
+    """
+    references the extension_uri_anchor defined for a specific extension URI.
+    """
+
+    type_variation_anchor: "int" = betterproto2.field(2, betterproto2.TYPE_UINT32)
+    """
+    A surrogate key used in the context of a single plan to reference a
+    specific type variation
+    """
+
+    name: "str" = betterproto2.field(3, betterproto2.TYPE_STRING)
+    """
+    the name of the type in the defined extension YAML.
+    """
+
+
+default_message_pool.register_message(
+    "substrait.extensions",
+    "SimpleExtensionDeclaration.ExtensionTypeVariation",
+    SimpleExtensionDeclarationExtensionTypeVariation,
+)
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionUri(betterproto2.Message):
+    extension_uri_anchor: "int" = betterproto2.field(1, betterproto2.TYPE_UINT32)
+    """
+    A surrogate key used in the context of a single plan used to reference the
+    URI associated with an extension.
+    """
+
+    uri: "str" = betterproto2.field(2, betterproto2.TYPE_STRING)
+    """
+    The URI where this extension YAML can be retrieved. This is the "namespace"
+    of this extension.
+    """
+
+
+default_message_pool.register_message("substrait.extensions", "SimpleExtensionURI", SimpleExtensionUri)
+
+
+from ...google import protobuf as __google__protobuf__

spiral/protogen/util.py

@@ -0,0 +1,41 @@
+import betterproto
+from betterproto.grpc.grpclib_server import ServiceBase
+
+
+def patch_protos(proto_module, our_module_globals):
+    """Calculate __all__ to re-export protos from a module."""
+
+    betterproto_types = (betterproto.Message, betterproto.Enum, betterproto.ServiceStub, ServiceBase)
+
+    proto_overrides = {}
+    missing = set()
+    for ident in dir(proto_module):
+        var = getattr(proto_module, ident)
+        if isinstance(var, type) and issubclass(var, betterproto_types):
+            if ident in our_module_globals:
+                override = id(our_module_globals.get(ident)) != id(var)
+            else:
+                override = False
+                missing.add(ident)
+            proto_overrides[ident] = override
+
+    if missing:
+        print(f"from {proto_module.__name__} import (")
+        for ident, override in proto_overrides.items():
+            if override:
+                print(f"    {ident} as {ident}_,")
+            else:
+                print(f"    {ident},")
+        print(")")
+        print("\n")
+        print("__all__ = [")
+        for ident in proto_overrides:
+            print(f'    "{ident}",')
+        print("]")
+
+        raise ValueError(f"Missing types that need to be re-exported: {missing}")
+
+    # Patch any local subclasses back into the original module so the gRPC client will use them
+    for ident, override in proto_overrides.items():
+        if override:
+            setattr(proto_module, ident, our_module_globals[ident])

spiral/scan.py

@@ -0,0 +1,285 @@
+from typing import TYPE_CHECKING, Any, Optional
+
+import pyarrow as pa
+
+from spiral.core.client import Shard, ShuffleConfig
+from spiral.core.table import Scan as CoreScan
+from spiral.core.table.spec import Schema
+from spiral.settings import CI, DEV
+
+if TYPE_CHECKING:
+    import dask.dataframe as dd
+    import datasets.iterable_dataset as hf  # noqa
+    import pandas as pd
+    import polars as pl
+    import streaming  # noqa
+    import torch.utils.data as torchdata  # noqa
+
+    from spiral.dataloader import SpiralDataLoader, World  # noqa
+
+
+class Scan:
+    """Scan object."""
+
+    def __init__(self, core: CoreScan):
+        self.core = core
+
+    @property
+    def metrics(self) -> dict[str, Any]:
+        """Returns metrics about the scan."""
+        return self.core.metrics()
+
+    @property
+    def schema(self) -> Schema:
+        """Returns the schema of the scan."""
+        return self.core.schema()
+
+    def is_empty(self) -> bool:
+        """Check if the Spiral is empty for the given key range.
+
+        **IMPORTANT**: False negatives are possible, but false positives are not,
+            i.e. is_empty can return False and scan can return zero rows.
+        """
+        return self.core.is_empty()
+
+    def to_record_batches(
+        self,
+        key_table: pa.Table | pa.RecordBatchReader | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+    ) -> pa.RecordBatchReader:
+        """Read as a stream of RecordBatches.
+
+        Args:
+            key_table: a table of keys to "take" (including aux columns for cell-push-down).
+                If None, the scan will be executed without a key table.
+            batch_size: the maximum number of rows per returned batch.
+                IMPORTANT: This is currently only respected when the key_table is used. If key table is a
+                    RecordBatchReader, the batch_size argument must be None, and the existing batching is respected.
+            batch_readahead: the number of batches to prefetch in the background.
+        """
+        if isinstance(key_table, pa.RecordBatchReader):
+            if batch_size is not None:
+                raise ValueError(
+                    "batch_size must be None when key_table is a RecordBatchReader, the existing batching is respected."
+                )
+        elif isinstance(key_table, pa.Table):
+            key_table = key_table.to_reader(max_chunksize=batch_size)
+
+        return self.core.to_record_batches(key_table=key_table, batch_readahead=batch_readahead)
+
+    def to_table(
+        self,
+        key_table: pa.Table | pa.RecordBatchReader | None = None,
+    ) -> pa.Table:
+        """Read into a single PyArrow Table.
+
+        Args:
+            key_table: a table of keys to "take" (including aux columns for cell-push-down).
+                If None, the scan will be executed without a key table.
+        """
+        # NOTE: Evaluates fully on Rust side which improved debuggability.
+        if DEV and not CI and key_table is None:
+            rb = self.core.to_record_batch()
+            return pa.Table.from_batches([rb])
+
+        return self.to_record_batches(key_table=key_table).read_all()
+
+    def to_dask(self) -> "dd.DataFrame":
+        """Read into a Dask DataFrame.
+
+        Requires the `dask` package to be installed.
+        """
+        import dask.dataframe as dd
+        import pandas as pd
+
+        def _read_shard(shard: Shard) -> pd.DataFrame:
+            # TODO(ngates): we need a way to preserve the existing asofs?
+            raise NotImplementedError()
+
+        # Fetch a set of partition ranges
+        return dd.from_map(_read_shard, self.shards())
+
+    def to_pandas(self) -> "pd.DataFrame":
+        """Read into a Pandas DataFrame.
+
+        Requires the `pandas` package to be installed.
+        """
+        return self.to_table().to_pandas()
+
+    def to_polars(self) -> "pl.DataFrame":
+        """Read into a Polars DataFrame.
+
+        Requires the `polars` package to be installed.
+        """
+        import polars as pl
+
+        return pl.from_arrow(self.to_record_batches())
+
+    def to_data_loader(
+        self, seed: int = 42, shuffle_buffer_size: int = 8192, batch_size: int = 32, **kwargs
+    ) -> "SpiralDataLoader":
+        """Read into a Torch-compatible DataLoader for single-node training.
+
+        Args:
+            seed: Random seed for reproducibility.
+            shuffle_buffer_size: Size of shuffle buffer.
+            batch_size: Batch size.
+            **kwargs: Additional arguments passed to SpiralDataLoader constructor.
+
+        Returns:
+            SpiralDataLoader with shuffled shards.
+        """
+        from spiral.dataloader import SpiralDataLoader
+
+        return SpiralDataLoader(
+            self, seed=seed, shuffle_buffer_size=shuffle_buffer_size, batch_size=batch_size, **kwargs
+        )
+
+    def to_distributed_data_loader(
+        self,
+        world: Optional["World"] = None,
+        shards: list[Shard] | None = None,
+        seed: int = 42,
+        shuffle_buffer_size: int = 8192,
+        batch_size: int = 32,
+        **kwargs,
+    ) -> "SpiralDataLoader":
+        """Read into a Torch-compatible DataLoader for distributed training.
+
+        Args:
+            world: World configuration with rank and world_size.
+                If None, auto-detects from torch.distributed.
+            shards: Optional sharding. Sharding is global, i.e. the world will be used to select
+                the shards for this rank. If None, uses scan's natural sharding.
+            seed: Random seed for reproducibility.
+            shuffle_buffer_size: Size of shuffle buffer.
+                Use zero to skip shuffling with shuffle buffer.
+            batch_size: Batch size.
+            **kwargs: Additional arguments passed to SpiralDataLoader constructor.
+
+        Returns:
+            SpiralDataLoader with shards partitioned for this rank.
+        """
+        # Example usage:
+        #
+        # Auto-detect from PyTorch distributed:
+        #   loader: SpiralDataLoader = scan.to_distributed_data_loader(batch_size=32)
+        #
+        # Explicit world configuration:
+        #   world = World(rank=0, world_size=4)
+        #   loader: SpiralDataLoader = scan.to_distributed_data_loader(world=world, batch_size=32)
+
+        from spiral.dataloader import SpiralDataLoader, World
+
+        if world is None:
+            world = World.from_torch()
+
+        shards = shards or self.shards()
+        # Apply world partitioning to shards.
+        shards = world.shards(shards, seed)
+
+        return SpiralDataLoader(
+            self,
+            shards=shards,
+            shuffle_shards=False,  # Shards are shuffled before selected for the world.
+            seed=seed,
+            shuffle_buffer_size=shuffle_buffer_size,
+            batch_size=batch_size,
+            **kwargs,
+        )
+
+    def resume_data_loader(self, state: dict[str, Any], **kwargs) -> "SpiralDataLoader":
+        """Create a DataLoader from checkpoint state, resuming from where it left off.
+
+        This is the recommended way to resume training from a checkpoint. It extracts
+        the seed, samples_yielded, and shards from the state dict and creates a new
+        DataLoader that will skip the already-processed samples.
+
+        Args:
+            state: Checkpoint state from state_dict().
+            **kwargs: Additional arguments to pass to SpiralDataLoader constructor.
+                These will override values in the state dict where applicable.
+
+        Returns:
+            New SpiralDataLoader instance configured to resume from the checkpoint.
+        """
+        # Example usage:
+        #
+        # Save checkpoint during training:
+        #   loader = scan.to_distributed_data_loader(batch_size=32, seed=42)
+        #   checkpoint = loader.state_dict()
+        #
+        # Resume later - uses same shards from checkpoint:
+        #   resumed_loader = scan.resume_data_loader(
+        #       checkpoint,
+        #       batch_size=32,
+        #       transform_fn=my_transform,
+        #   )
+        from spiral.dataloader import SpiralDataLoader
+
+        return SpiralDataLoader.from_state_dict(self, state, **kwargs)
+
+    def to_iterable_dataset(
+        self,
+        shards: list[Shard] | None = None,
+        shuffle: ShuffleConfig | None = None,
+        batch_readahead: int | None = None,
+        infinite: bool = False,
+    ) -> "hf.IterableDataset":
+        """Returns a Huggingface's IterableDataset.
+
+        Requires `datasets` package to be installed.
+
+        Note: For new code, consider using SpiralDataLoader instead.
+
+        Args:
+            shards: Optional list of shards to read. If None, uses scan's natural sharding.
+            shuffle: Optional ShuffleConfig for configuring within-shard sample shuffling.
+                If None, no shuffling is performed.
+            batch_readahead: Controls how many batches to read ahead concurrently.
+                If pipeline includes work after reading (e.g. decoding, transforming, ...) this can be set higher.
+                Otherwise, it should be kept low to reduce next batch latency. Defaults to 2.
+            infinite: If True, the returned IterableDataset will loop infinitely over the data,
+                re-shuffling ranges after exhausting all data.
+        """
+        stream = self.core.to_shuffled_record_batches(
+            shards=shards,
+            shuffle=shuffle,
+            batch_readahead=batch_readahead,
+            infinite=infinite,
+        )
+
+        from spiral.iterable_dataset import to_iterable_dataset
+
+        return to_iterable_dataset(stream)
+
+    def shards(self) -> list[Shard]:
+        """Get list of shards for this scan.
+
+        The shards are based on the scan's physical data layout (file fragments).
+        Each shard contains a key range and cardinality (set to None when unknown).
+
+        Returns:
+            List of Shard objects with key range and cardinality (if known).
+
+        """
+        return self.core.shards()
+
+    def _debug(self):
+        # Visualizes the scan, mainly for debugging purposes.
+        from spiral.debug.scan import show_scan
+
+        show_scan(self.core)
+
+    def _dump_manifests(self):
+        # Print manifests in a human-readable format.
+        from spiral.debug.manifests import display_scan_manifests
+
+        display_scan_manifests(self.core)
+
+    def _dump_metrics(self):
+        # Print metrics in a human-readable format.
+        from spiral.debug.metrics import display_metrics
+
+        display_metrics(self.metrics)

spiral/server.py

@@ -0,0 +1,17 @@
+import socket
+import time
+
+
+def wait_for_port(port: int, host: str = "localhost", timeout: float = 5.0):
+    """Wait until a port starts accepting TCP connections."""
+    start_time = time.time()
+    while True:
+        try:
+            with socket.create_connection((host, port), timeout=timeout):
+                break
+        except OSError as ex:
+            time.sleep(0.01)
+            if time.time() - start_time >= timeout:
+                raise TimeoutError(
+                    f"Waited too long for the port {port} on host {host} to start accepting connections."
+                ) from ex

spiral/settings.py

@@ -0,0 +1,114 @@
+import functools
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated
+
+import typer
+from pydantic import Field, ValidatorFunctionWrapHandler, WrapValidator
+from pydantic_settings import (
+    BaseSettings,
+    InitSettingsSource,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+)
+
+from spiral.core.authn import Authn, DeviceCodeAuth, Token
+from spiral.core.client import Spiral
+
+if TYPE_CHECKING:
+    from spiral.api import SpiralAPI
+
+DEV = "PYTEST_VERSION" in os.environ or bool(os.environ.get("SPIRAL_DEV", None))
+CI = "GITHUB_ACTIONS" in os.environ
+
+APP_DIR = Path(typer.get_app_dir("pyspiral"))
+LOG_DIR = APP_DIR / "logs"
+
+PACKAGE_NAME = "pyspiral"
+
+
+def validate_token(v, handler: ValidatorFunctionWrapHandler):
+    if isinstance(v, str):
+        return Token(v)
+    else:
+        raise ValueError("Token value must be a string")
+
+
+TokenType = Annotated[Token, WrapValidator(validate_token)]
+
+
+class SpiralDBSettings(BaseSettings):
+    model_config = SettingsConfigDict(frozen=True)
+
+    host: str = "localhost" if DEV else "api.spiraldb.com"
+    port: int = 4279 if DEV else 443
+    ssl: bool = not DEV
+    token: TokenType | None = None
+
+    @property
+    def uri(self) -> str:
+        return f"{'https' if self.ssl else 'http'}://{self.host}:{self.port}"
+
+
+class SpfsSettings(BaseSettings):
+    model_config = SettingsConfigDict(frozen=True)
+
+    host: str = "localhost" if DEV else "spfs.spiraldb.dev"
+    port: int = 4295 if DEV else 443
+    ssl: bool = not DEV
+
+    @property
+    def uri(self) -> str:
+        return f"{'https' if self.ssl else 'http'}://{self.host}:{self.port}"
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_nested_delimiter="__",
+        env_prefix="SPIRAL__",
+        frozen=True,
+    )
+
+    spiraldb: SpiralDBSettings = Field(default_factory=SpiralDBSettings)
+    spfs: SpfsSettings = Field(default_factory=SpfsSettings)
+    file_format: str = Field(default="vortex")
+
+    @functools.cached_property
+    def api(self) -> "SpiralAPI":
+        from spiral.api import SpiralAPI
+
+        return SpiralAPI(self.authn, base_url=self.spiraldb.uri)
+
+    @functools.cached_property
+    def core(self) -> Spiral:
+        return Spiral(
+            api_url=self.spiraldb.uri,
+            spfs_url=self.spfs.uri,
+            authn=self.authn,
+        )
+
+    @functools.cached_property
+    def authn(self):
+        if self.spiraldb.token:
+            return Authn.from_token(self.spiraldb.token)
+        return Authn.from_fallback(self.spiraldb.uri)
+
+    @functools.cached_property
+    def device_code_auth(self) -> DeviceCodeAuth:
+        return DeviceCodeAuth.default()
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        init_settings: InitSettingsSource,
+        **kwargs,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return env_settings, dotenv_settings, init_settings
+
+
+@functools.cache
+def settings() -> Settings:
+    return Settings()

spiral/snapshot.py

@@ -0,0 +1,56 @@
+from typing import TYPE_CHECKING
+
+from spiral.core.table import Snapshot as CoreSnapshot
+from spiral.core.table.spec import Schema
+from spiral.types_ import Timestamp
+
+if TYPE_CHECKING:
+    import duckdb
+    import polars as pl
+    import pyarrow.dataset as ds
+    import torch.utils.data as torchdata  # noqa
+
+    from spiral.table import Table
+
+
+class Snapshot:
+    """Spiral table snapshot.
+
+    A snapshot represents a point-in-time view of a table.
+    """
+
+    def __init__(self, table: "Table", core: CoreSnapshot):
+        self.core = core
+        self._table = table
+
+    @property
+    def asof(self) -> Timestamp:
+        """Returns the asof timestamp of the snapshot."""
+        return self.core.asof
+
+    def schema(self) -> Schema:
+        """Returns the schema of the snapshot."""
+        return self.core.table.get_schema(asof=self.asof)
+
+    @property
+    def table(self) -> "Table":
+        """Returns the table associated with the snapshot."""
+        return self._table
+
+    def to_dataset(self) -> "ds.Dataset":
+        """Returns a PyArrow Dataset representing the table."""
+        from spiral.dataset import Dataset
+
+        return Dataset(self)
+
+    def to_polars(self) -> "pl.LazyFrame":
+        """Returns a Polars LazyFrame for the Spiral table."""
+        import polars as pl
+
+        return pl.scan_pyarrow_dataset(self.to_dataset())
+
+    def to_duckdb(self) -> "duckdb.DuckDBPyRelation":
+        """Returns a DuckDB relation for the Spiral table."""
+        import duckdb
+
+        return duckdb.from_arrow(self.to_dataset())

spiral/streaming_/__init__.py

@@ -0,0 +1,3 @@
+from .stream import SpiralStream
+
+__all__ = ["SpiralStream"]