pyspiral 0.7.18__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/dataloader.py

@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+import os
+import random
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from functools import partial
+from multiprocessing import Pool
+from typing import Any
+
+import pyarrow as pa
+
+from spiral.core.client import Shard
+from spiral.scan import Scan
+
+
+@dataclass(frozen=True)
+class World:
+    """Distributed training configuration.
+    Attributes:
+        rank: Process rank (0 to world_size-1).
+        world_size: Total number of processes.
+    """
+
+    rank: int
+    world_size: int
+
+    def shards(
+        self,
+        shards: list[Shard],
+        shuffle_seed: int | None = None,
+    ) -> list[Shard]:
+        """Partition shards for distributed training.
+
+        Args:
+            shards: List of Shard objects to partition.
+            shuffle_seed: Optional seed to shuffle before partitioning.
+
+        Returns:
+            Subset of shards for this rank (round-robin partitioning).
+        """
+        if shuffle_seed is not None:
+            shards = World._shuffle(shards, shuffle_seed)
+
+        return shards[self.rank :: self.world_size]
+
+    @classmethod
+    def from_torch(cls) -> World:
+        """Auto-detect world configuration from PyTorch distributed."""
+        try:
+            import torch.distributed as dist
+
+            if dist.is_available() and dist.is_initialized():
+                return cls(
+                    rank=dist.get_rank(),
+                    world_size=dist.get_world_size(),
+                )
+        except ImportError:
+            pass
+
+        return cls(
+            rank=int(os.environ.get("RANK", 0)),
+            world_size=int(os.environ.get("WORLD_SIZE", 1)),
+        )
+
+    @classmethod
+    def _shuffle(cls, shards: list[Shard], seed: int) -> list[Shard]:
+        """Shuffle shards deterministically with given seed."""
+        shuffled = list(shards)
+        random.Random(seed).shuffle(shuffled)
+        return shuffled
+
+
+# Top level so we can pickle this function
+def _len_and_transform(batch: pa.RecordBatch, transform_fn: Callable) -> tuple[int, Any]:
+    return (len(batch), transform_fn(batch))
+
+
+class SpiralDataLoader:
+    """DataLoader optimized for Spiral's multi-threaded streaming architecture.
+
+    Unlike PyTorch's DataLoader which uses multiprocessing for I/O (num_workers),
+    SpiralDataLoader leverages Spiral's efficient Rust-based streaming and only
+    uses multiprocessing for CPU-bound post-processing transforms.
+
+    Key differences from PyTorch DataLoader:
+    - No num_workers for I/O (Spiral's Rust layer is already multi-threaded)
+    - map_workers for parallel post-processing (tokenization, decoding, etc.)
+    - Built-in checkpoint support via skip_samples
+    - Explicit shard-based architecture for distributed training
+
+    Simple usage:
+    ```python
+    loader = SpiralDataLoader(scan, batch_size=32)
+    for batch in loader:
+        train_step(batch)
+    ```
+
+    With parallel transforms:
+    ```python
+    loader = SpiralDataLoader(
+        scan,
+        batch_size=32,
+        transform_fn=tokenize_batch,
+        map_workers=4,
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        scan: Scan,
+        *,
+        shards: list[Shard] | None = None,
+        shuffle_shards: bool = True,
+        seed: int = 42,
+        skip_samples: int = 0,
+        shuffle_buffer_size: int = 0,
+        batch_size: int = 32,
+        batch_readahead: int | None = None,
+        # TODO(os): accept vortex arrays here instead of Arrow
+        transform_fn: Callable[[pa.RecordBatch], Any] | None = None,
+        map_workers: int = 0,
+        infinite: bool = False,
+    ):
+        """Initialize SpiralDataLoader.
+
+        Args:
+            scan: Spiral scan to load data from.
+            shards: Optional list of Shard objects to read. If None, uses
+                scan's natural sharding based on physical layout.
+            shuffle_shards: Whether to shuffle the list of shards.
+                Uses the provided seed.
+            seed: Base random seed for deterministic shuffling and checkpointing.
+            skip_samples: Number of samples to skip at the beginning (for resuming
+                from checkpoint).
+            shuffle_buffer_size: Size of shuffle buffer for within-shard shuffling.
+                0 means no shuffling.
+            batch_size: Number of rows per batch.
+            batch_readahead: Number of batches to prefetch in background. If None,
+                uses a sensible default based on whether transforms are applied.
+            transform_fn: Optional function to transform each batch. Takes a PyArrow
+                RecordBatch and returns any type. Users can call batch.to_pydict()
+                inside the function if they need a dict. If map_workers > 0, this
+                function must be picklable.
+            map_workers: Number of worker processes for parallel transform_fn
+                application. 0 means single-process (no parallelism). Use this for
+                CPU-bound transforms like tokenization or audio decoding.
+            infinite: Whether to cycle through the dataset infinitely. If True,
+                the dataloader will repeat the dataset indefinitely. If False,
+                the dataloader will stop after going through the dataset once.
+        """
+        self.scan = scan
+        self.shards = shards if shards is not None else scan.shards()
+        if shuffle_shards:
+            self.shards = World._shuffle(self.shards, seed)
+        self.seed = seed
+        self.skip_samples = skip_samples
+        self.shuffle_buffer_size = shuffle_buffer_size
+        self.batch_size = batch_size
+        self.batch_readahead = batch_readahead
+        self.transform_fn = transform_fn
+        self.map_workers = map_workers
+        self.infinite = infinite
+
+        self._samples_yielded = 0
+
+    def __iter__(self) -> Iterator[Any]:
+        """Iterate over batches."""
+        from spiral.core.client import ShuffleConfig
+
+        shuffle = None
+        if self.shuffle_buffer_size > 0:
+            shuffle = ShuffleConfig(
+                buffer_size=self.shuffle_buffer_size,
+                seed=self.seed,
+            )
+
+        stream = self.scan.core.to_shuffled_record_batches(
+            shards=self.shards,
+            shuffle=shuffle,
+            max_batch_size=self.batch_size,
+            batch_readahead=self.batch_readahead,
+            infinite=self.infinite,
+        )
+
+        if self.skip_samples > 0:
+
+            def skip(s: Iterator[pa.RecordBatch], skip_count: int) -> Iterator[pa.RecordBatch]:
+                """Skip samples from stream, yielding remaining batches."""
+                skipped = 0
+                for batch in s:
+                    batch_size = len(batch)
+                    if skipped + batch_size <= skip_count:
+                        # Skip entire batch
+                        skipped += batch_size
+                        continue
+                    elif skipped < skip_count:
+                        # Partial skip - discard first N samples, yield remainder
+                        skip_in_batch = skip_count - skipped
+                        skipped = skip_count
+                        yield batch[skip_in_batch:]
+                    else:
+                        # take the entire batch
+                        yield batch
+
+            stream = skip(stream, self.skip_samples)
+
+        if self.transform_fn is None:
+            for batch in stream:
+                self._samples_yielded += len(batch)
+                yield batch
+        elif self.map_workers == 0:
+            # Single-process transform
+            for batch in stream:
+                result = self.transform_fn(batch)
+                self._samples_yielded += len(batch)
+                yield result
+        else:
+            with Pool(self.map_workers) as pool:
+                for batch_len, result in pool.imap(partial(_len_and_transform, transform_fn=self.transform_fn), stream):
+                    self._samples_yielded += batch_len
+                    yield result
+
+    def state_dict(self) -> dict[str, Any]:
+        """Get checkpoint state for resuming.
+
+        Returns:
+            Dictionary containing samples_yielded, seed, and shards.
+
+        Example checkpoint:
+        ```python
+        loader = SpiralDataLoader(scan, batch_size=32, seed=42)
+        for i, batch in enumerate(loader):
+            if i == 10:
+                checkpoint = loader.state_dict()
+                break
+        ```
+
+        Example resume:
+        ```python
+        loader = SpiralDataLoader.from_state_dict(scan, checkpoint, batch_size=32)
+        ```
+        """
+        return {
+            "samples_yielded": self._samples_yielded,
+            "seed": self.seed,
+            "shards": self.shards,  # Will be pickled automatically
+        }
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        scan: Scan,
+        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:
+            scan: Spiral scan to load data from.
+            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.
+
+        Save checkpoint during training:
+        ```python
+        loader = scan.to_distributed_data_loader(scan, batch_size=32, seed=42)
+        checkpoint = loader.state_dict()
+        ```
+
+        Resume later using the same shards from checkpoint:
+        ```python
+        resumed_loader = SpiralDataLoader.from_state_dict(
+            scan,
+            checkpoint,
+            batch_size=32,
+            transform_fn=my_transform,
+        )
+        """
+
+        # Extract resume parameters from state
+        seed = state.get("seed", 42)
+        skip_samples = state.get("samples_yielded", 0)
+        shards = state.get("shards")
+
+        # Allow kwargs to override state dict values
+        seed = kwargs.pop("seed", seed)
+        skip_samples = kwargs.pop("skip_samples", skip_samples)
+        shards = kwargs.pop("shards", shards)
+
+        return cls(scan, seed=seed, skip_samples=skip_samples, shards=shards, **kwargs)

spiral/dataset.py

@@ -0,0 +1,264 @@
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.dataset as ds
+
+from spiral.scan import Scan
+from spiral.snapshot import Snapshot
+
+
+class Dataset(ds.Dataset):
+    def __init__(self, snapshot: Snapshot):
+        self._snapshot = snapshot
+        self._table = snapshot.table
+        self._schema: pa.Schema = self._snapshot.schema().to_arrow()
+
+        # We don't actually initialize a Dataset, we just implement enough of the API
+        # to fool both DuckDB and Polars.
+        # super().__init__()
+        self._last_scan = None
+
+    @property
+    def schema(self) -> pa.Schema:
+        return self._schema
+
+    def count_rows(
+        self,
+        filter: pc.Expression | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        fragment_readahead: int | None = None,
+        fragment_scan_options: ds.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ):
+        return self.scanner(
+            None,
+            filter,
+            batch_size,
+            batch_readahead,
+            fragment_readahead,
+            fragment_scan_options,
+            use_threads,
+            memory_pool,
+        ).count_rows()
+
+    def filter(self, expression: pc.Expression) -> "Dataset":
+        raise NotImplementedError("filter not implemented")
+
+    def get_fragments(self, filter: pc.Expression | None = None):
+        """TODO(ngates): perhaps we should return ranges as per our split API?"""
+        raise NotImplementedError("get_fragments not implemented")
+
+    def head(
+        self,
+        num_rows: int,
+        columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        fragment_readahead: int | None = None,
+        fragment_scan_options: ds.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ):
+        return self.scanner(
+            columns,
+            filter,
+            batch_size,
+            batch_readahead,
+            fragment_readahead,
+            fragment_scan_options,
+            use_threads,
+            memory_pool,
+        ).head(num_rows)
+
+    def join(
+        self,
+        right_dataset,
+        keys,
+        right_keys=None,
+        join_type=None,
+        left_suffix=None,
+        right_suffix=None,
+        coalesce_keys=True,
+        use_threads=True,
+    ):
+        raise NotImplementedError("join not implemented")
+
+    def join_asof(self, right_dataset, on, by, tolerance, right_on=None, right_by=None):
+        raise NotImplementedError("join_asof not implemented")
+
+    def replace_schema(self, schema: pa.Schema) -> "Dataset":
+        raise NotImplementedError("replace_schema not implemented")
+
+    def scanner(
+        self,
+        columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        fragment_readahead: int | None = None,
+        fragment_scan_options: ds.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ) -> "TableScanner":
+        from spiral.substrait_ import SubstraitConverter
+
+        # Extract the substrait expression so we can convert it to a Spiral expression
+        if filter is not None:
+            filter = SubstraitConverter(self._table, self._schema, self._table.key_schema.to_arrow()).convert(
+                filter.to_substrait(self._schema, allow_arrow_extensions=True),
+            )
+
+        scan = (
+            self._table.spiral.scan(
+                {c: self._table[c] for c in columns},
+                where=filter,
+                asof=self._snapshot.asof,
+            )
+            if columns
+            else self._table.spiral.scan(
+                self._table,
+                where=filter,
+                asof=self._snapshot.asof,
+            )
+        )
+        self._last_scan = scan
+
+        return TableScanner(scan)
+
+    def sort_by(self, sorting, **kwargs):
+        raise NotImplementedError("sort_by not implemented")
+
+    def take(
+        self,
+        indices: pa.Array | Any,
+        columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        fragment_readahead: int | None = None,
+        fragment_scan_options: ds.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ):
+        return self.scanner(
+            columns,
+            filter,
+            batch_size,
+            batch_readahead,
+            fragment_readahead,
+            fragment_scan_options,
+            use_threads,
+            memory_pool,
+        ).take(indices)
+
+    def to_batches(
+        self,
+        columns: list[str] | None = None,
+        filter: pc.Expression | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        fragment_readahead: int | None = None,
+        fragment_scan_options: ds.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ):
+        return self.scanner(
+            columns,
+            filter,
+            batch_size,
+            batch_readahead,
+            fragment_readahead,
+            fragment_scan_options,
+            use_threads,
+            memory_pool,
+        ).to_batches()
+
+    def to_table(
+        self,
+        columns=None,
+        filter: pc.Expression | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+        fragment_readahead: int | None = None,
+        fragment_scan_options: ds.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ):
+        return self.scanner(
+            columns,
+            filter,
+            batch_size,
+            batch_readahead,
+            fragment_readahead,
+            fragment_scan_options,
+            use_threads,
+            memory_pool,
+        ).to_table()
+
+
+class TableScanner(ds.Scanner):
+    """A PyArrow Dataset Scanner that reads from a Spiral Table."""
+
+    def __init__(
+        self,
+        scan: Scan,
+        key_table: pa.Table | pa.RecordBatchReader | None = None,
+    ):
+        self._scan = scan
+        self._schema = scan.schema
+        self.key_table = key_table
+
+        # We don't actually initialize a Dataset, we just implement enough of the API
+        # to fool both DuckDB and Polars.
+        # super().__init__()
+
+    @property
+    def schema(self):
+        return self._schema
+
+    def count_rows(self):
+        # TODO(ngates): is there a faster way to count rows?
+        return sum(len(batch) for batch in self.to_reader())
+
+    def head(self, num_rows: int):
+        """Return the first `num_rows` rows of the dataset."""
+
+        kwargs = {}
+        if num_rows <= 10_000:
+            # We are unlikely to need more than a couple batches
+            kwargs["batch_readahead"] = 1
+            # The progress bar length is the total number of splits in this dataset. We will likely
+            # stop streaming early. As a result, the progress bar is misleading.
+            kwargs["hide_progress_bar"] = True
+
+        reader = self._scan.to_unordered_record_batches(key_table=self.key_table, **kwargs)
+        batches = []
+        row_count = 0
+        for batch in reader:
+            if row_count + len(batch) > num_rows:
+                batches.append(batch.slice(0, num_rows - row_count))
+                break
+            row_count += len(batch)
+            batches.append(batch)
+        return pa.Table.from_batches(batches, schema=reader.schema)
+
+    def scan_batches(self):
+        raise NotImplementedError("scan_batches not implemented")
+
+    def take(self, indices):
+        # TODO(ngates): can we defer take until after we've constructed the scan?
+        #  Or should this we delay constructing the Spiral Table.scan?
+        raise NotImplementedError("take not implemented")
+
+    def to_batches(self):
+        return self.to_reader()
+
+    def to_reader(self):
+        return self._scan.to_unordered_record_batches(key_table=self.key_table)
+
+    def to_table(self):
+        return self.to_reader().read_all()

spiral/datetime_.py

@@ -0,0 +1,27 @@
+import warnings
+from datetime import UTC, datetime, timedelta, tzinfo
+
+_THE_EPOCH = datetime.fromtimestamp(0, tz=UTC)
+
+
+def local_tz() -> tzinfo:
+    """Determine this machine's local timezone."""
+    tz = datetime.now().astimezone().tzinfo
+    if tz is None:
+        raise ValueError("Could not determine this machine's local timezone.")
+    return tz
+
+
+def timestamp_micros(instant: datetime) -> int:
+    """The number of microseconds between the epoch and the given instant."""
+    if instant.tzinfo is None:
+        warnings.warn("assuming timezone-naive datetime is local time", stacklevel=2)
+        instant = instant.replace(tzinfo=local_tz())
+    return (instant - _THE_EPOCH) // timedelta(microseconds=1)
+
+
+def from_timestamp_micros(ts: int) -> datetime:
+    """Convert a timestamp in microseconds to a datetime."""
+    if ts < 0:
+        raise ValueError("Timestamp must be non-negative")
+    return _THE_EPOCH + timedelta(microseconds=ts)

spiral/debug/manifests.py

@@ -0,0 +1,87 @@
+from rich.console import Console
+from rich.table import Table
+
+from spiral import datetime_
+from spiral.core.table import Scan
+from spiral.core.table.manifests import FragmentManifest
+from spiral.core.table.spec import ColumnGroup
+from spiral.debug.metrics import _format_bytes
+
+
+def display_scan_manifests(scan: Scan):
+    """Display all manifests in a scan."""
+    if len(scan.table_ids()) != 1:
+        raise NotImplementedError("Multiple table scans are not supported.")
+    table_id = scan.table_ids()[0]
+    key_space_manifest = scan.key_space_state(table_id).manifest
+    column_group_manifests = [
+        (column_group, scan.column_group_state(column_group).manifest) for column_group in scan.column_groups()
+    ]
+
+    display_manifests(key_space_manifest, column_group_manifests)
+
+
+def display_manifests(
+    key_space_manifest: FragmentManifest, column_group_manifests: list[tuple[ColumnGroup, FragmentManifest]]
+):
+    _table_of_fragments(
+        key_space_manifest,
+        title="Key Space manifest",
+    )
+
+    for column_group, column_group_manifest in column_group_manifests:
+        _table_of_fragments(
+            column_group_manifest,
+            title=f"Column Group manifest for {str(column_group)}",
+        )
+
+
+def _table_of_fragments(manifest: FragmentManifest, title: str):
+    """Display fragments in a formatted table."""
+    # Calculate summary statistics
+    total_size = sum(fragment.size_bytes for fragment in manifest)
+    total_metadata_size = sum(len(fragment.format_metadata or b"") for fragment in manifest)
+    fragment_count = len(manifest)
+    avg_size = total_size / fragment_count if fragment_count > 0 else 0
+
+    # Print title and summary
+    console = Console()
+    console.print(f"\n\n{title}")
+    console.print(
+        f"{fragment_count} fragments, "
+        f"total: {_format_bytes(total_size)}, "
+        f"avg: {_format_bytes(int(avg_size))}, "
+        f"metadata: {_format_bytes(total_metadata_size)}"
+    )
+
+    # Create rich table
+    table = Table(title=None, show_header=True, header_style="bold")
+    table.add_column("ID", style="cyan", no_wrap=True)
+    table.add_column("Size (Metadata)", justify="right")
+    table.add_column("Format", justify="center")
+    table.add_column("Key Span", justify="center")
+    table.add_column("Level", justify="center")
+    table.add_column("Committed At", justify="center")
+    table.add_column("Compacted At", justify="center")
+
+    # Add each fragment as a row
+    for fragment in manifest:
+        committed_str = str(datetime_.from_timestamp_micros(fragment.committed_at)) if fragment.committed_at else "N/A"
+        compacted_str = str(datetime_.from_timestamp_micros(fragment.compacted_at)) if fragment.compacted_at else "N/A"
+
+        size_with_metadata = (
+            f"{_format_bytes(fragment.size_bytes)} ({_format_bytes(len(fragment.format_metadata or b''))})"
+        )
+        key_span = f"{fragment.key_span.begin}..{fragment.key_span.end}"
+
+        table.add_row(
+            fragment.id,
+            size_with_metadata,
+            str(fragment.format),
+            key_span,
+            str(fragment.level),
+            committed_str,
+            compacted_str,
+        )
+
+    console.print(table)