xarray_sql 0.3.0__cp310-abi3-win_amd64.whl

xarray_sql/reader.py

@@ -0,0 +1,332 @@
+"""Lazy Arrow stream reader for xarray Datasets.
+
+This module provides XarrayRecordBatchReader, which implements the Arrow
+PyCapsule Interface (__arrow_c_stream__) to enable zero-copy, lazy streaming
+of xarray data to DataFusion and other Arrow consumers.
+
+The implementation delegates to PyArrow's RecordBatchReader for the
+actual stream implementation, wrapping xarray block iteration in a generator.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pyarrow as pa
+import xarray as xr
+
+from .df import (
+    Block,
+    Chunks,
+    DEFAULT_BATCH_SIZE,
+    _block_metadata,
+    _block_slices_from_resolved,
+    _parse_schema,
+    block_slices,
+    iter_record_batches,
+    resolve_chunks,
+)
+
+if TYPE_CHECKING:
+    from ._native import LazyArrowStreamTable
+
+
+class XarrayRecordBatchReader:
+    """A lazy Arrow stream reader for xarray Datasets.
+
+    Implements the Arrow PyCapsule Interface (__arrow_c_stream__) to enable
+    zero-copy, lazy streaming of xarray data to DataFusion and other Arrow
+    consumers.
+
+    The key property is that xarray blocks are only converted to Arrow
+    RecordBatches when the consumer calls get_next (e.g., during DataFusion's
+    collect()), NOT when the reader is created or registered.
+
+    Attributes:
+        schema: The Arrow schema for the stream.
+
+    Example:
+        >>> import xarray as xr
+        >>> from xarray_sql import XarrayRecordBatchReader
+        >>> ds = xr.tutorial.open_dataset('air_temperature')
+        >>> reader = XarrayRecordBatchReader(ds, chunks={'time': 240})
+        >>> # At this point, NO data has been read from xarray
+        >>> # Data is only read when consumed:
+        >>> import pyarrow as pa
+        >>> pa_reader = pa.RecordBatchReader.from_stream(reader)
+        >>> for batch in pa_reader:
+        ...     print(batch.num_rows)  # Data read here
+    """
+
+    def __init__(
+        self,
+        ds: xr.Dataset,
+        chunks: Chunks = None,
+        *,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        _iteration_callback: (
+            Callable[[Block, list[str] | None], None] | None
+        ) = None,
+    ):
+        """Initialize the lazy reader.
+
+        Args:
+            ds: An xarray Dataset. All data_vars must share the same dimensions.
+            chunks: Xarray-like chunks specification. If not provided, uses
+                the Dataset's existing chunks.
+            batch_size: Maximum rows per emitted Arrow RecordBatch.  Smaller
+                values let DataFusion start processing earlier at the cost of
+                more Python→Arrow conversion calls.
+            _iteration_callback: Internal callback for testing. Called with
+                each block dict just before it's converted to Arrow. This
+                allows tests to track when iteration actually occurs.
+        """
+        self._ds = ds
+        self._chunks = chunks
+        self._batch_size = batch_size
+        self._schema = _parse_schema(ds)
+        self._iteration_callback = _iteration_callback
+        self._consumed = False
+
+        # Validate dimensions
+        fst = next(iter(ds.values())).dims
+        if not all(da.dims == fst for da in ds.values()):
+            raise ValueError(
+                "All dimensions must be equal. Please filter data_vars in the Dataset."
+            )
+
+    @property
+    def schema(self) -> pa.Schema:
+        """The Arrow schema for this stream."""
+        return self._schema
+
+    def _generate_batches(self) -> Iterator[pa.RecordBatch]:
+        """Generate RecordBatches lazily from xarray blocks.
+
+        This generator is only consumed when the Arrow stream's get_next
+        is called, ensuring true lazy evaluation.  Each xarray block is
+        emitted as one or more RecordBatches of at most self._batch_size rows.
+        """
+        for block in block_slices(self._ds, self._chunks):
+            # Call the iteration callback if provided (for testing).
+            # XarrayRecordBatchReader has no projection concept, so always passes None.
+            if self._iteration_callback is not None:
+                self._iteration_callback(block, None)
+
+            yield from iter_record_batches(
+                self._ds.isel(block), self._schema, self._batch_size
+            )
+
+    def __arrow_c_stream__(
+        self, requested_schema: object | None = None
+    ) -> object:
+        """Export as Arrow C Stream via PyCapsule.
+
+        This method is called by Arrow consumers (like DataFusion) to get
+        a C-level stream interface. The actual data iteration only begins
+        when the consumer calls get_next on the stream.
+
+        Args:
+            requested_schema: Optional schema for type casting. Currently
+                passed through to PyArrow's implementation.
+
+        Returns:
+            PyCapsule containing ArrowArrayStream pointer with name
+            "arrow_array_stream".
+
+        Raises:
+            RuntimeError: If the stream has already been consumed.
+        """
+        if self._consumed:
+            raise RuntimeError(
+                "Stream already consumed. XarrayRecordBatchReader can only "
+                "be iterated once. Create a new reader for additional iterations."
+            )
+        self._consumed = True
+
+        # Create a PyArrow RecordBatchReader from our generator
+        # The generator is NOT consumed here - only when get_next is called
+        reader = pa.RecordBatchReader.from_batches(
+            self._schema, self._generate_batches()
+        )
+
+        # Delegate to PyArrow's __arrow_c_stream__ implementation
+        return reader.__arrow_c_stream__(requested_schema)
+
+    def __arrow_c_schema__(
+        self, requested_schema: object | None = None
+    ) -> object:
+        """Export the schema as Arrow C Schema via PyCapsule.
+
+        This allows consumers to inspect the schema without consuming the stream.
+
+        Args:
+            requested_schema: Optional schema for negotiation (unused).
+
+        Returns:
+            PyCapsule containing ArrowSchema pointer.
+        """
+        return self._schema.__arrow_c_schema__()
+
+
+def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader:
+    """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks.
+
+    Args:
+      ds: An Xarray Dataset. All `data_vars` must share the same dimensions.
+      chunks: Xarray-like chunks. If not provided, will default to the
+        Dataset's chunks. The product of the chunk sizes becomes the
+        standard length of each dataframe partition.
+
+    Returns:
+      A PyArrow RecordBatchReader, which is a table representation of the input
+      Dataset.
+    """
+    reader = XarrayRecordBatchReader(ds, chunks=chunks)
+    return pa.RecordBatchReader.from_stream(reader)
+
+
+def read_xarray_table(
+    ds: xr.Dataset,
+    chunks: Chunks = None,
+    *,
+    batch_size: int = DEFAULT_BATCH_SIZE,
+    coord_arrays: dict[str, np.ndarray] | None = None,
+    _iteration_callback: (
+        Callable[[Block, list[str] | None], None] | None
+    ) = None,
+) -> "LazyArrowStreamTable":
+    """Create a lazy DataFusion table from an xarray Dataset.
+
+    This is the simplest way to register xarray data with DataFusion.
+    Data is only read when queries are executed, not during registration.
+    The table can be queried multiple times.
+
+    Each chunk becomes a separate partition, enabling DataFusion's parallel
+    execution across multiple cores.
+
+    Note:
+        SQL queries with WHERE clauses on dimension columns (time, lat, lon, etc.)
+        automatically prune partitions that can't contain matching rows — this is
+        called *filter pushdown*. For example:
+
+            # This query will skip loading partitions with time < '2020-02-01'
+            result = ctx.sql('SELECT * FROM air WHERE time > "2020-02-01"').collect()
+
+        Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`.
+
+    Args:
+        ds: An xarray Dataset. All data_vars must share the same dimensions.
+        chunks: Xarray-like chunks specification. If not provided, uses
+            the Dataset's existing chunks.
+        batch_size: Maximum rows per Arrow RecordBatch emitted per partition.
+            Smaller values let DataFusion start processing earlier; the default
+            (65 536) works well for most datasets.
+        coord_arrays: Pre-materialised coordinate arrays keyed by dim-name
+            string.  Hand in to share a single read across multiple tables
+            built from the same parent Dataset (e.g. surface + atmosphere
+            from ARCO-ERA5); the dim coords are otherwise read once per
+            ``read_xarray_table`` call, which is a network round-trip for
+            Zarr-backed datasets.
+        _iteration_callback: Internal callback for testing. Called with
+            each block dict just before it's converted to Arrow.
+
+    Returns:
+        A LazyArrowStreamTable ready for registration with DataFusion.
+
+    Example:
+        >>> from datafusion import SessionContext
+        >>> import xarray as xr
+        >>> from xarray_sql import read_xarray_table
+        >>>
+        >>> ds = xr.tutorial.open_dataset('air_temperature')
+        >>> table = read_xarray_table(ds, chunks={'time': 240})
+        >>>
+        >>> ctx = SessionContext()
+        >>> ctx.register_table('air', table)
+        >>>
+        >>> # Data is only read here, during query execution
+        >>> # Filters on 'time' will prune partitions automatically!
+        >>> result = ctx.sql('SELECT AVG(air) FROM air').collect()
+    """
+    from ._native import LazyArrowStreamTable
+
+    schema = _parse_schema(ds)
+
+    # Hoist coordinate reads once; avoids N_partitions remote I/O calls for
+    # Zarr-backed datasets (e.g. ARCO-ERA5 on GCS).  When the caller supplies
+    # pre-materialised arrays (e.g. shared across surface + atmosphere
+    # tables), reuse them and skip the extra read.
+    if coord_arrays is None:
+        coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims}
+
+    # Determine which column names are data variables (not dimension coordinates).
+    # Used by the factory to skip loading unrequested variables.
+    data_var_names = set(ds.data_vars.keys())
+
+    def make_partition_factory(
+        block: Block,
+    ) -> Callable[[list[str] | None], pa.RecordBatchReader]:
+        def make_stream(
+            projection_names: list[str] | None,
+        ) -> pa.RecordBatchReader:
+            if _iteration_callback is not None:
+                _iteration_callback(block, projection_names)
+
+            if projection_names is not None:
+                # Restrict to the data variables mentioned in the projection.
+                # Dimension coordinates come along automatically via coords.
+                data_vars_needed = [
+                    c for c in projection_names if c in data_var_names
+                ]
+                if data_vars_needed:
+                    ds_block = ds[data_vars_needed].isel(block)
+                else:
+                    # Only dimension coords requested — drop all data vars to avoid
+                    # loading them unnecessarily (e.g. for queries like SELECT lat, lon).
+                    ds_block = ds.drop_vars(list(ds.data_vars)).isel(block)
+                batch_schema = pa.schema(
+                    [schema.field(name) for name in projection_names]
+                )
+            else:
+                ds_block = ds.isel(block)
+                batch_schema = schema
+
+            return pa.RecordBatchReader.from_batches(
+                batch_schema,
+                iter_record_batches(ds_block, batch_schema, batch_size),
+            )
+
+        return make_stream
+
+    # Separate dims whose chunk bounds vary across partitions from those
+    # whose bounds are constant (one chunk spanning the whole axis).  For the
+    # latter we compute min/max once instead of re-scanning the full coord
+    # array on every partition — dominant cost when registering hundreds of
+    # thousands of single-time-step partitions on a 4-D dataset like ERA5.
+    resolved = resolve_chunks(ds, chunks)
+    varying_dims = [d for d, tup in resolved.items() if len(tup) > 1]
+    static_dims = [d for d in ds.dims if d not in varying_dims]
+    static_block: Block = {d: slice(None) for d in static_dims}
+    static_ranges = _block_metadata(
+        coord_arrays, static_block, dims=static_dims
+    )
+
+    def partition_pairs():
+        """Lazily yield (factory, metadata) for each partition.
+
+        Consuming this generator one item at a time means Python never holds
+        all N block dicts, metadata dicts, and factory closures simultaneously.
+        Peak Python memory during registration is O(1) per partition instead
+        of O(N_partitions).
+        """
+        for block in _block_slices_from_resolved(ds, resolved):
+            dynamic = _block_metadata(coord_arrays, block, dims=varying_dims)
+            yield (
+                make_partition_factory(block),
+                {**static_ranges, **dynamic},
+            )
+
+    return LazyArrowStreamTable(partition_pairs(), schema)

xarray_sql/sql.py

@@ -0,0 +1,191 @@
+import xarray as xr
+from datafusion import SessionContext
+from datafusion.catalog import Schema
+from collections import defaultdict
+
+from . import cftime as cft
+from .df import Chunks
+from .ds import XarrayDataFrame
+from .reader import read_xarray_table
+
+
+class XarrayContext(SessionContext):
+    """A datafusion `SessionContext` that also supports `xarray.Dataset`s."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Track registered xarray Datasets so XarrayDataFrame can recover
+        # defaults (dimension_columns) and metadata (var/dataset attrs,
+        # non-dim coords, dim-coord dtype) that the forward pivot drops.
+        # Keys are the fully-qualified table names users will reference
+        # in SQL (e.g. ``"air"`` for a uniform-dim Dataset, or
+        # ``"era5.surface"`` for one entry from a multi-dim-group split).
+        self._registered_datasets: dict[str, xr.Dataset] = {}
+
+    def from_dataset(
+        self,
+        name: str,
+        input_table: xr.Dataset,
+        *,
+        table_names: dict[tuple[str, ...], str] | None = None,
+        chunks: Chunks = None,
+    ):
+        """Register an xarray Dataset as one or more queryable SQL tables.
+
+        When all data variables share the same dimensions, the dataset is
+        registered as a single table named ``name``. When variables have
+        differing dimensions (e.g. some on a 3D grid and others on a 4D
+        grid), the dataset is split into one table per dimension group.
+        The tables are registered under a SQL schema (namespace) named
+        ``name`` and named ``<dim1>_<dim2>_...`` by default::
+
+            ctx.from_dataset('era5', ds, chunks={'time': 24})
+            # registers tables: 'era5.time_lat_lon' and
+            #                   'era5.time_lat_lon_level'
+            ctx.sql('SELECT AVG(temperature_2m) FROM era5.time_lat_lon')
+
+        Use ``table_names`` to override the name for specific dimension
+        tuples::
+
+            ctx.from_dataset(
+                'era5', ds,
+                table_names={('time', 'lat', 'lon'): 'surface'},
+            )
+            ctx.sql('SELECT * FROM era5.surface')
+
+        For datasets with non-Gregorian cftime coordinates (e.g. 360_day,
+        julian), a ``cftime()`` scalar UDF is automatically registered so
+        you can write ergonomic SQL filters::
+
+            ctx.from_dataset("ds360", ds, chunks={"time": 6})
+            ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')")
+
+        .. note::
+
+            Only one ``cftime()`` UDF is registered per context, using the
+            units and calendar of the *first* non-Gregorian coordinate
+            encountered. If you register multiple datasets with *different*
+            non-Gregorian calendars (e.g. one 360_day and one julian), the
+            UDF from the first registration will be used for all subsequent
+            ``cftime()`` calls and may produce incorrect offsets for the
+            other dataset. In that case, create a separate ``XarrayContext``
+            for each calendar.
+
+        Args:
+            name: The SQL identifier under which the dataset is registered.
+                For datasets with uniform dimensions, this is the table
+                name. For datasets with mixed dimensions, this is the name
+                of a SQL schema (namespace) containing one table per
+                dimension group.
+            input_table: An xarray Dataset.
+            table_names: Optional mapping from dimension tuples to custom
+                table names within the schema, used when the dataset has
+                variables with differing dimensions.
+            chunks: Xarray-like chunks specification. If not provided, uses
+                the Dataset's existing chunks.
+
+        Returns:
+            self, to allow chaining.
+        """
+        groups = _group_vars_by_dims(input_table)
+
+        # Materialise dim coordinates once and share across every sub-table.
+        # For Zarr-backed parents (e.g. ARCO-ERA5 on GCS) this saves one
+        # network round-trip per dim per dim-group.
+        coord_arrays = {
+            str(dim): input_table.coords[dim].values for dim in input_table.dims
+        }
+
+        if len(groups) <= 1:
+            self._registered_datasets[name] = input_table
+            return self._from_dataset(
+                name, input_table, chunks, coord_arrays=coord_arrays
+            )
+
+        table_names = table_names or {}
+        schema = Schema.memory_schema(self)
+        self.catalog().register_schema(name, schema)
+
+        for dims, var_names in groups.items():
+            # Scalar variables group under empty dims, where "_".join(()) is
+            # the empty string; fall back to a valid default table name.
+            sub_name = table_names.get(dims, "_".join(dims) or "scalar")
+            sub_ds = input_table[var_names]
+            self._from_dataset(
+                sub_name,
+                sub_ds,
+                chunks,
+                schema=schema,
+                coord_arrays=coord_arrays,
+            )
+            # Track the fully-qualified name so XarrayDataFrame metadata
+            # recovery can find this Dataset on round-trip.
+            self._registered_datasets[f"{name}.{sub_name}"] = sub_ds
+
+        return self
+
+    def _from_dataset(
+        self,
+        table_name: str,
+        input_table: xr.Dataset,
+        chunks: Chunks = None,
+        schema: Schema | None = None,
+        coord_arrays: dict | None = None,
+    ):
+        """Register a Dataset as a single SQL table.
+
+        Registers a top-level table by default, or a table inside ``schema``
+        (a SQL namespace) when one is given.
+        """
+        register = (
+            self.register_table if schema is None else schema.register_table
+        )
+        register(
+            table_name,
+            read_xarray_table(input_table, chunks, coord_arrays=coord_arrays),
+        )
+        self._maybe_register_cftime_udf(input_table)
+        return self
+
+    def _maybe_register_cftime_udf(self, ds: xr.Dataset) -> None:
+        """Auto-register a cftime() UDF for non-Gregorian cftime coordinates."""
+        for coord_name in ds.dims:
+            if cft.is_cftime_index(ds, coord_name):
+                units, cal = cft.encoding(ds, coord_name)
+                if not cft.is_gregorian_like(cal):
+                    self.register_udf(cft.make_cftime_udf(units, cal))
+                    break  # One UDF per context is enough.
+
+    def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
+        """Run a SQL query, returning an :class:`XarrayDataFrame` wrapper.
+
+        Identical to ``datafusion.SessionContext.sql`` except the returned
+        object wraps the DataFusion DataFrame. The wrapper exposes
+        ``.to_pandas()`` (unchanged), forwards every other DataFusion
+        method via ``__getattr__``, and adds
+        ``.to_dataset(dimension_columns=[...])`` for round-tripping the
+        result back to an ``xr.Dataset``.
+
+        Args:
+            query: A SQL query string.
+            *args: Forwarded to ``SessionContext.sql``.
+            **kwargs: Forwarded to ``SessionContext.sql``.
+
+        Returns:
+            An :class:`XarrayDataFrame` wrapping the DataFusion DataFrame.
+        """
+        inner = super().sql(query, *args, **kwargs)
+        return XarrayDataFrame(inner, templates=self._registered_datasets)
+
+
+def _group_vars_by_dims(ds: xr.Dataset) -> dict[tuple[str, ...], list[str]]:
+    """Group variables in the dataset based on shared dims.
+
+    ("time", "lat", "lon"):          ["temperature_2m", "wind_speed"],
+    ("time", "lat", "lon", "level"): ["pressure", "humidity"]
+    """
+    groups = defaultdict(list)
+    for var_name, var in ds.data_vars.items():
+        dims = var.dims
+        groups[dims].append(var_name)
+    return groups