pyspiral 0.1.0__cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

spiral/proto/substrait/extensions/__init__.py

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

spiral/proto/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,168 @@
+from collections.abc import Iterator
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+from opentelemetry import trace
+
+from spiral.core.core import TableScan
+from spiral.core.spec import KeyRange, Schema
+from spiral.expressions.base import ExprLike
+
+if TYPE_CHECKING:
+    import dask.dataframe as dd
+    import pandas as pd
+    import polars as pl
+    from datasets import iterable_dataset
+
+tracer = trace.get_tracer("pyspiral.client.scan")
+
+
+def scan(
+    *projections: ExprLike,
+    where: ExprLike | None = None,
+    asof: datetime | int | str = None,
+    exclude_keys: bool = False,
+    # TODO(marko): Support config.
+    # config: Config | 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.
+        exclude_keys: whether to exclude the key columns in the scan result, defaults to False.
+    """
+    from spiral import expressions as se
+
+    # Combine all projections into a single struct.
+    projection = se.merge(*projections)
+    if where is not None:
+        where = se.lift(where)
+
+    return Scan(
+        TableScan(
+            projection.__expr__,
+            filter=where.__expr__ if where else None,
+            asof=asof,
+            exclude_keys=exclude_keys,
+        ),
+        # config=config,
+    )
+
+
+class Scan:
+    """Scan object."""
+
+    def __init__(
+        self,
+        scan: TableScan,
+        # TODO(marko): Support config.
+        # config: Config | None = None,
+    ):
+        # NOTE(ngates): this API is a little weird. e.g. if the query doesn't define an asof, it is resolved
+        #  when we wrap it into a core.Scan. Should we expose a Query object in the Python API that's reusable
+        #  and will re-resolve the asof? Or should we just expose a scan that fixes the asof at construction time?
+        self._scan = scan
+
+    @property
+    def metrics(self) -> dict[str, Any]:
+        """Returns metrics about the scan."""
+        return self._scan.metrics()
+
+    @property
+    def schema(self) -> Schema:
+        """Returns the schema of the scan."""
+        return self._scan.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._scan.is_empty()
+
+    def to_record_batches(self, key_table: pa.Table | pa.RecordBatchReader | 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 isinstance(key_table, pa.RecordBatchReader):
+            raise NotImplementedError("RecordBatchReader is not supported as key_table")
+
+        # Prefix non-key columns in the key table with # (auxiliary) to avoid conflicts with the scan schema.
+        if key_table is not None:
+            key_columns = list(self._scan.key_schema().to_arrow().names)
+            key_table = key_table.rename_columns(
+                {name: f"#{name}" if name not in key_columns else name for name in key_table.schema.names}
+            )
+
+        return self._scan.to_record_batches(aux_table=key_table)
+
+    def to_table(self) -> pa.Table:
+        """Read into a single PyArrow Table."""
+        return self.to_record_batches().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_key_range(key_range: KeyRange) -> pd.DataFrame:
+            # TODO(ngates): we need a way to preserve the existing asofs? Should we copy CoreScan instead of Query?
+            raise NotImplementedError()
+
+        # Fetch a set of partition ranges
+        return dd.from_map(_read_key_range, self.split())
+
+    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
+
+        # TODO(ngates): PR PyArrow to support lazy datasets
+        return pl.from_arrow(self.to_record_batches())
+
+    def to_pytorch(self) -> "iterable_dataset.IterableDataset":
+        """Returns an iterable dataset that can be used to build a `pytorch.DataLoader`.
+
+        Requires the `datasets` package to be installed.
+        """
+        from datasets.iterable_dataset import ArrowExamplesIterable, IterableDataset
+
+        def _generate_tables(**kwargs) -> Iterator[tuple[int, pa.Table]]:
+            stream = self.to_record_batches()
+
+            # This key is unused when training with IterableDataset.
+            # Default implementation returns shard id, e.g. parquet row group id.
+            for i, rb in enumerate(stream):
+                yield i, pa.Table.from_batches([rb], stream.schema)
+
+        # NOTE: Type annotation Callable[..., tuple[str, pa.Table]] is wrong. The return value must be iterable.
+        ex_iterable = ArrowExamplesIterable(generate_tables_fn=_generate_tables, kwargs={})
+        return IterableDataset(ex_iterable=ex_iterable)
+
+    def split(self) -> list[KeyRange]:
+        return self._scan.split()
+
+    def debug(self):
+        # Visualizes the scan, mainly for debugging purposes.
+        # NOTE: This is not part of the API and may disappear at any moment.
+        from spiral.debug import show_scan
+
+        show_scan(self._scan)

spiral/settings.py

@@ -0,0 +1,157 @@
+import functools
+import hashlib
+import os
+from collections.abc import Callable, Mapping
+from pathlib import Path
+from typing import Any
+
+import jwt
+import typer
+from pydantic import BaseModel, Field
+from pydantic_settings import (
+    BaseSettings,
+    InitSettingsSource,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+    TomlConfigSettingsSource,
+)
+
+from spiral.api import SpiralAPI
+from spiral.authn.authn import (
+    DeviceAuthProvider,
+    EnvironmentAuthn,
+    FallbackAuthn,
+    TokenAuthn,
+    TokenExchangeProvider,
+)
+from spiral.authn.device import DeviceAuth
+from spiral.authn.github_ import GitHubActionsProvider
+from spiral.authn.modal_ import ModalProvider
+
+DEV = "PYTEST_VERSION" in os.environ or bool(os.environ.get("SPIRAL_DEV", None))
+APP_DIR = Path(typer.get_app_dir("pyspiral"))
+LOG_DIR = APP_DIR / "logs"
+CONFIG_FILE = APP_DIR / "config.toml"
+
+
+class AuthSettings(BaseSettings):
+    model_config = SettingsConfigDict(frozen=True)
+
+    domain: str = "https://device.spiraldb.com"
+    client_id: str = "client_01J1CRS967RFQY7JSE8XZ44ATR" if DEV else "client_01J1CRS9MGF103FKTKZDCNJVM3"
+
+
+class AccessToken(BaseModel):
+    token: str
+
+    @functools.cached_property
+    def payload(self) -> dict:
+        return jwt.decode(self.token, options={"verify_signature": False})
+
+
+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
+    auth: AuthSettings = Field(default_factory=AuthSettings)
+    token: str | None = None
+
+    @property
+    def uri(self) -> str:
+        return f"{'https' if self.ssl else 'http'}://{self.host}:{self.port}"
+
+    @property
+    def uri_scandal(self) -> str:
+        # TODO(marko): Scandal will be a different service. For now, gRPC API is hosted on the SpiralDB service.
+        return f"{'grpc+tls' if self.ssl else 'grpc'}://{self.host}:{self.port}"
+
+    def device_auth(self) -> DeviceAuth:
+        auth_file = (
+            APP_DIR / hashlib.md5(f"{self.auth.domain}/{self.auth.client_id}".encode()).hexdigest() / "auth.json"
+        )
+        return DeviceAuth(auth_file=auth_file, domain=self.auth.domain, client_id=self.auth.client_id)
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        toml_file=CONFIG_FILE,
+        env_nested_delimiter="__",
+        env_prefix="SPIRAL__",
+        frozen=True,
+    )
+
+    spiraldb: SpiralDBSettings = Field(default_factory=SpiralDBSettings)
+
+    @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 authn(self):
+        if self.spiraldb.token:
+            return TokenAuthn(self.spiraldb.token)
+
+        return FallbackAuthn(
+            [
+                GitHubActionsProvider(),
+                ModalProvider(),
+                # TODO(marko): Github and Modal should also be behind token exchanged.
+                TokenExchangeProvider(EnvironmentAuthn(), self.spiraldb.uri),
+                DeviceAuthProvider(self.spiraldb.device_auth()),
+            ]
+        )
+
+    @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, TomlConfigSettingsSource(settings_cls), init_settings
+
+
+class _LazyDict(Mapping):
+    def __init__(self, lazy_values: dict[str, Any | Callable[[], Any]]):
+        self._lazy_values = lazy_values
+
+    def _dict(self):
+        return {k: v() if callable(v) else v for k, v in self._lazy_values.items()}
+
+    def __getitem__(self, key, /):
+        return self._lazy_values[key]()
+
+    def get(self, key, /, default=None):
+        return self._lazy_values.get(key, lambda: default)()
+
+    def items(self):
+        return self._dict().items()
+
+    def keys(self):
+        return self._lazy_values.keys()
+
+    def values(self):
+        return self._dict().values()
+
+    def __contains__(self, key, /):
+        return key in self._lazy_values
+
+    def __eq__(self, other, /):
+        return False
+
+    def __iter__(self):
+        return iter(self._dict())
+
+    def __len__(self):
+        return len(self._lazy_values)
+
+
+@functools.cache
+def settings() -> Settings:
+    return Settings()