pyspiral 0.6.8__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/streaming_/reader.py

@@ -0,0 +1,133 @@
+import dataclasses
+import functools
+import os
+from typing import Any
+
+import vortex as vx
+
+from spiral.core.client import Shard
+
+
+# Fake streaming.base.format.base.reader.FileInfo
+# Dataset manages decompression instead of the Stream in MDS.
+# So we return our own fake FileInfo that has None for compressed file.
+@dataclasses.dataclass
+class FileInfo:
+    basename: str
+    hashes: dict[str, str] = dataclasses.field(default_factory=dict)
+
+    @property
+    def bytes(self):
+        raise NotImplementedError("FileInfo.bytes should NOT be called.")
+
+
+class SpiralReader:
+    """
+    An MDS (streaming) compatible Reader.
+    """
+
+    def __init__(self, shard: Shard, basepath):
+        self._shard = shard
+        if shard.cardinality is None:
+            raise ValueError("Shard cardinality must be known for `streaming`.")
+        self._cardinality = shard.cardinality
+        self._basepath = basepath
+        self._scan: vx.RepeatedScan | None = None
+
+    @property
+    def shard(self) -> Shard:
+        return self._shard
+
+    @property
+    def size(self):
+        """Get the number of samples in this shard.
+
+        Returns:
+            int: Sample count.
+        """
+        return self._cardinality
+
+    @property
+    def samples(self):
+        """Get the number of samples in this shard.
+
+        Returns:
+            int: Sample count.
+        """
+        return self._cardinality
+
+    def __len__(self) -> int:
+        """Get the number of samples in this shard.
+
+        Returns:
+            int: Sample count.
+        """
+        return self._cardinality
+
+    @property
+    def file_pairs(self) -> list[tuple[FileInfo, FileInfo | None]]:
+        """Get the infos from raw and compressed file.
+
+        MDS uses this because dataset manages decompression of the shards, not stream...
+        """
+        return [(FileInfo(basename=self.filename), None)]
+
+    def get_max_size(self) -> int:
+        """Get the full size of this shard.
+
+        "Max" in this case means both the raw (decompressed) and zip (compressed) versions are
+        resident (assuming it has a zip form). This is the maximum disk usage the shard can reach.
+        When compressed was used, even if keep_zip is ``False``, the zip form must still be
+        resident at the same time as the raw form during shard decompression.
+
+        Returns:
+            int: Size in bytes.
+        """
+        # TODO(marko): This is used to check cache limit is possible...
+        return 0
+
+    @functools.cached_property
+    def filename(self) -> str:
+        """Used by SpiralStream to identify shard's file-on-disk, if it exists."""
+        # TODO(marko): This might be too long...
+        return (
+            bytes(self._shard.key_range.begin).hex()
+            + "_"
+            + bytes(self._shard.key_range.end).hex()
+            + "_"
+            + str(self._shard.cardinality)
+            + ".vortex"
+        )
+
+    @functools.cached_property
+    def filepath(self) -> str:
+        """Full path to the shard's file-on-disk, if it exists."""
+        return os.path.join(self._basepath, self.filename)
+
+    def evict(self) -> int:
+        """Remove all files belonging to this shard."""
+
+        # Clean up the scan handle first. This will make sure memory is freed.
+        self._scan = None
+
+        # Try to evict file.
+        try:
+            stat = os.stat(self.filepath)
+            os.remove(self.filepath)
+            return stat.st_size
+        except FileNotFoundError:
+            # Nothing to evict.
+            return 0
+
+    def __getitem__(self, item):
+        return self.get_item(item)
+
+    def get_item(self, idx: int) -> dict[str, Any]:
+        if self._scan is None:
+            # TODO(marko): vx.open should throw FileNotFoundError instead of
+            #   ValueError: No such file or directory (os error 2)
+            # Check if shard is ready on disk. This must throw FileNotFoundError.
+            if not os.path.exists(self.filepath):
+                raise FileNotFoundError(f"Shard not found: {self.filepath}")
+            self._scan = vx.open(self.filepath, without_segment_cache=True).to_repeated_scan()
+        return self._scan.scalar_at(idx).as_py()

spiral/streaming_/stream.py

@@ -0,0 +1,157 @@
+import os
+import tempfile
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from spiral import Scan, Spiral
+from spiral.core.client import Shard
+from spiral.streaming_.reader import SpiralReader
+
+if TYPE_CHECKING:
+    from streaming.base.array import NDArray
+    from streaming.base.format import Reader
+    from streaming.base.world import World
+
+
+class SpiralStream:
+    """
+    An MDS (streaming) compatible Stream.
+
+    The stream does not extend the default Stream class, but it is compactible with its API.
+
+    The stream is not registered with MDS, as the only way to construct the stream is through Spiral client.
+    Stream can be passed to MDS's StreamingDataset in `streams` argument.
+    """
+
+    def __init__(
+        self,
+        sp: Spiral,
+        scan: Scan,
+        shards: list[Shard],
+        cache_dir: str | None = None,
+        shard_row_block_size: int | None = None,
+    ):
+        self._sp = sp
+        self._scan = scan
+        # TODO(marko): Read shards only on world.is_local_leader in `get_shards` and materialize on disk.
+        self._shards = shards
+        self._shard_row_block_size = shard_row_block_size or 8192
+
+        if cache_dir is not None:
+            if not os.path.exists(cache_dir):
+                os.makedirs(cache_dir, exist_ok=True)
+            if not os.path.isdir(cache_dir):
+                raise ValueError(f"Cache dir {cache_dir} is not a directory.")
+        else:
+            cache_dir = os.path.join(tempfile.gettempdir(), "spiral-streaming")
+        self._cache_dir = cache_dir
+
+        # Enure split directory exists.
+        os.makedirs(os.path.join(self._cache_dir, self.split), exist_ok=True)
+
+    @property
+    def local(self) -> str:
+        # Dataset: Register/lookup our shared memory prefix and filelock root directory.
+        return self._cache_dir
+
+    @property
+    def remote(self) -> str | None:
+        # Dataset: Register/lookup our shared memory prefix and filelock root directory.
+        return None
+
+    @property
+    def split(self) -> str:
+        # Dataset: Register/lookup our shared memory prefix and filelock root directory.
+        return "default"
+
+    @classmethod
+    def validate_weights(cls, streams) -> tuple[bool, bool]:
+        from streaming.base.stream import Stream
+
+        return Stream.validate_weights(streams)
+
+    @classmethod
+    def apply_weights(cls, streams, samples_per_stream, choose_per_epoch, seed) -> int:
+        from streaming.base.stream import Stream
+
+        return Stream.apply_weights(streams, samples_per_stream, choose_per_epoch, seed)
+
+    def apply_default(self, default: dict):
+        # Applies defaults from the StreamingDataset.
+        # 'remote', 'local', 'split', 'download_retry', 'download_timeout', 'validate_hash', 'keep_zip'
+        if default["split"] is not None:
+            raise ValueError("SpiralStream does not support split, as the split is defined in the Scan.")
+
+    def prepare_shard(self, shard: "Reader") -> int:
+        """Ensure (download, validate, extract, etc.) that we have the given shard.
+
+        Args:
+            shard (Reader): Which shard.
+
+        Returns:
+            int: Change in cache usage.
+        """
+        if not isinstance(shard, SpiralReader):
+            raise ValueError("Only SpiralReader is supported in SpiralStream")
+
+        shard_path = os.path.join(self._cache_dir, self.split, shard.filename)
+        if os.path.exists(shard_path):
+            # Already exists.
+            return 0
+
+        # Prepare the shard, writing it to disk.
+        self._sp._ops().prepare_shard(
+            shard_path, self._scan.core, shard.shard, row_block_size=self._shard_row_block_size
+        )
+
+        # Get the size of the file on disk.
+        stat = os.stat(shard_path)
+        return stat.st_size
+
+    def get_shards(self, world: "World", allow_unsafe_types: bool) -> list["Reader"]:
+        """Load this Stream's index, retrieving its shard readers.
+
+        Args:
+            world (World): Distributed context.
+            allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code
+                execution during deserialization, whether to keep going if ``True`` or raise an error.
+                This argument is ignored as SpiralStream does not support Pickle.
+
+        Returns:
+            `List[Reader]: Shard readers.
+        """
+        basepath = os.path.join(self._cache_dir, self.split)
+        return [SpiralReader(shard, basepath) for shard in self._shards]  # type: ignore[return-value]
+
+    def set_up_local(self, shards: list["Reader"], cache_usage_per_shard: "NDArray[np.int64]") -> None:
+        """Bring a local directory into a consistent state, getting which shards are present.
+
+        Args:
+            shards (List[Reader]): List of this stream's shards.
+            cache_usage_per_shard (NDArray[np.int64]): Cache usage per shard of this stream.
+        """
+        listing = set()
+        for file in os.listdir(os.path.join(self._cache_dir, self.split)):
+            if os.path.isfile(os.path.join(self._cache_dir, self.split, file)) and file.endswith(".vortex"):
+                listing.add(file)
+
+        # Determine which shards are present, making local dir consistent.
+        for i, shard in enumerate(shards):
+            if not isinstance(shard, SpiralReader):
+                raise ValueError("Only SpiralReader is supported in SpiralStream")
+            if shard.filename in listing:
+                # Get the size of the file on disk.
+                stat = os.stat(os.path.join(self._cache_dir, self.split, shard.filename))
+                cache_usage_per_shard[i] = stat.st_size
+            else:
+                cache_usage_per_shard[i] = 0
+
+    def get_index_size(self) -> int:
+        """Get the size of the index file in bytes.
+
+        Returns:
+            int: Size in bytes.
+        """
+        # There is no index file stored on disk.
+        return 0

spiral/substrait_.py

@@ -0,0 +1,274 @@
+import betterproto2
+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 betterproto2.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 betterproto2.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 betterproto2.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 betterproto2.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 betterproto2.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 betterproto2.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
+                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)