pyspiral 0.7.18__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/debug/metrics.py

@@ -0,0 +1,56 @@
+from typing import Any
+
+
+def display_metrics(metrics: dict[str, Any]) -> None:
+    """Display metrics in a formatted table."""
+    print(
+        f"{'Metric':<40} {'Type':<10} {'Count':<8} {'Avg':<12} {'Min':<12} "
+        f"{'Max':<12} {'P95':<12} {'P99':<12} {'StdDev':<12}"
+    )
+    print("=" * 140)
+
+    for metric_name, data in sorted(metrics.items()):
+        metric_type = data["type"]
+        count = data["count"]
+        avg = _format_value(data["avg"], metric_type, metric_name)
+        min_val = _format_value(data["min"], metric_type, metric_name)
+        max_val = _format_value(data["max"], metric_type, metric_name)
+        p95 = _format_value(data["p95"], metric_type, metric_name)
+        p99 = _format_value(data["p99"], metric_type, metric_name)
+        stddev = _format_value(data["stddev"], metric_type, metric_name)
+
+        print(
+            f"{metric_name:<40} {metric_type:<10} {count:<8} {avg:<12} {min_val:<12} "
+            f"{max_val:<12} {p95:<12} {p99:<12} {stddev:<12}"
+        )
+
+
+def _format_duration(nanoseconds: float) -> str:
+    """Convert nanoseconds to human-readable duration."""
+    if nanoseconds >= 1_000_000_000:
+        return f"{nanoseconds / 1_000_000_000:.2f}s"
+    elif nanoseconds >= 1_000_000:
+        return f"{nanoseconds / 1_000_000:.2f}ms"
+    elif nanoseconds >= 1_000:
+        return f"{nanoseconds / 1_000:.2f}μs"
+    else:
+        return f"{nanoseconds:.0f}ns"
+
+
+def _format_bytes(bytes_value: float) -> str:
+    """Convert bytes to human-readable size."""
+    for unit in ["B", "KB", "MB", "GB"]:
+        if bytes_value < 1024:
+            return f"{bytes_value:.1f}{unit}"
+        bytes_value /= 1024
+    return f"{bytes_value:.1f}TB"
+
+
+def _format_value(value: float, metric_type: str, metric_name: str) -> str:
+    """Format a value based on metric type and name."""
+    if metric_type == "timer" or "duration" in metric_name:
+        return _format_duration(value)
+    elif "bytes" in metric_name:
+        return _format_bytes(value)
+    else:
+        return f"{value:,.0f}"

spiral/debug/scan.py

@@ -0,0 +1,266 @@
+from datetime import datetime
+
+from spiral.core.table import Scan
+from spiral.core.table.manifests import FragmentFile, FragmentManifest
+from spiral.core.table.spec import Key
+from spiral.types_ import Timestamp
+
+
+def show_scan(scan: Scan):
+    """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.splits()
+    key_space_state = scan.key_space_state(table_id)
+
+    # Collect all key bounds from all manifests. This makes sure all visualizations are aligned.
+    key_points = set()
+    key_space_manifest = key_space_state.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_state(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_state(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_human_size(size_bytes: int) -> str:
+    # Convert bytes to a human-readable format
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if size_bytes < 1024:
+            return f"{size_bytes:.2f} {unit}"
+        size_bytes /= 1024
+    return f"{size_bytes:.2f} PB"
+
+
+def _maybe_truncate(text, max_length: int = 30) -> str:
+    text = str(text)
+    if len(text) <= max_length:
+        return text
+
+    half_length = (max_length - 3) // 2
+    return text[:half_length] + "..." + text[-half_length:]
+
+
+def _get_fragment_legend(manifest_file: FragmentFile):
+    return "\n".join(
+        [
+            f"id: {manifest_file.id}",
+            f"size: {_get_human_size(manifest_file.size_bytes)} ({manifest_file.size_bytes} bytes)",
+            f"key_span: {manifest_file.key_span}",
+            f"key_min: {_maybe_truncate(manifest_file.key_extent.min)}",
+            f"key_max: {_maybe_truncate(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/enrichment.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+from functools import partial
+from typing import TYPE_CHECKING
+
+from spiral.core.client import KeyColumns, Shard
+from spiral.core.table import KeyRange
+from spiral.core.table.spec import Key, Operation
+from spiral.expressions import Expr
+
+if TYPE_CHECKING:
+    import dask.distributed
+
+    from spiral import Scan, Table
+
+logger = logging.getLogger(__name__)
+
+
+class Enrichment:
+    """
+    An enrichment is used to derive new columns from the existing once, such as fetching data from object storage
+    with `se.s3.get` or compute embeddings. With column groups design supporting 100s of thousands of columns,
+    horizontally expanding tables are a powerful primitive.
+
+    NOTE: Spiral aims to optimize enrichments where source and destination table are the same.
+    """
+
+    def __init__(
+        self,
+        table: Table,
+        projection: Expr,
+        where: Expr | None,
+    ):
+        self._table = table
+        self._projection = projection
+        self._where = where
+
+    @property
+    def table(self) -> Table:
+        """The table to write back into."""
+        return self._table
+
+    @property
+    def projection(self) -> Expr:
+        """The projection expression."""
+        return self._projection
+
+    @property
+    def where(self) -> Expr | None:
+        """The filter expression."""
+        return self._where
+
+    def _scan(self) -> Scan:
+        return self._table.spiral.scan(self._projection, where=self._where, _key_columns=KeyColumns.Included)
+
+    def apply(
+        self,
+        *,
+        batch_readahead: int | None = None,
+        partition_size_bytes: int | None = None,
+        txn_dump: str | None = None,
+    ) -> None:
+        """Apply the enrichment onto the table in a streaming fashion.
+
+        For large tables, consider using `apply_dask` for distributed execution.
+
+        Args:
+            index: Optional key space index to use for sharding the enrichment.
+                If not provided, the table's default sharding will be used.
+            partition_size_bytes: The maximum partition size in bytes.
+                If not provided, the default partition size is used.
+            txn_dump: Optional path to dump the transaction JSON for debugging.
+        """
+
+        txn = self._table.txn()
+
+        txn.writeback(
+            self._scan(),
+            partition_size_bytes=partition_size_bytes,
+            batch_readahead=batch_readahead,
+        )
+
+        if txn.is_empty():
+            logger.warning("Transaction not committed. No rows were read for enrichment.")
+            return
+
+        txn.commit(txn_dump=txn_dump)
+
+    def apply_dask(
+        self,
+        *,
+        partition_size_bytes: int | None = None,
+        max_task_size: int | None = None,
+        checkpoint_dump: str | None = None,
+        shards: list[Shard] | None = None,
+        txn_dump: str | None = None,
+        client: dask.distributed.Client | None = None,
+        **kwargs,
+    ) -> None:
+        """Use distributed Dask to apply the enrichment. Requires `dask[distributed]` to be installed.
+
+        If "address" of an existing Dask cluster is not provided in `kwargs`, a local cluster will be created.
+
+        Dask execution has some limitations, e.g. UDFs are not currently supported. These limitations
+        usually manifest as serialization errors when Dask workers attempt to serialize the state. If you are
+        encountering such issues, consider splitting the enrichment into UDF-only derivation that will be
+        executed in a streaming fashion, followed by a Dask enrichment for the rest of the computation.
+        If that is not possible, please reach out to the support for assistance.
+
+        How shards are determined:
+        - If `shards` is provided, those will be used directly.
+        - Else, if `checkpoint_dump` is provided, shards will be loaded from the checkpoint.
+        - Else, if `max_task_size` is provided, shards will be created based on the task size.
+        - Else, the scan's default sharding will be used.
+
+        Args:
+            partition_size_bytes: The maximum partition size in bytes.
+                If not provided, the default partition size is used.
+            max_task_size: Optional size task limit, in number of rows. Used for sharding.
+                If provided and checkpoint is present, the checkpoint shards will be used instead.
+                If not provided, the scan's default sharding will be used.
+            checkpoint_dump: Optional path to dump intermediate checkpoints for incremental progress.
+            shards: Optional list of shards to process.
+                If provided, `max_task_size` and `checkpoint_dump` are ignored.
+            txn_dump: Optional path to dump the transaction JSON for debugging.
+            client: Optional Dask distributed client. If not provided, a new client will be created
+            **kwargs: Additional keyword arguments to pass to `dask.distributed.Client`
+                such as `address` to connect to an existing cluster.
+        """
+        if client is None:
+            try:
+                from dask.distributed import Client
+            except ImportError:
+                raise ImportError("dask is not installed, please install dask[distributed] to use this feature.")
+
+            # Connect before doing any work.
+            client = Client(**kwargs)
+
+        # Start a transaction BEFORE the planning scan.
+        tx = self._table.txn()
+        plan_scan = self._scan()
+
+        # Determine the "tasks". Start from provided shards.
+        task_shards = shards
+        # If shards are not provided, try loading from checkpoint.
+        if task_shards is None and checkpoint_dump is not None:
+            checkpoint: list[KeyRange] | None = _checkpoint_load_key_ranges(checkpoint_dump)
+            if checkpoint is None:
+                logger.info(f"No existing checkpoint found at {checkpoint_dump}. Starting from scratch.")
+            else:
+                logger.info(f"Resuming enrichment from checkpoint at {checkpoint_dump} with {len(checkpoint)} ranges.")
+                task_shards = [Shard(kr, None) for kr in checkpoint]
+        # If still no shards, try creating from max task size.
+        if task_shards is None and max_task_size is not None:
+            task_shards = self._table.spiral.compute_shards(max_task_size, self.projection, self.where)
+        # Fallback to default sharding in the scan.
+        if task_shards is None:
+            task_shards = plan_scan.shards()
+
+        # Partially bind the enrichment function.
+        _compute = partial(
+            _enrichment_task,
+            settings_json=self._table.spiral.config.to_json(),
+            state_json=plan_scan.core.plan_state().to_json(),
+            output_table_id=self._table.table_id,
+            partition_size_bytes=partition_size_bytes,
+            incremental=checkpoint_dump is not None,
+        )
+        enrichments = client.map(_compute, task_shards)
+
+        logger.info(f"Applying enrichment with {len(task_shards)} shards. Follow progress at {client.dashboard_link}")
+
+        failed_ranges = []
+        try:
+            for result, shard in zip(client.gather(enrichments), task_shards):
+                result: EnrichmentTaskResult
+
+                if result.error is not None:
+                    logger.error(f"Enrichment task failed for range {shard.key_range}: {result.error}")
+                    failed_ranges.append(shard.key_range)
+                    continue
+
+                tx.include(result.ops)
+        except Exception as e:
+            # If not incremental, re-raise the exception.
+            if checkpoint_dump is None:
+                raise e
+
+            # Handle worker failures (e.g., KilledWorker from Dask)
+            from dask.distributed import KilledWorker
+
+            if isinstance(e, KilledWorker):
+                logger.error(f"Dask worker was killed during enrichment: {e}")
+
+            # Try to gather partial results and mark remaining tasks as failed
+            for future, shard in zip(enrichments, task_shards):
+                if future.done() and not future.exception():
+                    try:
+                        result = future.result()
+
+                        if result.error is not None:
+                            logger.error(f"Enrichment task failed for range {shard.key_range}: {result.error}")
+                            failed_ranges.append(shard.key_range)
+                            continue
+
+                        tx.include(result.ops)
+                    except Exception:
+                        # Task failed or incomplete, add to failed ranges
+                        failed_ranges.append(shard.key_range)
+                else:
+                    # Task didn't complete, add to failed ranges
+                    failed_ranges.append(shard.key_range)
+
+        # Dump checkpoint of failed ranges, if any.
+        if checkpoint_dump is not None:
+            logger.info(
+                f"Dumping checkpoint with failed {len(failed_ranges)}/{len(task_shards)} ranges to {checkpoint_dump}."
+            )
+            _checkpoint_dump_key_ranges(checkpoint_dump, failed_ranges)
+
+        if tx.is_empty():
+            logger.warning("Transaction not committed. No rows were read for enrichment.")
+            return
+
+        # Always compact in distributed enrichment.
+        tx.commit(compact=True, txn_dump=txn_dump)
+
+
+def _checkpoint_load_key_ranges(checkpoint_dump: str) -> list[KeyRange] | None:
+    import json
+    import os
+
+    if not os.path.exists(checkpoint_dump):
+        return None
+
+    with open(checkpoint_dump) as f:
+        data = json.load(f)
+        return [
+            KeyRange(begin=Key(bytes.fromhex(r["begin"])), end=Key(bytes.fromhex(r["end"])))
+            for r in data.get("key_ranges", [])
+        ]
+
+
+def _checkpoint_dump_key_ranges(checkpoint_dump: str, ranges: list[KeyRange]):
+    import json
+    import os
+
+    os.makedirs(os.path.dirname(checkpoint_dump), exist_ok=True)
+    with open(checkpoint_dump, "w") as f:
+        json.dump(
+            {"key_ranges": [{"begin": bytes(r.begin).hex(), "end": bytes(r.end).hex()} for r in ranges]},
+            f,
+        )
+
+
+@dataclasses.dataclass
+class EnrichmentTaskResult:
+    ops: list[Operation]
+    error: str | None = None
+
+    def __getstate__(self):
+        return {
+            "ops": [op.to_json() for op in self.ops],
+            "error": self.error,
+        }
+
+    def __setstate__(self, state):
+        self.ops = [Operation.from_json(op_json) for op_json in state["ops"]]
+        self.error = state["error"]
+
+
+# NOTE(marko): This function must be picklable!
+def _enrichment_task(
+    shard: Shard,
+    *,
+    settings_json: str,
+    state_json: str,
+    output_table_id,
+    partition_size_bytes: int | None,
+    incremental: bool,
+) -> EnrichmentTaskResult:
+    # Returns operations that can be included in a transaction.
+    from spiral import Scan, Spiral
+    from spiral.core.table import ScanState
+    from spiral.settings import ClientSettings
+
+    settings = ClientSettings.from_json(settings_json)
+    sp = Spiral(config=settings)
+    state = ScanState.from_json(state_json)
+    task_scan = Scan(sp, sp.core.load_scan(state))
+    table = sp.table(output_table_id)
+    task_tx = table.txn()
+
+    try:
+        task_tx.writeback(task_scan, key_range=shard.key_range, partition_size_bytes=partition_size_bytes)
+        return EnrichmentTaskResult(ops=task_tx.take())
+    except Exception as e:
+        task_tx.abort()
+
+        if incremental:
+            return EnrichmentTaskResult(ops=[], error=str(e))
+
+        logger.error(f"Enrichment task failed for shard {shard}: {e}")
+        raise e