import itertools

from collections.abc import Callable, Hashable, Iterator, Mapping

from typing import Any

import numpy as np

import pandas as pd

import pyarrow as pa

import xarray as xr

from . import cftime as cft

Block = dict[Hashable, slice]

Chunks = dict[str, int] | None

# Borrowed from Xarray

def _get_chunk_slicer(

dim: Hashable, chunk_index: Mapping, chunk_bounds: Mapping

):

if dim in chunk_index:

which_chunk = chunk_index[dim]

return slice(

chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1]

)

return slice(None)

# Adapted from Xarray `map_blocks` implementation.

def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]:

"""Compute block slices for a chunked Dataset."""

if chunks is not None:

for_chunking = ds.copy(data=None, deep=False).chunk(chunks)

chunks = for_chunking.chunks

del for_chunking

else:

chunks = ds.chunks

assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided."

# chunks is Dict[str, Tuple[int, ...]] from xarray

chunk_bounds = {

dim: np.cumsum((0,) + tuple(c)) # type: ignore[arg-type]

for dim, c in chunks.items()

}

ichunk = {dim: range(len(tuple(c))) for dim, c in chunks.items()} # type: ignore[arg-type]

ick, icv = zip(*ichunk.items()) # Makes same order of keys and val.

chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv))

blocks = (

{

dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds)

for dim in ds.dims

}

for chunk_index in chunk_idxs

)

yield from blocks

def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]:

"""Explodes a dataset into its chunks."""

yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks))

def _block_len(block: Block) -> int:

return int(np.prod([v.stop - v.start for v in block.values()]))

def from_map_batched(

func: Callable[..., pd.DataFrame],

*iterables: tuple[Any, ...],

args: tuple | None = None,

schema: pa.Schema = None,

**kwargs: dict[str, Any],

) -> pa.RecordBatchReader:

"""Create a PyArrow RecordBatchReader by mapping a function over iterables.

This is equivalent to dask's from_map but returns a PyArrow

RecordBatchReader that can be used with DataFusion. It iterates over

RecordBatches which are created via the `func` one-at-a-time.

Args:

func: Function to apply to each element of the iterables. Currently, the

function must return a Pandas DataFrame.

*iterables: Iterable objects to map the function over.

schema: Optional schema needed for the RecordBatchReader.

args: Additional positional arguments to pass to func.

**kwargs: Additional keyword arguments to pass to func.

Returns:

A PyArrow RecordBatchReader containing the stream of RecordBatches.

"""

if args is None:

args = ()

def map_batches():

for items in zip(*iterables):

df = func(*items, *args, **kwargs)

yield pa.RecordBatch.from_pandas(df, schema=schema)

return pa.RecordBatchReader.from_batches(schema, map_batches())

def from_map(

func: Callable,

*iterables: tuple[Any, ...],

args: tuple | None = None,

**kwargs: dict[str, Any],

) -> pa.Table:

"""Create a PyArrow Table by mapping a function over iterables.

This is equivalent to dask's from_map but returns a PyArrow Table

that can be used with DataFusion instead of a Dask DataFrame.

Args:

func: Function to apply to each element of the iterables.

*iterables: Iterable objects to map the function over.

args: Additional positional arguments to pass to func.

**kwargs: Additional keyword arguments to pass to func.

Returns:

A PyArrow Table containing the concatenated results.

"""

if args is None:

args = ()

# Apply the function to each combination of iterable elements

results = []

for items in zip(*iterables) if len(iterables) > 1 else iterables[0]:

if isinstance(items, tuple):

result = func(*items, *args, **kwargs)

else:

result = func(items, *args, **kwargs)

# Convert result to PyArrow Table

if isinstance(result, pd.DataFrame):

pa_table = pa.Table.from_pandas(result)

elif isinstance(result, pa.Table):

pa_table = result

else:

# Try to convert to pandas first, then to PyArrow

try:

df = pd.DataFrame(result)

pa_table = pa.Table.from_pandas(df)

except Exception as e:

raise ValueError(

f"Cannot convert function result to PyArrow Table: {e}"

)

results.append(pa_table)

# Concatenate all results

if not results:

raise ValueError("No results to concatenate")

return pa.concat_tables(results)

def pivot(ds: xr.Dataset) -> pd.DataFrame:

"""Converts an xarray Dataset to a pandas DataFrame."""

return ds.to_dataframe().reset_index() # type: ignore[no-any-return]

def dataset_to_record_batch(

ds: xr.Dataset, schema: pa.Schema

) -> pa.RecordBatch:

"""Convert an xarray Dataset partition to an Arrow RecordBatch.

Builds the RecordBatch directly from numpy arrays, bypassing the pandas

round-trip (to_dataframe → reset_index → from_pandas) used by pivot().

For large partitions this reduces peak memory from ~5× to ~2× the

partition size.

Dimension coordinates are broadcast to the full partition shape and

ravelled. np.broadcast_to() is zero-copy; the ravel() forces one copy

per coordinate (unavoidable, since broadcast arrays are non-contiguous).

Data variable arrays are ravelled in-place — a zero-copy view when the

underlying array is already C-contiguous (the common case for numpy-backed

xarray datasets).

Args:

ds: A partition-sized xarray Dataset (already sliced via isel).

schema: The Arrow schema for the output, as produced by _parse_schema.

Column order in the output matches schema field order.

Returns:

A RecordBatch with one column per dimension coordinate and data

variable, in schema order.

"""

# Use the data variable's dimension order as canonical so coordinate

# broadcasts and data variable ravels use the same layout. All data

# variables are validated to share the same dims tuple.

if ds.data_vars:

first_var = next(iter(ds.data_vars.values()))

dim_names = list(first_var.dims)

shape = first_var.shape

else:

dim_names = list(ds.sizes.keys())

shape = tuple(ds.sizes[d] for d in dim_names)

arrays = []

for field in schema:

name = field.name

if name in ds.coords and name in ds.dims:

# Broadcast 1-D coordinate to the full N-D partition shape, then ravel.

axis = dim_names.index(name)

coord = ds.coords[name].values

if cft.is_cftime(coord):

coord = cft.convert_for_field(coord, field)

reshape = [1] * len(shape)

reshape[axis] = coord.shape[0]

arr = np.broadcast_to(coord.reshape(reshape), shape).ravel()

arrays.append(pa.array(arr, type=field.type))

else:

# Data variable: ravel to 1-D (zero-copy for C-contiguous arrays).

raw = ds[name].values.ravel()

if cft.is_cftime(ds[name].values):

raw = cft.convert_for_field(ds[name].values, field)

# from_pandas=True maps NaN → Arrow null inside the C++ copy kernel,

# so SQL aggregates (MAX, MIN, AVG) skip missing values correctly.

arrays.append(pa.array(raw, type=field.type, from_pandas=True))

return pa.RecordBatch.from_arrays(arrays, schema=schema)

#: Default number of rows per emitted Arrow RecordBatch.

#: 64 K rows balances DataFusion pipeline depth against per-batch overhead.

DEFAULT_BATCH_SIZE: int = 65_536

def iter_record_batches(

ds: xr.Dataset,

schema: pa.Schema,

batch_size: int = DEFAULT_BATCH_SIZE,

) -> Iterator[pa.RecordBatch]:

"""Yield RecordBatches of at most *batch_size* rows from a partition Dataset.

Unlike `dataset_to_record_batch`, which materialises the entire

partition as one batch, this generator emits smaller batches so that

DataFusion can begin filtering and aggregating before the full partition

is loaded. Peak memory per batch is O(batch_size) for coordinate columns

and O(partition_size) for data-variable columns (which must be loaded in

full from storage).

Coordinate values are computed per batch via strided index arithmetic —

no broadcast array spanning the whole partition is ever allocated. Data

variable flat arrays are loaded once (triggering any remote I/O) and then

sliced as zero-copy views for each batch.

Args:

ds: A partition-sized xarray Dataset (already sliced via isel).

schema: The Arrow schema for the output, as produced by _parse_schema.

batch_size: Maximum number of rows per yielded RecordBatch.

Yields:

RecordBatches in schema column order, covering all rows of the

partition exactly once.

"""

if ds.data_vars:

first_var = next(iter(ds.data_vars.values()))

dim_names = list(first_var.dims)

shape = first_var.shape

else:

dim_names = list(ds.sizes.keys())

shape = tuple(ds.sizes[d] for d in dim_names)

total_rows = int(np.prod(shape))

# Preload small 1-D coordinate arrays (negligible memory).

# Convert cftime objects to numeric values matching the schema type.

coord_values = {}

for name in dim_names:

vals = ds.coords[name].values

if cft.is_cftime(vals):

coord_values[name] = cft.convert_for_field(vals, schema.field(name))

else:

coord_values[name] = vals

# C-order stride for each dimension: stride[k] = prod(shape[k+1:]).

# Flat row index i → coordinate index for dim k: (i // stride[k]) % shape[k].

strides = [int(np.prod(shape[k + 1 :])) for k in range(len(shape))]

# Load data-variable arrays fully (triggers Dask/Zarr compute once).

# ravel() is a zero-copy view for C-contiguous arrays.

data_arrays = {}

for field in schema:

if field.name not in ds.dims:

raw = ds[field.name].values

if cft.is_cftime(raw):

data_arrays[field.name] = cft.convert_for_field(raw, field)

else:

data_arrays[field.name] = raw.ravel()

for row_start in range(0, total_rows, batch_size):

row_end = min(row_start + batch_size, total_rows)

row_idx = np.arange(row_start, row_end)

arrays = []

for field in schema:

name = field.name

if name in ds.coords and name in ds.dims:

k = dim_names.index(name)

coord_idx = (row_idx // strides[k]) % shape[k]

arrays.append(

pa.array(coord_values[name][coord_idx], type=field.type)

)

else:

arrays.append(

pa.array(

data_arrays[name][row_start:row_end],

type=field.type,

from_pandas=True,

)

)

yield pa.RecordBatch.from_arrays(arrays, schema=schema)

def _parse_schema(ds: xr.Dataset) -> pa.Schema:

"""Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.

Uses the xarray index type to detect cftime coordinates without

materializing their data — important for Dask/Zarr-backed datasets

where .values would trigger eager computation.

cftime coordinates are mapped to one of two Arrow types:

* **Gregorian-like calendars** (standard, noleap, all_leap, etc.):

``pa.timestamp('us')`` so string-based SQL filters work naturally.

* **Non-Gregorian calendars** (360_day, julian):

``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata

on the field, preserving lossless CF-convention encoding.

"""

columns = []

for coord_name, coord_var in ds.coords.items():

# Only include dimension coordinates

if coord_name in ds.dims:

if cft.is_cftime_index(ds, coord_name):

units, calendar = cft.encoding(ds, coord_name)

columns.append(cft.arrow_field(coord_name, units, calendar))

else:

pa_type = pa.from_numpy_dtype(coord_var.dtype)

columns.append(pa.field(coord_name, pa_type))

for var_name, var in ds.data_vars.items():

# Data variables are virtually never cftime, but check dtype as a

# cheap guard. Only fall back to _is_cftime (which materializes

# element 0) when dtype is object.

if var.dtype == np.dtype("O") and cft.is_cftime(var.values):

# Rare: a data variable holding cftime objects. Use same encoding

# as the first cftime dimension coordinate, or default.

cal = var.values.ravel()[0].calendar

columns.append(cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal))

else:

pa_type = pa.from_numpy_dtype(var.dtype)

columns.append(pa.field(var_name, pa_type))

return pa.schema(columns)

# Type alias for partition metadata: maps dimension name to (min, max, dtype_str) values

PartitionBounds = dict[str, tuple[Any, Any, str]]

def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds:

"""Compute min/max coordinate values for a single partition block.

Args:

coord_arrays: Pre-materialised coordinate arrays keyed by dimension name

string. Hoist this outside any loop to avoid repeated remote I/O

for Zarr-backed datasets.

block: A single block slice dict from block_slices().

Returns:

Dict mapping dimension name to (min_value, max_value, dtype_str).

Dimensions with an empty slice are omitted; the Rust pruning logic

treats missing dimensions conservatively (never prunes on them).

"""

ranges: PartitionBounds = {}

for dim, slc in block.items():

coord_values = coord_arrays[str(dim)][slc]

if len(coord_values) == 0:

continue

# String/object dtypes are not representable as ScalarBound

# (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not

# support them. Skip so pruning treats the dimension conservatively.

if coord_values.dtype.kind in ("U", "S", "O"):

continue

if cft.is_cftime(coord_values):

ranges[str(dim)] = cft.partition_bounds(coord_values)

continue

# Use actual min/max rather than first/last so that non-monotonic

# coordinate axes (e.g. descending latitude 90→-90) are handled

# correctly. np.min/max work for both numeric and datetime64 arrays.

min_val = coord_values.min()

max_val = coord_values.max()

if isinstance(min_val, (np.datetime64, pd.Timestamp)):

min_val = int(pd.Timestamp(min_val).value)

max_val = int(pd.Timestamp(max_val).value)

ranges[str(dim)] = (min_val, max_val, "timestamp_ns")

elif hasattr(min_val, "item"):

min_val = min_val.item()

max_val = max_val.item()

dtype = "float64" if isinstance(min_val, float) else "int64"

ranges[str(dim)] = (min_val, max_val, dtype)

else:

dtype = "float64" if isinstance(min_val, float) else "int64"

ranges[str(dim)] = (min_val, max_val, dtype)

return ranges

def partition_metadata(

ds: xr.Dataset, blocks: list[Block]

) -> list[PartitionBounds]:

"""Compute min/max coordinate values for each partition.

This metadata enables filter pushdown: SQL queries with WHERE clauses

on dimension columns can prune partitions that can't contain matching rows.

Args:

ds: The xarray Dataset containing coordinate values.

blocks: List of block slices from block_slices().

Returns:

List of dicts mapping dimension name to

(min_value, max_value, dtype_str) tuples.

- For datetime64, values are nanoseconds since Unix epoch

(int64), dtype_str is "timestamp_ns"

- For numeric types, values are Python int or float,

dtype_str is "int64" or "float64"

Note:

If a partition has an empty slice for a dimension, that dimension is

omitted from the partition's metadata. The Rust pruning logic treats

missing dimensions conservatively (never prunes on them).

"""

# Hoist coordinate array reads outside the partition loop.

# ds.coords[dim].values materializes the full array on every call; doing it

# N_partitions × N_dims times is wasteful and, for remote Zarr-backed datasets

# (e.g. ARCO-ERA5 on GCS), may trigger repeated network I/O.

coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims}

return [_block_metadata(coord_arrays, block) for block in blocks]