pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/protogen/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/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/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,101 @@
+import functools
+import os
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from pydantic import Field, ValidatorFunctionWrapHandler, WrapValidator
+from pydantic_settings import (
+    BaseSettings,
+    InitSettingsSource,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+)
+
+from spiral.api import SpiralAPI
+from spiral.core.client import Authn, DeviceCodeAuth, Token
+
+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"
+
+
+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 authn(self):
+        if self.spiraldb.token:
+            return Authn.from_token(self.spiraldb.token)
+        return Authn.from_fallback()
+
+    @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/substrait_.py

@@ -0,0 +1,279 @@
+import betterproto
+import pyarrow as pa
+
+import spiral.expressions as se
+from spiral.expressions.base import Expr
+from spiral.protogen.substrait import (
+    Expression,
+    ExpressionFieldReference,
+    ExpressionLiteral,
+    ExpressionLiteralList,
+    ExpressionLiteralStruct,
+    ExpressionLiteralUserDefined,
+    ExpressionMaskExpression,
+    ExpressionReferenceSegment,
+    ExpressionReferenceSegmentListElement,
+    ExpressionReferenceSegmentStructField,
+    ExpressionScalarFunction,
+    ExtendedExpression,
+)
+from spiral.protogen.substrait.extensions import (
+    SimpleExtensionDeclaration,
+    SimpleExtensionDeclarationExtensionFunction,
+    SimpleExtensionDeclarationExtensionType,
+    SimpleExtensionDeclarationExtensionTypeVariation,
+)
+
+
+class SubstraitConverter:
+    def __init__(self, scope: Expr, schema: pa.Schema, key_schema: pa.Schema):
+        self.scope = scope
+        self.schema = schema
+        self.key_names = set(key_schema.names)
+
+        # Extension URIs, keyed by extension URI anchor
+        self.extension_uris = {}
+
+        # Functions, keyed by function_anchor
+        self.functions = {}
+
+        # Types, keyed by type anchor
+        self.type_factories = {}
+
+    def convert(self, buffer: pa.Buffer) -> Expr:
+        """Convert a Substrait Extended Expression into a Spiral expression."""
+
+        expr: ExtendedExpression = ExtendedExpression().parse(buffer)
+        assert len(expr.referred_expr) == 1, "Only one expression is supported"
+
+        # Parse the extension URIs from the plan.
+        for ext_uri in expr.extension_uris:
+            self.extension_uris[ext_uri.extension_uri_anchor] = ext_uri.uri
+
+        # Parse the extensions from the plan.
+        for ext in expr.extensions:
+            self._extension_declaration(ext)
+
+        # Convert the expression
+        return self._expr(expr.referred_expr[0].expression)
+
+    def _extension_declaration(self, ext: SimpleExtensionDeclaration):
+        match betterproto.which_one_of(ext, "mapping_type"):
+            case "extension_function", ext_func:
+                self._extension_function(ext_func)
+            case "extension_type", ext_type:
+                self._extension_type(ext_type)
+            case "extension_type_variation", ext_type_variation:
+                self._extension_type_variation(ext_type_variation)
+            case _:
+                raise AssertionError("Invalid substrait plan")
+
+    def _extension_function(self, ext: SimpleExtensionDeclarationExtensionFunction):
+        ext_uri: str = self.extension_uris[ext.extension_uri_reference]
+        match ext_uri:
+            case "https://github.com/substrait-io/substrait/blob/main/extensions/functions_boolean.yaml":
+                match ext.name:
+                    case "or":
+                        self.functions[ext.function_anchor] = se.or_
+                    case "and":
+                        self.functions[ext.function_anchor] = se.and_
+                    case "xor":
+                        self.functions[ext.function_anchor] = se.xor
+                    case "not":
+                        self.functions[ext.function_anchor] = se.not_
+                    case _:
+                        raise NotImplementedError(f"Function name {ext.name} not supported")
+            case "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml":
+                match ext.name:
+                    case "equal":
+                        self.functions[ext.function_anchor] = se.eq
+                    case "not_equal":
+                        self.functions[ext.function_anchor] = se.neq
+                    case "lt":
+                        self.functions[ext.function_anchor] = se.lt
+                    case "lte":
+                        self.functions[ext.function_anchor] = se.lte
+                    case "gt":
+                        self.functions[ext.function_anchor] = se.gt
+                    case "gte":
+                        self.functions[ext.function_anchor] = se.gte
+                    case "is_null":
+                        self.functions[ext.function_anchor] = se.is_null
+                    case "is_not_null":
+                        self.functions[ext.function_anchor] = se.is_not_null
+                    case _:
+                        raise NotImplementedError(f"Function name {ext.name} not supported")
+            case uri:
+                raise NotImplementedError(f"Function extension URI {uri} not supported")
+
+    def _extension_type(self, ext: SimpleExtensionDeclarationExtensionType):
+        ext_uri: str = self.extension_uris[ext.extension_uri_reference]
+        match ext_uri:
+            case "https://github.com/apache/arrow/blob/main/format/substrait/extension_types.yaml":
+                match ext.name:
+                    case "null":
+                        self.type_factories[ext.type_anchor] = pa.null
+                    case "interval_month_day_nano":
+                        self.type_factories[ext.type_anchor] = pa.month_day_nano_interval
+                    case "u8":
+                        self.type_factories[ext.type_anchor] = pa.uint8
+                    case "u16":
+                        self.type_factories[ext.type_anchor] = pa.uint16
+                    case "u32":
+                        self.type_factories[ext.type_anchor] = pa.uint32
+                    case "u64":
+                        self.type_factories[ext.type_anchor] = pa.uint64
+                    case "fp16":
+                        self.type_factories[ext.type_anchor] = pa.float16
+                    case "date_millis":
+                        self.type_factories[ext.type_anchor] = pa.date64
+                    case "time_seconds":
+                        self.type_factories[ext.type_anchor] = lambda: pa.time32("s")
+                    case "time_millis":
+                        self.type_factories[ext.type_anchor] = lambda: pa.time32("ms")
+                    case "time_nanos":
+                        self.type_factories[ext.type_anchor] = lambda: pa.time64("ns")
+                    case "large_string":
+                        self.type_factories[ext.type_anchor] = pa.large_string
+                    case "large_binary":
+                        self.type_factories[ext.type_anchor] = pa.large_binary
+                    case "decimal256":
+                        self.type_factories[ext.type_anchor] = pa.decimal256
+                    case "large_list":
+                        self.type_factories[ext.type_anchor] = pa.large_list
+                    case "fixed_size_list":
+                        self.type_factories[ext.type_anchor] = pa.list_
+                    case "duration":
+                        self.type_factories[ext.type_anchor] = pa.duration
+            case uri:
+                raise NotImplementedError(f"Type extension URI {uri} not support")
+
+    def _extension_type_variation(self, ext: SimpleExtensionDeclarationExtensionTypeVariation):
+        raise NotImplementedError()
+
+    def _expr(self, expr: Expression) -> Expr:
+        match betterproto.which_one_of(expr, "rex_type"):
+            case "literal", e:
+                return self._expr_literal(e)
+            case "selection", e:
+                return self._expr_selection(e)
+            case "scalar_function", e:
+                return self._expr_scalar_function(e)
+            case "window_function", _:
+                raise ValueError("Window functions are not supported in Spiral push-down")
+            case "if_then", e:
+                return self._expr_if_then(e)
+            case "switch", e:
+                return self._expr_switch(e)
+            case "singular_or_list", _:
+                raise ValueError("singular_or_list is not supported in Spiral push-down")
+            case "multi_or_list", _:
+                raise ValueError("multi_or_list is not supported in Spiral push-down")
+            case "cast", e:
+                return self._expr_cast(e)
+            case "subquery", _:
+                raise ValueError("Subqueries are not supported in Spiral push-down")
+            case "nested", e:
+                return self._expr_nested(e)
+            case _:
+                raise NotImplementedError(f"Expression type {expr.rex_type} not implemented")
+
+    def _expr_literal(self, expr: ExpressionLiteral):
+        # TODO(ngates): the Spiral literal expression is quite weakly typed...
+        #  Maybe we can switch to Vortex?
+        simple = {
+            "boolean",
+            "i8",
+            "i16",
+            "i32",
+            "i64",
+            "fp32",
+            "fp64",
+            "string",
+            "binary",
+            "fixed_char",
+            "var_char",
+            "fixed_binary",
+        }
+
+        match betterproto.which_one_of(expr, "literal_type"):
+            case type_, v if type_ in simple:
+                return se.scalar(pa.scalar(v))
+            case "timestamp", v:
+                return se.scalar(pa.scalar(v, type=pa.timestamp("us")))
+            case "date", v:
+                return se.scalar(pa.scalar(v, type=pa.date32()))
+            case "time", v:
+                # Substrait time is us since midnight. PyArrow only supports ms.
+                v: int
+                v = int(v / 1000)
+                return se.scalar(pa.scalar(v, type=pa.time32("ms")))
+            case "null", _null_type:
+                # We need a typed null value
+                raise NotImplementedError()
+            case "struct", v:
+                v: ExpressionLiteralStruct
+                # Hmm, v has fields, but no field names. I guess we return a list and the type is applied later?
+                raise NotImplementedError()
+            case "list", v:
+                v: ExpressionLiteralList
+                return pa.scalar([self._expr_literal(e) for e in v.values])
+            case "user_defined", v:
+                v: ExpressionLiteralUserDefined
+                raise NotImplementedError()
+            case literal_type, _:
+                raise NotImplementedError(f"Literal type not supported: {literal_type}")
+
+    def _expr_selection(self, expr: ExpressionFieldReference):
+        match betterproto.which_one_of(expr, "root_type"):
+            case "root_reference", _:
+                # The reference is relative to the root
+                base_expr = self.scope
+                base_type = pa.struct(self.schema)
+            case _:
+                raise NotImplementedError("Only root_reference expressions are supported")
+
+        match betterproto.which_one_of(expr, "reference_type"):
+            case "direct_reference", direct_ref:
+                return self._expr_direct_reference(base_expr, base_type, direct_ref)
+            case "masked_reference", masked_ref:
+                return self._expr_masked_reference(base_expr, base_type, masked_ref)
+            case _:
+                raise NotImplementedError()
+
+    def _expr_direct_reference(self, scope: Expr, scope_type: pa.StructType, expr: ExpressionReferenceSegment):
+        match betterproto.which_one_of(expr, "reference_type"):
+            case "map_key", ref:
+                raise NotImplementedError("Map types not yet supported in Spiral")
+            case "struct_field", ref:
+                ref: ExpressionReferenceSegmentStructField
+                field_name = scope_type.field(ref.field).name
+
+                if field_name in self.key_names:
+                    # This is a key column, so we need to select it from the scope.
+                    return se.key(field_name)
+
+                scope = se.getitem(scope, field_name)
+                scope_type = scope_type.field(ref.field).type
+                if ref.is_set("child"):
+                    return self._expr_direct_reference(scope, scope_type, ref.child)
+                return scope
+            case "list_element", ref:
+                ref: ExpressionReferenceSegmentListElement
+                scope = se.getitem(scope, ref.offset)
+                scope_type = scope_type.field(ref.field).type
+                if ref.is_set("child"):
+                    return self._expr_direct_reference(scope, scope_type, ref.child)
+                return scope
+            case "", ref:
+                # Because Proto... we hit this case when we recurse into a child node and it's actually "None".
+                return scope
+            case _:
+                raise NotImplementedError()
+
+    def _expr_masked_reference(self, scope: Expr, scope_type: pa.StructType, expr: ExpressionMaskExpression):
+        raise NotImplementedError("Masked references are not yet supported in Spiral push-down")
+
+    def _expr_scalar_function(self, expr: ExpressionScalarFunction):
+        args = [self._expr(arg.value) for arg in expr.arguments]
+        return self.functions[expr.function_reference](*args)

spiral/tables/__init__.py

@@ -0,0 +1,12 @@
+from spiral import _lib
+from spiral.tables.client import Tables
+from spiral.tables.maintenance import Maintenance
+from spiral.tables.scan import Scan
+from spiral.tables.snapshot import Snapshot
+from spiral.tables.table import Table
+from spiral.tables.transaction import Transaction
+
+# Eagerly import the Spiral library
+assert _lib, "Spiral library"
+
+__all__ = ["Tables", "Table", "Snapshot", "Scan", "Transaction", "Maintenance"]