pyspiral 0.6.3__cp310-abi3-macosx_11_0_arm64.whl → 0.6.4__cp310-abi3-macosx_11_0_arm64.whl

spiral/iterable_dataset.py

@@ -0,0 +1,106 @@
+from collections.abc import Callable, Iterator
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+
+if TYPE_CHECKING:
+    import datasets.iterable_dataset as hf  # noqa
+    import streaming  # noqa
+    import torch.utils.data as torchdata  # noqa
+
+
+def _hf_compatible_schema(schema: pa.Schema) -> pa.Schema:
+    """
+    Replace string-view and binary-view columns in the schema with strings/binary.
+    Recursively handles nested types (struct, list, etc).
+    We use this converted schema as Features in the returned Dataset.
+    Remove this method once we have https://github.com/huggingface/datasets/pull/7718
+    """
+
+    def _convert_type(dtype: pa.DataType) -> pa.DataType:
+        if dtype == pa.string_view():
+            return pa.string()
+        elif dtype == pa.binary_view():
+            return pa.binary()
+        elif pa.types.is_struct(dtype):
+            new_fields = [
+                pa.field(field.name, _convert_type(field.type), nullable=field.nullable, metadata=field.metadata)
+                for field in dtype
+            ]
+            return pa.struct(new_fields)
+        elif pa.types.is_list(dtype):
+            return pa.list_(_convert_type(dtype.value_type))
+        elif pa.types.is_large_list(dtype):
+            return pa.large_list(_convert_type(dtype.value_type))
+        elif pa.types.is_fixed_size_list(dtype):
+            return pa.list_(_convert_type(dtype.value_type), dtype.list_size)
+        elif pa.types.is_map(dtype):
+            return pa.map_(_convert_type(dtype.key_type), _convert_type(dtype.item_type))
+        else:
+            return dtype
+
+    new_fields = []
+    for field in schema:
+        new_type = _convert_type(field.type)
+        new_fields.append(pa.field(field.name, new_type, nullable=field.nullable, metadata=field.metadata))
+
+    return pa.schema(new_fields)
+
+
+def to_iterable_dataset(stream: pa.RecordBatchReader) -> "hf.IterableDataset":
+    from datasets import DatasetInfo, Features
+    from datasets.builder import ArrowExamplesIterable
+    from datasets.iterable_dataset import IterableDataset
+
+    def _generate_tables(**kwargs) -> Iterator[tuple[int, pa.Table]]:
+        # 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)
+
+    # TODO(marko): This is temporary until we stop returning IterableDataset from this function.
+    class _IterableDataset(IterableDataset):
+        # Diff with datasets.iterable_dataset.IterableDataset:
+        # - Removes torch handling which attempts to handle worker processes.
+        # - Assumes arrow iterator.
+        def __iter__(self):
+            from datasets.formatting import get_formatter
+
+            prepared_ex_iterable = self._prepare_ex_iterable_for_iteration()
+            if self._formatting and (prepared_ex_iterable.iter_arrow or self._formatting.is_table):
+                formatter = get_formatter(self._formatting.format_type, features=self.features)
+                iterator = prepared_ex_iterable.iter_arrow()
+                for key, pa_table in iterator:
+                    yield formatter.format_row(pa_table)
+                return
+
+            for key, example in prepared_ex_iterable:
+                # no need to format thanks to FormattedExamplesIterable
+                yield example
+
+        def map(self, *args, **kwargs):
+            # Map constructs a new IterableDataset, so we need to "patch" it
+            base = super().map(*args, **kwargs)
+            if isinstance(base, IterableDataset):
+                # Patch __iter__ to avoid torch handling
+                base.__class__ = _IterableDataset  # type: ignore
+            return base
+
+    class _ArrowExamplesIterable(ArrowExamplesIterable):
+        def __init__(self, generate_tables_fn: Callable[..., Iterator[tuple[int, pa.Table]]], features: Features):
+            # NOTE: generate_tables_fn type annotations are wrong, return type must be an iterable of tuples.
+            super().__init__(generate_tables_fn, kwargs={})  # type: ignore
+            self._features = features
+
+        @property
+        def is_typed(self) -> bool:
+            return True
+
+        @property
+        def features(self) -> Features:
+            return self._features
+
+    target_features = Features.from_arrow_schema(_hf_compatible_schema(stream.schema))
+    ex_iterable = _ArrowExamplesIterable(_generate_tables, target_features)
+    info = DatasetInfo(features=target_features)
+    return _IterableDataset(ex_iterable=ex_iterable, info=info)

spiral/protogen/_/arrow/flight/protocol/sql/__init__.py

@@ -69,7 +69,7 @@ import betterproto2
 
 from .....message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 

spiral/protogen/_/google/protobuf/__init__.py

@@ -21,12 +21,15 @@ __all__ = (
     "FeatureSet",
     "FeatureSetDefaults",
     "FeatureSetDefaultsFeatureSetEditionDefault",
+    "FeatureSetEnforceNamingStyle",
     "FeatureSetEnumType",
     "FeatureSetFieldPresence",
     "FeatureSetJsonFormat",
     "FeatureSetMessageEncoding",
     "FeatureSetRepeatedFieldEncoding",
     "FeatureSetUtf8Validation",
+    "FeatureSetVisibilityFeature",
+    "FeatureSetVisibilityFeatureDefaultSymbolVisibility",
     "FieldDescriptorProto",
     "FieldDescriptorProtoLabel",
     "FieldDescriptorProtoType",
@@ -54,6 +57,7 @@ __all__ = (
     "ServiceOptions",
     "SourceCodeInfo",
     "SourceCodeInfoLocation",
+    "SymbolVisibility",
     "UninterpretedOption",
     "UninterpretedOptionNamePart",
 )
@@ -66,7 +70,7 @@ import betterproto2
 
 from ...message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 
@@ -174,6 +178,14 @@ class ExtensionRangeOptionsVerificationState(betterproto2.Enum):
     UNVERIFIED = 1
 
 
+class FeatureSetEnforceNamingStyle(betterproto2.Enum):
+    ENFORCE_NAMING_STYLE_UNKNOWN = 0
+
+    STYLE2024 = 1
+
+    STYLE_LEGACY = 2
+
+
 class FeatureSetEnumType(betterproto2.Enum):
     ENUM_TYPE_UNKNOWN = 0
 
@@ -224,6 +236,32 @@ class FeatureSetUtf8Validation(betterproto2.Enum):
     NONE = 3
 
 
+class FeatureSetVisibilityFeatureDefaultSymbolVisibility(betterproto2.Enum):
+    DEFAULT_SYMBOL_VISIBILITY_UNKNOWN = 0
+
+    EXPORT_ALL = 1
+    """
+    Default pre-EDITION_2024, all UNSET visibility are export.
+    """
+
+    EXPORT_TOP_LEVEL = 2
+    """
+    All top-level symbols default to export, nested default to local.
+    """
+
+    LOCAL_ALL = 3
+    """
+    All symbols default to local.
+    """
+
+    STRICT = 4
+    """
+    All symbols local by default. Nested types cannot be exported.
+    With special case caveat for message { enum {} reserved 1 to max; }
+    This is the recommended setting for new protos.
+    """
+
+
 class FieldDescriptorProtoLabel(betterproto2.Enum):
     OPTIONAL = 1
     """
@@ -512,6 +550,22 @@ class MethodOptionsIdempotencyLevel(betterproto2.Enum):
     """
 
 
+class SymbolVisibility(betterproto2.Enum):
+    """
+    Describes the 'visibility' of a symbol with respect to the proto import
+    system. Symbols can only be imported when the visibility rules do not prevent
+    it (ex: local symbols cannot be imported).  Visibility modifiers can only set
+    on `message` and `enum` as they are the only types available to be referenced
+    from other files.
+    """
+
+    VISIBILITY_UNSET = 0
+
+    VISIBILITY_LOCAL = 1
+
+    VISIBILITY_EXPORT = 2
+
+
 @dataclass(eq=False, repr=False)
 class Any(betterproto2.Message):
     """
@@ -744,6 +798,13 @@ class DescriptorProto(betterproto2.Message):
     A given name may only be reserved once.
     """
 
+    visibility: "SymbolVisibility" = betterproto2.field(
+        11, betterproto2.TYPE_ENUM, default_factory=lambda: SymbolVisibility(0)
+    )
+    """
+    Support for `export` and `local` keywords on enums.
+    """
+
 
 default_message_pool.register_message("google.protobuf", "DescriptorProto", DescriptorProto)
 
@@ -835,6 +896,13 @@ class EnumDescriptorProto(betterproto2.Message):
     be reserved once.
     """
 
+    visibility: "SymbolVisibility" = betterproto2.field(
+        6, betterproto2.TYPE_ENUM, default_factory=lambda: SymbolVisibility(0)
+    )
+    """
+    Support for `export` and `local` keywords on enums.
+    """
+
 
 default_message_pool.register_message("google.protobuf", "EnumDescriptorProto", EnumDescriptorProto)
 
@@ -895,6 +963,9 @@ class EnumOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(7, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     uninterpreted_option: "list[UninterpretedOption]" = betterproto2.field(
@@ -942,6 +1013,9 @@ class EnumValueOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(2, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     debug_redact: "bool" = betterproto2.field(3, betterproto2.TYPE_BOOL)
@@ -1082,10 +1156,26 @@ class FeatureSet(betterproto2.Message):
         6, betterproto2.TYPE_ENUM, default_factory=lambda: FeatureSetJsonFormat(0)
     )
 
+    enforce_naming_style: "FeatureSetEnforceNamingStyle" = betterproto2.field(
+        7, betterproto2.TYPE_ENUM, default_factory=lambda: FeatureSetEnforceNamingStyle(0)
+    )
+
+    default_symbol_visibility: "FeatureSetVisibilityFeatureDefaultSymbolVisibility" = betterproto2.field(
+        8, betterproto2.TYPE_ENUM, default_factory=lambda: FeatureSetVisibilityFeatureDefaultSymbolVisibility(0)
+    )
+
 
 default_message_pool.register_message("google.protobuf", "FeatureSet", FeatureSet)
 
 
+@dataclass(eq=False, repr=False)
+class FeatureSetVisibilityFeature(betterproto2.Message):
+    pass
+
+
+default_message_pool.register_message("google.protobuf", "FeatureSet.VisibilityFeature", FeatureSetVisibilityFeature)
+
+
 @dataclass(eq=False, repr=False)
 class FeatureSetDefaults(betterproto2.Message):
     """
@@ -1340,6 +1430,9 @@ class FieldOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(21, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     feature_support: "FieldOptionsFeatureSupport | None" = betterproto2.field(
@@ -1438,6 +1531,12 @@ class FileDescriptorProto(betterproto2.Message):
     For Google-internal migration only. Do not use.
     """
 
+    option_dependency: "list[str]" = betterproto2.field(15, betterproto2.TYPE_STRING, repeated=True)
+    """
+    Names of files imported by this file purely for the purpose of providing
+    option extensions. These are excluded from the dependency list above.
+    """
+
     message_type: "list[DescriptorProto]" = betterproto2.field(4, betterproto2.TYPE_MESSAGE, repeated=True)
     """
     All top-level definitions in this file.
@@ -1465,11 +1564,17 @@ class FileDescriptorProto(betterproto2.Message):
     The supported values are "proto2", "proto3", and "editions".
 
     If `edition` is present, this value must be "editions".
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     edition: "Edition" = betterproto2.field(14, betterproto2.TYPE_ENUM, default_factory=lambda: Edition(0))
     """
     The edition of the proto file.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
 
@@ -1665,6 +1770,9 @@ class FileOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(50, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     uninterpreted_option: "list[UninterpretedOption]" = betterproto2.field(
@@ -1817,6 +1925,9 @@ class MessageOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(12, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     uninterpreted_option: "list[UninterpretedOption]" = betterproto2.field(
@@ -1889,6 +2000,9 @@ class MethodOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(35, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     uninterpreted_option: "list[UninterpretedOption]" = betterproto2.field(
@@ -1921,6 +2035,9 @@ class OneofOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(1, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     uninterpreted_option: "list[UninterpretedOption]" = betterproto2.field(
@@ -1955,6 +2072,9 @@ class ServiceOptions(betterproto2.Message):
     features: "FeatureSet | None" = betterproto2.field(34, betterproto2.TYPE_MESSAGE, optional=True)
     """
     Any features defined in the specific edition.
+    WARNING: This field should only be used by protobuf plugins or special
+    cases like the proto compiler. Other uses are discouraged and
+    developers should rely on the protoreflect APIs for their client language.
     """
 
     deprecated: "bool" = betterproto2.field(33, betterproto2.TYPE_BOOL)

spiral/protogen/_/scandal/__init__.py

@@ -25,7 +25,7 @@ import grpc
 
 from ..message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 

spiral/protogen/_/spfs/__init__.py

@@ -16,7 +16,7 @@ import betterproto2
 
 from ..message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 

spiral/protogen/_/spql/__init__.py

@@ -17,7 +17,7 @@ import betterproto2
 
 from ..message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 

spiral/protogen/_/substrait/__init__.py

@@ -246,7 +246,7 @@ import betterproto2
 
 from ..message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 

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

@@ -18,7 +18,7 @@ import betterproto2
 
 from ...message_pool import default_message_pool
 
-_COMPILER_VERSION = "0.8.0"
+_COMPILER_VERSION = "0.9.0"
 betterproto2.check_compiler_version(_COMPILER_VERSION)
 
 

spiral/scan.py

@@ -1,4 +1,3 @@
-from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
@@ -120,8 +119,11 @@ class Scan:
         self,
         shuffle: ShuffleStrategy | None = None,
         batch_readahead: int | None = None,
+        num_workers: int | None = None,
+        worker_id: int | None = None,
+        infinite: bool = False,
     ) -> "hf.IterableDataset":
-        """Returns an Huggingface's IterableDataset.
+        """Returns a Huggingface's IterableDataset.
 
         Requires `datasets` package to be installed.
 
@@ -130,39 +132,25 @@ class Scan:
             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.
+            num_workers: If not None, shards the scan across multiple workers.
+                Must be used together with worker_id.
+            worker_id: If not None, the id of the current worker.
+                Scan will only return a subset of the data corresponding to the worker_id.
+            infinite: If True, the returned IterableDataset will loop infinitely over the data,
+                re-shuffling ranges after exhausting all data.
         """
-        from datasets import DatasetInfo, Features
-        from datasets.iterable_dataset import ArrowExamplesIterable, IterableDataset
-
-        def _generate_tables(**kwargs) -> Iterator[tuple[int, pa.Table]]:
-            stream = self.core.to_shuffled_record_batches(
-                shuffle,
-                batch_readahead,
-            )
-
-            # 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)
-
-        def _hf_compatible_schema(schema: pa.Schema) -> pa.Schema:
-            """
-            Replace string-view columns in the schema with strings. We do use this converted schema
-            as Features in the returned Dataset.
-            Remove this method once we have https://github.com/huggingface/datasets/pull/7718
-            """
-            new_fields = [
-                pa.field(field.name, pa.string(), nullable=field.nullable, metadata=field.metadata)
-                if field.type == pa.string_view()
-                else field
-                for field in schema
-            ]
-            return pa.schema(new_fields)
-
-        # NOTE: generate_tables_fn type annotations are wrong, return type must be an iterable of tuples.
-        ex_iterable = ArrowExamplesIterable(generate_tables_fn=_generate_tables, kwargs={})  # type: ignore
-        info = DatasetInfo(features=Features.from_arrow_schema(_hf_compatible_schema(self.schema.to_arrow())))
-        return IterableDataset(ex_iterable=ex_iterable, info=info)
+
+        stream = self.core.to_shuffled_record_batches(
+            shuffle,
+            batch_readahead,
+            num_workers,
+            worker_id,
+            infinite,
+        )
+
+        from spiral.iterable_dataset import to_iterable_dataset
+
+        return to_iterable_dataset(stream)
 
     def _splits(self) -> list[KeyRange]:
         # Splits the scan into a set of key ranges.

spiral/settings.py

@@ -24,6 +24,8 @@ 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):

spiral/snapshot.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING
 
+from spiral import ShuffleStrategy
 from spiral.core.table import Snapshot as CoreSnapshot
 from spiral.core.table.spec import Schema
 from spiral.types_ import Timestamp
@@ -8,6 +9,7 @@ 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
 
@@ -53,3 +55,17 @@ class Snapshot:
         import duckdb
 
         return duckdb.from_arrow(self.to_dataset())
+
+    def to_iterable_dataset(
+        self,
+        *,
+        shuffle: ShuffleStrategy | None = None,
+        batch_readahead: int | None = None,
+        infinite: bool = False,
+    ) -> "torchdata.IterableDataset":
+        """Returns an iterable dataset compatible with `torch.IterableDataset`.
+
+        See `Table` docs for details on the parameters.
+        """
+        # TODO(marko): WIP.
+        raise NotImplementedError

spiral/streaming_/stream.py

@@ -25,12 +25,16 @@ class SpiralStream:
     """
 
     def __init__(
-        self, scan: CoreScan, shards: list[Shard], cache_dir: str | None = None, shard_row_block_size: int = 8192
+        self,
+        scan: CoreScan,
+        shards: list[Shard],
+        cache_dir: str | None = None,
+        shard_row_block_size: int | None = None,
     ):
         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
+        self._shard_row_block_size = shard_row_block_size or 8192
 
         if cache_dir is not None:
             if not os.path.exists(cache_dir):
@@ -99,7 +103,7 @@ class SpiralStream:
             shard_path,
             shard.shard.key_range,
             expected_cardinality=shard.shard.cardinality,
-            shard_row_block_size=self.shard_row_block_size,
+            shard_row_block_size=self._shard_row_block_size,
         )
 
         # Get the size of the file on disk.