[Copilot]

Skim (multi-dimension) data currently must be read from OMX or ZARR, or from a multi-indexed dataframe via xarray, which is slow. Parquet storage is substantially smaller and faster to read, so this adds a dedicated parquet reader for multi-dimension data.

from sharrow.dataset import from_parquet_3d

skims = from_parquet_3d(
    ["skims1.parquet", "skims2.parquet"],
    index_names=("otaz", "dtaz", "time_period"),
    time_periods=["EA", "AM", "MD", "PM", "EV"],
)

Input files have two index columns (otaz, dtaz by default) plus one column per matrix table, named exactly as in OMX including __ time period suffixes, which are assembled into the third dimension.

Changes

• sharrow/dataset.py
  • from_parquet_3d(...): accepts one path or an iterable of paths; mirrors the from_omx_3d signature (index_names, time_periods, time_period_sep, max_float_precision, ignore) and its handling of 2D vars, missing time period pages (zero-filled), and float downcasting.
  • _parquet_layout(...): inspects the index columns of each file and classifies the layout as row-major, column-major, unsorted-dense, or sparse. Dense detection is by first-occurrence unique labels; when neither dense ordering matches, uniqueness of index pairs is checked to distinguish a genuinely dense-but-misordered file from a coincidental row count.
  • Read pathways: dense files read column-by-column with pyarrow and reshape (transposing for column-major); sparse files (sorted or not) fall back to the existing xarray loader via set_index(...).to_xarray(); dense-but-unsorted raises ValueError rather than silently taking the slow path.
  • Multi-file: each file is classified independently, so layouts may differ between files. Files whose index labels are ordered differently from the first file are permuted onto the common coordinates; a mismatched label set raises.
• sharrow/tests/test_datasets.py: coverage for row-major, column-major, sparse with gaps and shuffling, dense-but-unsorted error, multiple files with mixed layouts, mismatched index ordering across files, and ignore filtering.
• docs/api.rst: lists the new constructor.

Notes for review

• Unlike from_omx_3d, the returned dataset is numpy-backed rather than dask-backed; the parquet read is eager.
• Sparse inputs inherit xarray's behavior of filling absent cells with NaN (promoting integer columns to float).
• Directories of parquet parts are not handled; each path is opened as a single ParquetFile, matching how the OMX reader takes discrete files.

• Fixes Read multi-index dense data from parquet #3

[jpn--]

Code Review (by AI)

Bugs / Correctness

Silent zero-fill for missing 3D time periods (dataset.py, from_parquet_3d):
When a time period declared in time_periods is absent from all input files, the slot is filled with np.zeros_like(prototype). For float data this should be np.full_like(prototype, np.nan) so callers can distinguish missing data from a true zero.

# current
arrs_ = [(i if i is not None else np.zeros_like(prototype)) for i in arrs]

# suggested
fill = np.full_like(prototype, np.nan) if np.issubdtype(prototype.dtype, np.floating) else np.zeros_like(prototype)
arrs_ = [(i if i is not None else fill) for i in arrs]

_parquet_column_to_numpy double-check falls through unsafely:
The second isinstance(column, pa.ChunkedArray) branch calls .to_numpy() without zero_copy_only=False, which will raise on non-contiguous arrays in older PyArrow. Should mirror the outer call:

if isinstance(column, pa.ChunkedArray):
    return column.to_numpy(zero_copy_only=False)  # not .to_numpy()

Empty parquet list returns a silent empty Dataset: no guard or error. Worth adding if not parquet: raise ValueError (or at least documenting the behaviour).

---

Performance

One pf.read() call per data column (_read_one_parquet_3d):

for k in data_names:
    content = _parquet_column_to_numpy(pf.read(columns=[k]), k)

For files with many columns (common for full skim tables) this causes N extra seeks/reads. A single pf.read(columns=data_names) followed by iterating over the returned table columns would be substantially faster, especially on cloud/remote storage.

---

Minor

• The prototype search loop can be simplified:

  prototype = next((i for i in arrs if i is not None), None)

• No test covers max_float_precision downcasting, or the zero-fill / NaN-fill behaviour for a declared but absent time period.

---

What's good

• _parquet_layout is clean and well-documented with clear detection of all four cases.
• The sparse fallback via pandas/xarray is correct and safe.
• The multi-file index-reordering logic (take_0 / take_1) is correct.
• Test coverage is broad: row-major, column-major, sparse, unsorted-dense, multi-file, mismatched index order, and ignore patterns are all exercised.

[jpn--]

superceded