pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/tables/debug/scan.py

@@ -0,0 +1,248 @@
+from datetime import datetime
+
+from spiral.core.table import TableScan
+from spiral.core.table.manifests import FragmentFile, FragmentManifest
+from spiral.core.table.spec import Key
+from spiral.types_ import Timestamp
+
+
+def show_scan(scan: TableScan):
+    """Displays a scan in a way that is useful for debugging."""
+    table_ids = scan.table_ids()
+    if len(table_ids) > 1:
+        raise NotImplementedError("Multiple table scan is not supported.")
+    table_id = table_ids[0]
+    column_groups = scan.column_groups()
+
+    splits = scan.split()
+    key_space_scan = scan.key_space_scan(table_id)
+
+    # Collect all key bounds from all manifests. This makes sure all visualizations are aligned.
+    key_points = set()
+    key_space_manifest = key_space_scan.manifest
+    for i in range(len(key_space_manifest)):
+        fragment_file = key_space_manifest[i]
+        key_points.add(fragment_file.key_extent.min)
+        key_points.add(fragment_file.key_extent.max)
+    for cg in column_groups:
+        cg_scan = scan.column_group_scan(cg)
+        cg_manifest = cg_scan.manifest
+        for i in range(len(cg_manifest)):
+            fragment_file = cg_manifest[i]
+            key_points.add(fragment_file.key_extent.min)
+            key_points.add(fragment_file.key_extent.max)
+
+    # Make sure split points exist in all key points.
+    for s in splits[:-1]:  # Don't take the last end.
+        key_points.add(s.end)
+    key_points = list(sorted(key_points))
+
+    show_manifest(key_space_manifest, scope="Key space", key_points=key_points, splits=splits)
+    for cg in scan.column_groups():
+        cg_scan = scan.column_group_scan(cg)
+        # Skip table id from the start of the column group.
+        show_manifest(cg_scan.manifest, scope=".".join(cg.path[1:]), key_points=key_points, splits=splits)
+
+
+def show_manifest(manifest: FragmentManifest, scope: str = None, key_points: list[Key] = None, splits: list = None):
+    try:
+        import matplotlib.patches as patches
+        import matplotlib.pyplot as plt
+    except ImportError:
+        raise ImportError("matplotlib is required for debug")
+
+    total_fragments = len(manifest)
+
+    size_points = set()
+    for i in range(total_fragments):
+        manifest_file: FragmentFile = manifest[i]
+        size_points.add(manifest_file.size_bytes)
+    size_points = list(sorted(size_points))
+
+    if key_points is None:
+        key_points = set()
+
+        for i in range(total_fragments):
+            manifest_file: FragmentFile = manifest[i]
+
+            key_points.add(manifest_file.key_extent.min)
+            key_points.add(manifest_file.key_extent.max)
+
+        if splits is not None:
+            for split in splits[:-1]:
+                key_points.add(split.end)
+
+        key_points = list(sorted(key_points))
+
+    # Create figure and axis with specified size
+    fig, ax = plt.subplots(figsize=(12, 8))
+
+    # Plot each rectangle
+    for i in range(total_fragments):
+        manifest_file: FragmentFile = manifest[i]
+
+        left = key_points.index(manifest_file.key_extent.min)
+        right = key_points.index(manifest_file.key_extent.max)
+        height = size_points.index(manifest_file.size_bytes) + 1
+
+        color = _get_fragment_color(manifest_file, i, total_fragments)
+
+        # Create rectangle patch
+        rect = patches.Rectangle(
+            (left, 0),  # (x, y)
+            right - left,  # width
+            height,  # height
+            facecolor=color,  # fill color
+            edgecolor="black",  # border color
+            alpha=0.5,  # transparency
+            linewidth=1,  # border width
+            label=manifest_file.id,  # label for legend
+        )
+
+        ax.add_patch(rect)
+
+    # Set axis limits with some padding
+    ax.set_xlim(-0.5, len(key_points) - 1 + 0.5)
+    ax.set_ylim(-0.5, len(size_points) + 0.5)
+
+    # Create split markers on x-axis
+    if splits is not None:
+        split_positions = [key_points.index(split.end) for split in splits[:-1]]
+
+        # Add split markers at the bottom
+        for pos in split_positions:
+            ax.annotate("▲", xy=(pos, 0), ha="center", va="top", color="red", annotation_clip=False)
+
+    # Add grid
+    ax.grid(True, linestyle="--", alpha=0.7, zorder=0)
+
+    # Add labels and title
+    ax.set_title("Fragment Distribution" if scope is None else f"{scope} Fragment Distribution")
+    ax.set_xlabel("Key Index")
+    ax.set_ylabel("Size Index")
+
+    # Add legend
+    ax.legend(bbox_to_anchor=(1, 1), loc="upper left", fontsize="small")
+
+    # Adjust layout to prevent label cutoff
+    plt.tight_layout()
+
+    plot = FragmentManifestPlot(fig, ax, manifest)
+    fig.canvas.mpl_connect("motion_notify_event", plot.hover)
+
+    plt.show()
+
+
+def _get_fragment_color(manifest_file: FragmentFile, color_index, total_colors):
+    import matplotlib.cm as cm
+
+    if manifest_file.compacted_at is not None:
+        # Use a shade of gray for compacted fragments
+        # Vary the shade based on the index to distinguish different compacted fragments
+        gray_value = 0.3 + (0.5 * (color_index / total_colors))
+        return (gray_value, gray_value, gray_value)
+    else:
+        # Use viridis colormap for non-compacted fragments
+        return cm.viridis(color_index / total_colors)
+
+
+def _get_fragment_legend(manifest_file: FragmentFile):
+    return "\n".join(
+        [
+            f"id: {manifest_file.id}",
+            f"size: {manifest_file.size_bytes:,} bytes",
+            f"key_span: {manifest_file.key_span}",
+            f"key_min: {manifest_file.key_extent.min}",
+            f"key_max: {manifest_file.key_extent.max}",
+            f"format: {manifest_file.format}",
+            f"level: {manifest_file.level}",
+            f"committed_at: {_format_timestamp(manifest_file.committed_at)}",
+            f"compacted_at: {_format_timestamp(manifest_file.compacted_at)}",
+            f"ks_id: {manifest_file.ks_id}",
+        ]
+    )
+
+
+def _format_timestamp(ts: Timestamp | None) -> str:
+    # Format timestamp or show None
+    if ts is None:
+        return "None"
+    try:
+        return datetime.fromtimestamp(ts / 1e6).strftime("%Y-%m-%d %H:%M:%S")
+    except ValueError:
+        return str(ts)
+
+
+class FragmentManifestPlot:
+    def __init__(self, fig, ax, manifest: FragmentManifest):
+        self.fig = fig
+        self.ax = ax
+        self.manifest = manifest
+
+        # Position the annotation in the bottom right corner
+        self.annotation = ax.annotate(
+            "",
+            xy=(0.98, 0.02),  # Position in axes coordinates
+            xycoords="axes fraction",
+            bbox=dict(boxstyle="round,pad=0.5", fc="white", ec="gray", alpha=0.8),
+            ha="right",  # Right-align text
+            va="bottom",  # Bottom-align text
+            visible=False,
+        )
+        self.highlighted_rect = None
+        self.highlighted_legend = None
+
+    def hover(self, event):
+        if event.inaxes != self.ax:
+            # Check if we're hovering over the legend
+            legend = self.ax.get_legend()
+            if legend and legend.contains(event)[0]:
+                # Find which legend item we're hovering over
+                for i, legend_text in enumerate(legend.get_texts()):
+                    if legend_text.contains(event)[0]:
+                        manifest_file = self.manifest[i]
+                        self._show_legend(manifest_file, i, legend_text)
+                        return
+            self._hide_legend()
+            return
+
+        # Check rectangles in the main plot
+        for i, rect in enumerate(self.ax.patches):
+            if rect.contains(event)[0]:
+                manifest_file = self.manifest[i]
+                self._show_legend(manifest_file, i, rect)
+                return
+
+        self._hide_legend()
+
+    def _show_legend(self, manifest_file, index, highlight_obj):
+        import matplotlib.patches as patches
+
+        # Update tooltip text
+        self.annotation.set_text(_get_fragment_legend(manifest_file))
+        self.annotation.set_visible(True)
+
+        # Handle highlighting
+        if isinstance(highlight_obj, patches.Rectangle):
+            # Highlighting rectangle in main plot
+            if self.highlighted_rect and self.highlighted_rect != highlight_obj:
+                self.highlighted_rect.set_alpha(0.5)
+            highlight_obj.set_alpha(0.8)
+            self.highlighted_rect = highlight_obj
+        else:
+            # Highlighting legend text
+            if self.highlighted_rect:
+                self.highlighted_rect.set_alpha(0.5)
+            # Find and highlight corresponding rectangle
+            rect = self.ax.patches[index]
+            rect.set_alpha(0.8)
+            self.highlighted_rect = rect
+
+        self.fig.canvas.draw_idle()
+
+    def _hide_legend(self):
+        if self.annotation.get_visible():
+            self.annotation.set_visible(False)
+            if self.highlighted_rect:
+                self.highlighted_rect.set_alpha(0.5)
+            self.fig.canvas.draw_idle()

spiral/tables/maintenance.py

@@ -0,0 +1,12 @@
+from spiral.core.table import TableMaintenance
+
+
+class Maintenance:
+    """Spiral table maintenance."""
+
+    def __init__(self, maintenance: TableMaintenance):
+        self._maintenance = maintenance
+
+    def flush_wal(self):
+        """Flush the write-ahead log."""
+        self._maintenance.flush_wal()

spiral/tables/scan.py

@@ -0,0 +1,193 @@
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+from datasets import DatasetInfo, Features
+
+from spiral.core.table import KeyRange, TableScan
+from spiral.core.table.spec import Schema
+from spiral.settings import CI, DEV
+
+if TYPE_CHECKING:
+    import dask.dataframe as dd
+    import pandas as pd
+    import polars as pl
+    from datasets import iterable_dataset
+
+
+class Scan:
+    """Scan object."""
+
+    def __init__(
+        self,
+        scan: TableScan,
+    ):
+        # NOTE(ngates): this API is a little weird. e.g. if the query doesn't define an asof, it is resolved
+        #  when we wrap it into a core.Scan. Should we expose a Query object in the Python API that's reusable
+        #  and will re-resolve the asof? Or should we just expose a scan that fixes the asof at construction time?
+        self._scan = scan
+
+    @property
+    def metrics(self) -> dict[str, Any]:
+        """Returns metrics about the scan."""
+        return self._scan.metrics()
+
+    @property
+    def schema(self) -> Schema:
+        """Returns the schema of the scan."""
+        return self._scan.schema()
+
+    def is_empty(self) -> bool:
+        """Check if the Spiral is empty for the given key range.
+
+        **IMPORTANT**: False negatives are possible, but false positives are not,
+            i.e. is_empty can return False and scan can return zero rows.
+        """
+        return self._scan.is_empty()
+
+    def to_record_batches(
+        self,
+        key_table: pa.Table | pa.RecordBatchReader | None = None,
+        batch_size: int | None = None,
+        batch_readahead: int | None = None,
+    ) -> pa.RecordBatchReader:
+        """Read as a stream of RecordBatches.
+
+        Args:
+            key_table: a table of keys to "take" (including aux columns for cell-push-down).
+                If None, the scan will be executed without a key table.
+            batch_size: the maximum number of rows per returned batch.
+                IMPORTANT: This is currently only respected when the key_table is used. If key table is a
+                    RecordBatchReader, the batch_size argument must be None, and the existing batching is respected.
+            batch_readahead: the number of batches to prefetch in the background.
+        """
+        if isinstance(key_table, pa.RecordBatchReader):
+            if batch_size is not None:
+                raise ValueError(
+                    "batch_size must be None when key_table is a RecordBatchReader, the existing batching is respected."
+                )
+        elif isinstance(key_table, pa.Table):
+            key_table = key_table.to_reader(max_chunksize=batch_size)
+
+        return self._scan.to_record_batches(key_table=key_table, batch_readahead=batch_readahead)
+
+    def to_table(
+        self,
+        key_table: pa.Table | pa.RecordBatchReader | None = None,
+    ) -> pa.Table:
+        """Read into a single PyArrow Table.
+
+        Args:
+            key_table: a table of keys to "take" (including aux columns for cell-push-down).
+                If None, the scan will be executed without a key table.
+        """
+        # NOTE: Evaluates fully on Rust side which improved debuggability.
+        if DEV and not CI and key_table is None:
+            rb = self._scan.to_record_batch()
+            return pa.Table.from_batches([rb])
+
+        return self.to_record_batches(key_table=key_table).read_all()
+
+    def to_dask(self) -> "dd.DataFrame":
+        """Read into a Dask DataFrame.
+
+        Requires the `dask` package to be installed.
+        """
+        import dask.dataframe as dd
+        import pandas as pd
+
+        def _read_key_range(key_range: KeyRange) -> pd.DataFrame:
+            # TODO(ngates): we need a way to preserve the existing asofs? Should we copy CoreScan instead of Query?
+            raise NotImplementedError()
+
+        # Fetch a set of partition ranges
+        return dd.from_map(_read_key_range, self.split())
+
+    def to_pandas(self) -> "pd.DataFrame":
+        """Read into a Pandas DataFrame.
+
+        Requires the `pandas` package to be installed.
+        """
+        return self.to_table().to_pandas()
+
+    def to_polars(self) -> "pl.DataFrame":
+        """Read into a Polars DataFrame.
+
+        Requires the `polars` package to be installed.
+        """
+        import polars as pl
+
+        # TODO(marko): This should support lazy dataframe.
+        return pl.from_arrow(self.to_record_batches())
+
+    def to_pytorch(
+        self,
+        batch_readahead: int | None = None,
+        shuffle_batch_size: int | None = None,
+        shuffle_pool_num_rows: int | None = None,
+    ) -> "iterable_dataset.IterableDataset":
+        """Returns an iterable dataset that can be used to build a PyTorch DataLoader.
+
+        Args:
+            batch_readahead: Number of batches to prefetch in the background.
+            shuffle_batch_size: read granularity of number of rows for a shuffled scan. If left as
+            None along with shuffle_pool_num_rows=None, shuffling is disabled.
+            shuffle_pool_num_rows: Pool size for shuffling batches.
+        """
+        from datasets.iterable_dataset import ArrowExamplesIterable, IterableDataset
+
+        def _generate_tables(**kwargs) -> Iterator[tuple[int, pa.Table]]:
+            if shuffle_batch_size is None and shuffle_pool_num_rows is None:
+                stream = self.to_record_batches(
+                    batch_readahead=batch_readahead,
+                )
+            else:
+                stream = self._scan.to_shuffled_record_batches(
+                    batch_readahead, shuffle_batch_size, shuffle_pool_num_rows
+                )
+
+            # 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={})
+        info = DatasetInfo(features=Features.from_arrow_schema(_hf_compatible_schema(self.schema.to_arrow())))
+        return IterableDataset(ex_iterable=ex_iterable, info=info)
+
+    def _split(self) -> list[KeyRange]:
+        # Splits the scan into a set of key ranges.
+        return self._scan.split()
+
+    def _debug(self):
+        # Visualizes the scan, mainly for debugging purposes.
+        from spiral.tables.debug.scan import show_scan
+
+        show_scan(self._scan)
+
+    def _dump_manifests(self):
+        # Print manifests in a human-readable format.
+        from spiral.tables.debug.manifests import display_manifests
+
+        display_manifests(self._scan)
+
+    def _dump_metrics(self):
+        # Print metrics in a human-readable format.
+        from spiral.tables.debug.metrics import display_metrics
+
+        display_metrics(self.metrics)

spiral/tables/snapshot.py

@@ -0,0 +1,78 @@
+from typing import TYPE_CHECKING
+
+from spiral.core.table import TableSnapshot
+from spiral.expressions import ExprLike
+from spiral.tables.scan import Scan
+from spiral.types_ import Timestamp
+
+if TYPE_CHECKING:
+    import duckdb
+    import polars as pl
+    import pyarrow.dataset
+
+    from spiral.tables import Tables
+    from spiral.tables.table import Table
+
+
+class Snapshot:
+    """Spiral table snapshot.
+
+    A snapshot represents a point-in-time view of a table.
+    """
+
+    def __init__(self, tables: "Tables", snapshot: TableSnapshot):
+        self._tables = tables
+        self._snapshot = snapshot
+
+    @property
+    def asof(self) -> Timestamp:
+        """Returns the asof timestamp of the snapshot."""
+        return self._snapshot.asof
+
+    @property
+    def client(self) -> "Tables":
+        """Returns the client used by the snapshot."""
+        return self._tables
+
+    @property
+    def table(self) -> "Table":
+        """Returns the table associated with the snapshot."""
+        from spiral.tables.table import Table
+
+        return Table(self._tables, self._snapshot.table)
+
+    def to_dataset(self) -> "pyarrow.dataset.Dataset":
+        """Returns a PyArrow Dataset representing the table."""
+        from .dataset import TableDataset
+
+        return TableDataset(self)
+
+    def to_polars(self) -> "pl.LazyFrame":
+        """Returns a Polars LazyFrame for the Spiral table."""
+        import polars as pl
+
+        return pl.scan_pyarrow_dataset(self.to_dataset())
+
+    def to_duckdb(self) -> "duckdb.DuckDBPyRelation":
+        """Returns a DuckDB relation for the Spiral table."""
+        import duckdb
+
+        return duckdb.from_arrow(self.to_dataset())
+
+    def scan(
+        self,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        exclude_keys: bool = False,
+    ) -> Scan:
+        """Reads the snapshot. If projections are not provided, the entire table is read."""
+        if not projections:
+            # Use table as the default projection.
+            projections = [self._snapshot.table.__expr__]
+
+        return self._tables.scan(
+            *projections,
+            where=where,
+            asof=self._snapshot.asof,
+            exclude_keys=exclude_keys,
+        )

spiral/tables/table.py

@@ -0,0 +1,157 @@
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from spiral.core.table import Table as CoreTable
+from spiral.core.table.spec import Schema
+from spiral.expressions.base import Expr, ExprLike
+from spiral.settings import settings
+from spiral.tables.maintenance import Maintenance
+from spiral.tables.scan import Scan
+from spiral.tables.snapshot import Snapshot
+from spiral.tables.transaction import Transaction
+
+if TYPE_CHECKING:
+    from spiral.tables import Tables
+
+
+class Table(Expr):
+    """API for interacting with a SpiralDB's Table.
+
+    Different catalog implementations should ultimately construct a Table object.
+    """
+
+    # TODO(marko): Make identifier required.
+    def __init__(self, tables: "Tables", table: CoreTable, *, identifier: str | None = None):
+        super().__init__(table.__expr__)
+
+        self._tables = tables
+        self._table = table
+        self._identifier = identifier
+        self._key_schema = self._table.key_schema
+        self._key_columns = set(self._key_schema.names)
+
+    @property
+    def client(self) -> "Tables":
+        """Returns the client used by the table."""
+        return self._tables
+
+    @property
+    def table_id(self) -> str:
+        return self._table.id
+
+    @property
+    def identifier(self) -> str:
+        """Returns the fully qualified identifier of the table."""
+        return self._identifier or self._table.id
+
+    @property
+    def dataset(self) -> str | None:
+        """Returns the dataset of the table."""
+        if self._identifier is None:
+            return None
+        _, dataset, _ = self._identifier.split(".")
+        return dataset
+
+    @property
+    def name(self) -> str | None:
+        """Returns the name of the table."""
+        if self._identifier is None:
+            return None
+        _, _, name = self._identifier.split(".")
+        return name
+
+    @property
+    def last_modified_at(self) -> int:
+        return self._table.get_wal(asof=None).last_modified_at
+
+    def __str__(self):
+        return self.identifier
+
+    def __repr__(self):
+        return f'Table("{self.identifier}")'
+
+    def __getitem__(self, item: str) -> Expr:
+        from spiral import expressions as se
+
+        if item in self._key_columns:
+            return se.key(name=item)
+
+        return super().__getitem__(item)
+
+    def select(self, *paths: str, exclude: list[str] = None) -> "Expr":
+        # Override an expression select in the root column group to split between keys and columns.
+        if exclude is not None:
+            if set(exclude) & self._key_columns:
+                raise ValueError(
+                    "Cannot use 'exclude' arg with key columns. Use 'exclude_keys' and an explicit select of keys."
+                )
+
+        key_paths = set(paths) & self._key_columns
+        other_paths = set(paths) - key_paths
+        if not key_paths:
+            return super().select(*paths, exclude=exclude)
+
+        from spiral import expressions as se
+
+        return se.merge(se.pack({key: se.key(key) for key in key_paths}), super().select(*other_paths, exclude=exclude))
+
+    @property
+    def key_schema(self) -> Schema:
+        """Returns the key schema of the table."""
+        return self._key_schema
+
+    @property
+    def schema(self) -> Schema:
+        """Returns the FULL schema of the table.
+
+        NOTE: This can be expensive for large tables.
+        """
+        return self._table.get_schema(asof=None)
+
+    def scan(
+        self,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        asof: datetime | int | None = None,
+        exclude_keys: bool = False,
+    ) -> Scan:
+        """Reads the table. If projections are not provided, the entire table is read."""
+        if not projections:
+            projections = [self]
+
+        return self._tables.scan(*projections, where=where, asof=asof, exclude_keys=exclude_keys)
+
+    def write(
+        self,
+        expr: ExprLike,
+        *,
+        partition_size_bytes: int | None = None,
+    ) -> None:
+        """Write an item to the table inside a single transaction.
+
+        :param expr: The expression to write. Must evaluate to a struct array.
+        :param partition_size_bytes: The maximum partition size in bytes.
+        """
+        with self.txn() as txn:
+            txn.write(
+                expr,
+                partition_size_bytes=partition_size_bytes,
+            )
+
+    def snapshot(self, asof: datetime | int | None = None) -> Snapshot:
+        """Returns a snapshot of the table at the given timestamp."""
+        if isinstance(asof, datetime):
+            asof = int(asof.timestamp() * 1_000_000)
+        return Snapshot(self._tables, self._table.get_snapshot(asof=asof))
+
+    def txn(self) -> Transaction:
+        """Begins a new transaction. Transaction must be committed for writes to become visible.
+
+        IMPORTANT: While transaction can be used to atomically write data to the table,
+             it is important that the primary key columns are unique within the transaction.
+        """
+        return Transaction(self._tables._spiral.open_transaction(self._table, settings().file_format))
+
+    def maintenance(self) -> Maintenance:
+        """Access maintenance operations for a table."""
+        return Maintenance(self._tables._spiral.open_maintenance(self._table, settings().file_format))