pyspiral 0.1.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/dataset.py

@@ -0,0 +1,239 @@
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+if TYPE_CHECKING:
+    import pyarrow.dataset
+
+from spiral import Scan, Table
+
+
+class TableDataset(pa.dataset.Dataset):
+    def __init__(self, table: Table):
+        self._table = table
+        self._schema: pa.Schema = table.scan().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__()
+
+    @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: pa.dataset.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) -> "TableDataset":
+        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: pa.dataset.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ):
+        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) -> "TableDataset":
+        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: pa.dataset.FragmentScanOptions | None = None,
+        use_threads: bool = True,
+        memory_pool: pa.MemoryPool = None,
+    ) -> "TableScanner":
+        from .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).convert(
+                filter.to_substrait(self._schema, allow_arrow_extensions=True),
+            )
+
+        scan = self._table.scan(
+            {c: self._table[c] for c in columns} if columns else self._table,
+            where=filter,
+            exclude_keys=True,
+        )
+        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: pa.dataset.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: pa.dataset.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: pa.dataset.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(pa.dataset.Scanner):
+    """A PyArrow Dataset Scanner that reads from a Spiral Table."""
+
+    def __init__(self, scan: Scan):
+        self._scan = scan
+        self._schema = scan.schema
+
+        # 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."""
+        reader = self.to_reader()
+        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_record_batches()
+
+    def to_table(self):
+        return self.to_reader().read_all()

spiral/debug.py

@@ -0,0 +1,251 @@
+from datetime import datetime
+
+from spiral.core.core import TableScan
+from spiral.core.manifests import FragmentFile, FragmentManifest
+from spiral.core.spec import Key, KeyRange
+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[KeyRange] = 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.fs_level}",
+            f"committed_at: {_format_timestamp(manifest_file.committed_at)}",
+            f"compacted_at: {_format_timestamp(manifest_file.compacted_at)}",
+            f"fs_id: {manifest_file.fs_id}",
+            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/expressions/__init__.py

@@ -0,0 +1,222 @@
+import builtins
+import functools
+import operator
+from typing import Any
+
+import pyarrow as pa
+
+from spiral import _lib, arrow
+
+from . import http as http
+from . import io as io
+from . import list_ as list
+from . import refs as refs
+from . import str_ as str
+from . import struct as struct
+from . import tiff as tiff
+from .base import Expr, ExprLike
+
+__all__ = [
+    "Expr",
+    "add",
+    "and_",
+    "deref",
+    "divide",
+    "eq",
+    "getitem",
+    "gt",
+    "gte",
+    "http",
+    "io",
+    "is_not_null",
+    "is_null",
+    "lift",
+    "list",
+    "lt",
+    "lte",
+    "merge",
+    "modulo",
+    "multiply",
+    "negate",
+    "neq",
+    "not_",
+    "or_",
+    "pack",
+    "ref",
+    "refs",
+    "scalar",
+    "select",
+    "str",
+    "struct",
+    "subtract",
+    "tiff",
+    "var",
+    "xor",
+]
+
+# Inline some of the struct expressions since they're so common
+getitem = struct.getitem
+merge = struct.merge
+pack = struct.pack
+select = struct.select
+ref = refs.ref
+deref = refs.deref
+
+
+def lift(expr: ExprLike) -> Expr:
+    # Convert an ExprLike into an Expr.
+    if isinstance(expr, Expr):
+        return expr
+
+    if isinstance(expr, dict):
+        # NOTE: we assume this is a struct expression. We could be smarter and be context aware to determine if
+        # this is in fact a struct scalar, but the user can always create one of those manually.
+
+        # First we un-nest any dot-separated field names
+        expr: dict = arrow.nest_structs(expr)
+
+        return pack({k: lift(v) for k, v in expr.items()})
+
+    if isinstance(expr, builtins.list):
+        return lift(pa.array(expr))
+
+    # Unpack tables and chunked arrays
+    if isinstance(expr, pa.Table):
+        expr = expr.to_struct_array()
+    if isinstance(expr, pa.ChunkedArray):
+        expr = expr.combine_chunks()
+
+    # If the value is struct-like, we un-nest any dot-separated field names
+    if isinstance(expr, pa.StructArray | pa.StructScalar):
+        return lift(arrow.nest_structs(expr))
+
+    if isinstance(expr, pa.Array):
+        return Expr(_lib.spql.expr.array_lit(expr))
+
+    # Otherwise, assume it's a scalar.
+    return scalar(expr)
+
+
+def var(name: builtins.str) -> Expr:
+    """Create a variable expression."""
+    return Expr(_lib.spql.expr.var(name))
+
+
+def keyed(name: builtins.str, dtype: pa.DataType) -> Expr:
+    """Create a variable expression referencing a column in the key table.
+
+    Key table is optionally given to `Scan#to_record_batches` function when reading only specific keys
+    or doing cell pushdown.
+
+    Args:
+        name: variable name
+        dtype: must match dtype of the column in the key table.
+    """
+    return Expr(_lib.spql.expr.keyed(f"#{name}", dtype))
+
+
+def scalar(value: Any) -> Expr:
+    """Create a scalar expression."""
+    if not isinstance(value, pa.Scalar):
+        value = pa.scalar(value)
+    return Expr(_lib.spql.expr.scalar(value))
+
+
+def cast(expr: ExprLike, dtype: pa.DataType) -> Expr:
+    """Cast an expression into another PyArrow DataType."""
+    expr = lift(expr)
+    return Expr(_lib.spql.expr.cast(expr.__expr__, dtype))
+
+
+def and_(expr: ExprLike, *exprs: ExprLike) -> Expr:
+    """Create a conjunction of one or more expressions."""
+
+    return functools.reduce(operator.and_, [lift(e) for e in exprs], lift(expr))
+
+
+def or_(expr: ExprLike, *exprs: ExprLike) -> Expr:
+    """Create a disjunction of one or more expressions."""
+    return functools.reduce(operator.or_, [lift(e) for e in exprs], lift(expr))
+
+
+def eq(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create an equality comparison."""
+    return operator.eq(lift(lhs), rhs)
+
+
+def neq(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a not-equal comparison."""
+    return operator.ne(lift(lhs), rhs)
+
+
+def xor(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a XOR comparison."""
+    return operator.xor(lift(lhs), rhs)
+
+
+def lt(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a less-than comparison."""
+    return operator.lt(lift(lhs), rhs)
+
+
+def lte(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a less-than-or-equal comparison."""
+    return operator.le(lift(lhs), rhs)
+
+
+def gt(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a greater-than comparison."""
+    return operator.gt(lift(lhs), rhs)
+
+
+def gte(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Create a greater-than-or-equal comparison."""
+    return operator.ge(lift(lhs), rhs)
+
+
+def negate(expr: ExprLike) -> Expr:
+    """Negate the given expression."""
+    return operator.neg(lift(expr))
+
+
+def not_(expr: ExprLike) -> Expr:
+    """Negate the given expression."""
+    expr = lift(expr)
+    return Expr(_lib.spql.expr.unary("not", expr.__expr__))
+
+
+def is_null(expr: ExprLike) -> Expr:
+    """Check if the given expression is null."""
+    expr = lift(expr)
+    return Expr(_lib.spql.expr.unary("is_null", expr.__expr__))
+
+
+def is_not_null(expr: ExprLike) -> Expr:
+    """Check if the given expression is null."""
+    expr = lift(expr)
+    return Expr(_lib.spql.expr.unary("is_not_null", expr.__expr__))
+
+
+def add(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Add two expressions."""
+    return operator.add(lift(lhs), rhs)
+
+
+def subtract(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Subtract two expressions."""
+    return operator.sub(lift(lhs), rhs)
+
+
+def multiply(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Multiply two expressions."""
+    return operator.mul(lift(lhs), rhs)
+
+
+def divide(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Divide two expressions."""
+    return operator.truediv(lift(lhs), rhs)
+
+
+def modulo(lhs: ExprLike, rhs: ExprLike) -> Expr:
+    """Modulo two expressions."""
+    return operator.mod(lift(lhs), rhs)