From f0e38902f92f343b3c5fe26e6ef89270034e338c Mon Sep 17 00:00:00 2001 From: Shane Date: Fri, 12 Jun 2026 16:11:28 -0700 Subject: [PATCH 1/8] phase 1 of issue #30: sort/hash grouping refactor --- src/zagg/processing.py | 94 ++++++++++++++++++++++++++++++---------- tests/test_processing.py | 87 +++++++++++++++++++++++++++++++++---- 2 files changed, 151 insertions(+), 30 deletions(-) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 8ad31e3b..da64f987 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -34,6 +34,42 @@ def _make_url_rewriter(driver: str | None): return lambda url: url.replace("s3://", "", 1) +def _build_groups( + df_all: pd.DataFrame, + cell_col: np.ndarray, +) -> tuple[dict[str, np.ndarray], dict[int, tuple[int, int]]]: + """Sort observations by cell id; return reordered column arrays and per-cell slice map. + + O(n log n) replacement for the O(n_children x n_obs) boolean-mask loop. The + returned arrays are sorted by ascending cell id; each cell's observations form + a contiguous slice, so ``col_arrays[col][start:end]`` is a zero-copy view. + + Parameters + ---------- + df_all : pd.DataFrame + Combined observation DataFrame (all beams / granules for this shard). + cell_col : np.ndarray + Cell id for each row in df_all (from ``grid.cells_of``). + + Returns + ------- + col_arrays : dict[str, np.ndarray] + Column arrays from df_all, sorted in ascending cell-id order. + cell_to_slice : dict[int, tuple[int, int]] + Maps each observed cell id to ``(start, end)`` indices into col_arrays. + """ + sort_idx = np.argsort(cell_col, kind="stable") + sorted_cells = cell_col[sort_idx] + col_arrays = {col: df_all[col].values[sort_idx] for col in df_all.columns} + boundaries = np.flatnonzero(np.diff(sorted_cells)) + 1 + starts = np.concatenate([[0], boundaries]) + ends = np.concatenate([boundaries, [len(sorted_cells)]]) + cell_to_slice = { + int(sorted_cells[s]): (int(s), int(e)) for s, e in zip(starts, ends) + } + return col_arrays, cell_to_slice + + def write_dataframe_to_zarr( df_out: pd.DataFrame, store: Store, @@ -89,9 +125,9 @@ def write_dataframe_to_zarr( def calculate_cell_statistics( - df_cell: pd.DataFrame, - value_col="h_li", - sigma_col="s_li", + cell_data: dict[str, np.ndarray], + value_col: str = "h_li", + sigma_col: str = "s_li", config: PipelineConfig | None = None, ) -> dict: """ @@ -99,19 +135,20 @@ def calculate_cell_statistics( Parameters ---------- - df_cell : pd.DataFrame - Dataframe containing observations for a single cell + cell_data : dict[str, np.ndarray] + Column arrays for a single cell. Keys are column names; values are + numpy arrays of equal length. value_col : str - Column name for elevation values + Column name for elevation values. sigma_col : str - Column name for uncertainty values + Column name for uncertainty values. config : PipelineConfig, optional Pipeline config to use for dispatch. Defaults to ``default_config()``. Returns ------- dict - Dictionary of statistics keyed by aggregation variable name + Dictionary of statistics keyed by aggregation variable name. """ from zagg.config import evaluate_expression, resolve_function @@ -119,7 +156,8 @@ def calculate_cell_statistics( config = default_config() agg_fields = get_agg_fields(config) - if len(df_cell) == 0: + n_obs = len(next(iter(cell_data.values()))) if cell_data else 0 + if n_obs == 0: return { name: (0 if meta.get("function") in ("len", "count") else np.nan) for name, meta in agg_fields.items() @@ -134,25 +172,26 @@ def calculate_cell_statistics( # Expression-based aggregation (e.g. h_sigma) if expression: - columns = {col: df_cell[col].values for col in df_cell.columns} - result[name] = evaluate_expression(expression, columns) + result[name] = evaluate_expression(expression, cell_data) continue - values = df_cell[source].values + values = cell_data[source] # Count via len if func_name in ("len", "count"): - result[name] = len(df_cell) + result[name] = n_obs continue # Resolve params: bare column name -> array, expression -> eval'd resolved_params = {} for pkey, pval in params.items(): - if isinstance(pval, str) and pval in df_cell.columns: - resolved_params[pkey] = df_cell[pval].values - elif isinstance(pval, str) and any(c in pval for c in df_cell.columns): - ns = {"__builtins__": {}, "np": np, "numpy": np, - **{c: df_cell[c].values for c in df_cell.columns}} + if isinstance(pval, str) and pval in cell_data: + resolved_params[pkey] = cell_data[pval] + elif isinstance(pval, str) and any(c in pval for c in cell_data): + ns = { + "__builtins__": {}, "np": np, "numpy": np, + **cell_data, + } resolved_params[pkey] = eval(pval, ns) # noqa: S307 else: resolved_params[pkey] = pval @@ -392,13 +431,24 @@ def process_shard( else: stats_arrays[name] = np.zeros(n_cells, dtype=zarr_dtype) - cells_with_data = 0 + # Sort/hash grouping: O(n_obs log n_obs) vs O(n_children x n_obs) cell_col = grid.cells_of(df_all["leaf_id"].values) + col_arrays, cell_to_slice = _build_groups(df_all, cell_col) + _empty: dict[str, np.ndarray] = {col: arr[:0] for col, arr in col_arrays.items()} + + cells_with_data = 0 for i, child_morton in enumerate(children): - df_cell = df_all[cell_col == child_morton] - if len(df_cell) > 0: + if child_morton in cell_to_slice: + start, end = cell_to_slice[child_morton] + cell_data: dict[str, np.ndarray] = { + col: arr[start:end] for col, arr in col_arrays.items() + } cells_with_data += 1 - stats = calculate_cell_statistics(df_cell, value_col="h_li", sigma_col="s_li", config=config) + else: + cell_data = _empty + stats = calculate_cell_statistics( + cell_data, value_col="h_li", sigma_col="s_li", config=config + ) for key, value in stats.items(): stats_arrays[key][i] = value diff --git a/tests/test_processing.py b/tests/test_processing.py index 51e20f55..5e5d59b2 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -7,6 +7,7 @@ from zagg.config import default_config, get_agg_fields, get_coords, get_data_vars from zagg.grids import HealpixGrid from zagg.processing import ( + _build_groups, calculate_cell_statistics, write_dataframe_to_zarr, ) @@ -65,21 +66,21 @@ def test_write_row_count_mismatch(self, mock_dataframe_factory): class TestCalculateCellStatistics: - def test_empty_df_returns_zeros_and_nans(self): - result = calculate_cell_statistics(pd.DataFrame(columns=["h_li", "s_li"])) + def test_empty_data_returns_zeros_and_nans(self): + result = calculate_cell_statistics({"h_li": np.array([]), "s_li": np.array([])}) assert result["count"] == 0 for name in get_agg_fields(default_config()): if name != "count": assert np.isnan(result[name]), f"{name} should be NaN for empty input" def test_result_keys_match_data_vars(self): - df = pd.DataFrame({"h_li": [1.0, 2.0, 3.0], "s_li": [0.1, 0.1, 0.1]}) - result = calculate_cell_statistics(df) + data = {"h_li": np.array([1.0, 2.0, 3.0]), "s_li": np.array([0.1, 0.1, 0.1])} + result = calculate_cell_statistics(data) assert list(result.keys()) == get_data_vars(default_config()) def test_basic_statistics(self): - df = pd.DataFrame({"h_li": [1.0, 2.0, 3.0], "s_li": [0.1, 0.1, 0.1]}) - result = calculate_cell_statistics(df) + data = {"h_li": np.array([1.0, 2.0, 3.0]), "s_li": np.array([0.1, 0.1, 0.1])} + result = calculate_cell_statistics(data) assert result["count"] == 3 assert result["h_min"] == 1.0 assert result["h_max"] == 3.0 @@ -87,8 +88,8 @@ def test_basic_statistics(self): def test_with_explicit_config(self): cfg = default_config() - df = pd.DataFrame({"h_li": [10.0, 20.0, 30.0], "s_li": [0.1, 0.2, 0.1]}) - result = calculate_cell_statistics(df, config=cfg) + data = {"h_li": np.array([10.0, 20.0, 30.0]), "s_li": np.array([0.1, 0.2, 0.1])} + result = calculate_cell_statistics(data, config=cfg) assert result["count"] == 3 assert result["h_min"] == 10.0 assert result["h_max"] == 30.0 @@ -98,6 +99,76 @@ def test_with_explicit_config(self): ) +class TestBuildGroups: + def test_slice_counts_match_per_cell_mask(self): + """_build_groups produces identical cell populations as the old boolean-mask loop.""" + rng = np.random.default_rng(42) + cells = np.array([10, 10, 20, 10, 30, 20, 30], dtype=np.int64) + h_vals = rng.standard_normal(len(cells)) + s_vals = np.abs(rng.standard_normal(len(cells))) + 0.01 + df = pd.DataFrame({"h_li": h_vals, "s_li": s_vals, "leaf_id": cells}) + + col_arrays, cell_to_slice = _build_groups(df, cells) + + for cell_id in [10, 20, 30]: + start, end = cell_to_slice[cell_id] + new_vals = col_arrays["h_li"][start:end] + old_vals = h_vals[cells == cell_id] + np.testing.assert_array_equal(np.sort(new_vals), np.sort(old_vals)) + + def test_boundary_positions(self): + cells = np.array([1, 1, 2, 3, 3, 3], dtype=np.int64) + df = pd.DataFrame({"h_li": np.zeros(6), "s_li": np.ones(6), "leaf_id": cells}) + col_arrays, cell_to_slice = _build_groups(df, cells) + start_1, end_1 = cell_to_slice[1] + start_2, end_2 = cell_to_slice[2] + start_3, end_3 = cell_to_slice[3] + assert end_1 - start_1 == 2 + assert end_2 - start_2 == 1 + assert end_3 - start_3 == 3 + + def test_absent_cell_not_in_map(self): + cells = np.array([1, 2], dtype=np.int64) + df = pd.DataFrame({"h_li": np.zeros(2), "s_li": np.ones(2), "leaf_id": cells}) + _, cell_to_slice = _build_groups(df, cells) + assert 99 not in cell_to_slice + + def test_statistics_match_old_approach(self): + """Sort-group statistics are identical to boolean-mask statistics.""" + rng = np.random.default_rng(7) + n = 200 + child_ids = np.array([100, 200, 300, 400], dtype=np.int64) + cells = rng.choice(child_ids, size=n) + h_vals = rng.standard_normal(n).astype(np.float32) + s_vals = np.abs(rng.standard_normal(n)).astype(np.float32) + 0.01 + df = pd.DataFrame({"h_li": h_vals, "s_li": s_vals, "leaf_id": cells}) + + cfg = default_config() + col_arrays, cell_to_slice = _build_groups(df, cells) + _empty = {col: arr[:0] for col, arr in col_arrays.items()} + + for child_id in child_ids: + # New sort/hash approach + if child_id in cell_to_slice: + s, e = cell_to_slice[child_id] + new_data = {col: arr[s:e] for col, arr in col_arrays.items()} + else: + new_data = _empty + new_stats = calculate_cell_statistics(new_data, config=cfg) + + # Reference: boolean-mask approach + mask = cells == child_id + old_data = {"h_li": h_vals[mask], "s_li": s_vals[mask], "leaf_id": cells[mask]} + old_stats = calculate_cell_statistics(old_data, config=cfg) + + for key in new_stats: + if np.isnan(new_stats[key]) and np.isnan(old_stats[key]): + continue + np.testing.assert_array_equal( + new_stats[key], old_stats[key], err_msg=f"{key} mismatch for cell {child_id}" + ) + + class TestDataSource: """Test data_source section of default config (replaces old DataSourceConfig tests).""" From faa09c43379b7c062828aefefdd4886cf6ddecee Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 12 Jun 2026 23:21:06 +0000 Subject: [PATCH 2/8] fix review findings from PR #33 --- src/zagg/processing.py | 2 ++ tests/test_config.py | 3 ++- tests/test_processing.py | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index da64f987..57d96ba3 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -61,6 +61,8 @@ def _build_groups( sort_idx = np.argsort(cell_col, kind="stable") sorted_cells = cell_col[sort_idx] col_arrays = {col: df_all[col].values[sort_idx] for col in df_all.columns} + if len(sorted_cells) == 0: + return col_arrays, {} boundaries = np.flatnonzero(np.diff(sorted_cells)) + 1 starts = np.concatenate([[0], boundaries]) ends = np.concatenate([boundaries, [len(sorted_cells)]]) diff --git a/tests/test_config.py b/tests/test_config.py index bea099bb..3a5122e1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -302,7 +302,8 @@ def test_expression_in_params(self, synthetic_df): assert result == pytest.approx(expected, rel=1e-5) def test_config_matches_calculate_cell_statistics(self, atl06_config, synthetic_df): - expected = calculate_cell_statistics(synthetic_df) + cell_data = {col: synthetic_df[col].values for col in synthetic_df.columns} + expected = calculate_cell_statistics(cell_data) agg_fields = get_agg_fields(atl06_config) for name, meta in agg_fields.items(): diff --git a/tests/test_processing.py b/tests/test_processing.py index 5e5d59b2..941fc044 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -114,7 +114,7 @@ def test_slice_counts_match_per_cell_mask(self): start, end = cell_to_slice[cell_id] new_vals = col_arrays["h_li"][start:end] old_vals = h_vals[cells == cell_id] - np.testing.assert_array_equal(np.sort(new_vals), np.sort(old_vals)) + np.testing.assert_array_equal(new_vals, old_vals) def test_boundary_positions(self): cells = np.array([1, 1, 2, 3, 3, 3], dtype=np.int64) From 8140437dc16760d10befadbe4f53d2f1d4649ee7 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 16:40:34 +0000 Subject: [PATCH 3/8] phase 2 of issue #30: arrow handoff path --- benchmarks/handoff_bench.py | 138 ++++++++++++++++++++++++++++++++++++ src/zagg/processing.py | 123 +++++++++++++++++++++++--------- tests/test_processing.py | 66 ++++++++++++++++- 3 files changed, 290 insertions(+), 37 deletions(-) create mode 100644 benchmarks/handoff_bench.py diff --git a/benchmarks/handoff_bench.py b/benchmarks/handoff_bench.py new file mode 100644 index 00000000..60d5cb40 --- /dev/null +++ b/benchmarks/handoff_bench.py @@ -0,0 +1,138 @@ +"""Synthetic benchmark for the per-cell aggregation handoff (issue #30). + +Times the per-shard grouping + aggregation for three approaches on synthetic +in-memory observations and asserts they produce byte-for-byte identical stats: + + * ``mask-loop`` -- the pre-#30 O(n_children x n_obs) boolean-mask loop (reference) + * ``pandas-group`` -- sort/hash grouping, pandas carrier (current default) + * ``arrow-group`` -- sort/hash grouping, Arrow carrier (opt-in, this phase) + +This is the CI-runnable half of #30's benchmark: it isolates the grouping +algorithm and the carrier representation cost with no I/O, so it runs anywhere +without credentials. The real-data (ATL03 region) timings — which also exercise +h5coro/S3 reads and real density regimes — land as phase 3 (needs earthaccess/S3). + +Memory is reported via ``tracemalloc`` (Python-domain peak): it does not capture +raw numpy data buffers, but it does capture the pandas BlockManager/Index and +Arrow wrapper overhead, which is where the two carriers actually differ. The +phase-3 real-shard script reports process RSS instead. + +Run:: + + uv run python benchmarks/handoff_bench.py --n-obs 2000000 --n-cells 4096 +""" + +import argparse +import time +import tracemalloc + +import numpy as np +import pandas as pd + +from zagg.config import default_config +from zagg.processing import _build_groups, _group_columns, calculate_cell_statistics + + +def make_synthetic(n_obs: int, n_cells: int, seed: int = 0): + """Random observations spread across ``n_cells`` cells (shuffled, not pre-sorted).""" + rng = np.random.default_rng(seed) + cells = rng.integers(0, n_cells, size=n_obs).astype(np.int64) + h_li = (rng.standard_normal(n_obs) * 50.0).astype(np.float32) + s_li = (np.abs(rng.standard_normal(n_obs)) + 0.01).astype(np.float32) + return {"h_li": h_li, "s_li": s_li, "leaf_id": cells}, cells + + +def agg_mask_loop(col_dict, cell_col, children, cfg): + """Reference: one boolean mask per child cell (the pre-#30 hot loop).""" + stats = {} + for child in children: + mask = cell_col == child + cell_data = {k: v[mask] for k, v in col_dict.items()} + stats[int(child)] = calculate_cell_statistics(cell_data, config=cfg) + return stats + + +def agg_grouped(col_arrays, cell_to_slice, children, cfg): + """Sort/hash grouping: one contiguous slice per child cell.""" + empty = {k: v[:0] for k, v in col_arrays.items()} + stats = {} + for child in children: + child = int(child) + if child in cell_to_slice: + s, e = cell_to_slice[child] + cell_data = {k: v[s:e] for k, v in col_arrays.items()} + else: + cell_data = empty + stats[child] = calculate_cell_statistics(cell_data, config=cfg) + return stats + + +def timed(fn): + tracemalloc.start() + t0 = time.perf_counter() + out = fn() + dt = time.perf_counter() - t0 + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + return out, dt, peak / 1e6 + + +def stats_equal(a, b): + for child in a: + for key in a[child]: + x, y = a[child][key], b[child][key] + if np.isnan(x) and np.isnan(y): + continue + if x != y: + return False, (child, key, x, y) + return True, None + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--n-obs", type=int, default=2_000_000) + ap.add_argument("--n-cells", type=int, default=4096) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + cfg = default_config() + col_dict, cell_col = make_synthetic(args.n_obs, args.n_cells, args.seed) + children = np.arange(args.n_cells, dtype=np.int64) + + ref, dt_mask, mem_mask = timed(lambda: agg_mask_loop(col_dict, cell_col, children, cfg)) + + def run_pandas(): + df = pd.DataFrame(col_dict) + col_arrays, cell_to_slice = _build_groups(df, cell_col) + return agg_grouped(col_arrays, cell_to_slice, children, cfg) + + res_pd, dt_pd, mem_pd = timed(run_pandas) + + def run_arrow(): + import pyarrow as pa + + table = pa.table(col_dict).combine_chunks() + leaf = table.column("leaf_id").to_numpy(zero_copy_only=False) + carrier = {n: table.column(n).to_numpy(zero_copy_only=False) for n in table.column_names} + col_arrays, cell_to_slice = _group_columns(carrier, leaf) + return agg_grouped(col_arrays, cell_to_slice, children, cfg) + + res_ar, dt_ar, mem_ar = timed(run_arrow) + + ok_pd, diff_pd = stats_equal(ref, res_pd) + ok_ar, diff_ar = stats_equal(ref, res_ar) + assert ok_pd, f"pandas grouping diverged from mask loop: {diff_pd}" + assert ok_ar, f"arrow grouping diverged from mask loop: {diff_ar}" + + print(f"n_obs={args.n_obs:,} n_cells={args.n_cells:,} parity: OK") + print(f"{'approach':<16}{'wall_s':>10}{'peak_MB':>12}") + for name, dt, mem in [ + ("mask-loop", dt_mask, mem_mask), + ("pandas-group", dt_pd, mem_pd), + ("arrow-group", dt_ar, mem_ar), + ]: + print(f"{name:<16}{dt:>10.3f}{mem:>12.1f}") + + +if __name__ == "__main__": + main() diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 57d96ba3..35cff70c 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -34,15 +34,41 @@ def _make_url_rewriter(driver: str | None): return lambda url: url.replace("s3://", "", 1) +def _group_columns( + col_dict: dict[str, np.ndarray], + cell_col: np.ndarray, +) -> tuple[dict[str, np.ndarray], dict[int, tuple[int, int]]]: + """Sort column arrays by cell id; return reordered arrays and per-cell slice map. + + Carrier-agnostic core shared by the pandas and Arrow handoff paths. ``col_dict`` + is a plain ``name -> ndarray`` mapping (extracted from a DataFrame or an Arrow + table); the math below is identical regardless of carrier, so both paths produce + byte-for-byte identical groupings and aggregations. + + O(n log n) replacement for the O(n_children x n_obs) boolean-mask loop. The + returned arrays are sorted (stably) by ascending cell id; each cell's + observations form a contiguous slice, so ``col_arrays[col][start:end]`` is a + view. + """ + sort_idx = np.argsort(cell_col, kind="stable") + sorted_cells = cell_col[sort_idx] + col_arrays = {col: arr[sort_idx] for col, arr in col_dict.items()} + if len(sorted_cells) == 0: + return col_arrays, {} + boundaries = np.flatnonzero(np.diff(sorted_cells)) + 1 + starts = np.concatenate([[0], boundaries]) + ends = np.concatenate([boundaries, [len(sorted_cells)]]) + cell_to_slice = {int(sorted_cells[s]): (int(s), int(e)) for s, e in zip(starts, ends)} + return col_arrays, cell_to_slice + + def _build_groups( df_all: pd.DataFrame, cell_col: np.ndarray, ) -> tuple[dict[str, np.ndarray], dict[int, tuple[int, int]]]: """Sort observations by cell id; return reordered column arrays and per-cell slice map. - O(n log n) replacement for the O(n_children x n_obs) boolean-mask loop. The - returned arrays are sorted by ascending cell id; each cell's observations form - a contiguous slice, so ``col_arrays[col][start:end]`` is a zero-copy view. + Pandas carrier wrapper over :func:`_group_columns` (extracts ``.values`` once). Parameters ---------- @@ -58,18 +84,8 @@ def _build_groups( cell_to_slice : dict[int, tuple[int, int]] Maps each observed cell id to ``(start, end)`` indices into col_arrays. """ - sort_idx = np.argsort(cell_col, kind="stable") - sorted_cells = cell_col[sort_idx] - col_arrays = {col: df_all[col].values[sort_idx] for col in df_all.columns} - if len(sorted_cells) == 0: - return col_arrays, {} - boundaries = np.flatnonzero(np.diff(sorted_cells)) + 1 - starts = np.concatenate([[0], boundaries]) - ends = np.concatenate([boundaries, [len(sorted_cells)]]) - cell_to_slice = { - int(sorted_cells[s]): (int(s), int(e)) for s, e in zip(starts, ends) - } - return col_arrays, cell_to_slice + col_dict = {col: df_all[col].values for col in df_all.columns} + return _group_columns(col_dict, cell_col) def write_dataframe_to_zarr( @@ -107,8 +123,7 @@ def write_dataframe_to_zarr( expected_count = int(np.prod(grid.chunk_shape)) if len(df_out) != expected_count: raise ValueError( - f"Expected {expected_count} rows for chunk_shape={grid.chunk_shape}, " - f"got {len(df_out)}" + f"Expected {expected_count} rows for chunk_shape={grid.chunk_shape}, got {len(df_out)}" ) chunk_idx = tuple(int(i) for i in chunk_idx) @@ -118,8 +133,10 @@ def write_dataframe_to_zarr( values = values.reshape(grid.chunk_shape) with config.set({"async.concurrency": 128}): array = open_array( - store, path=f"{grid.group_path}/{name}", - zarr_format=3, consolidated=False, + store, + path=f"{grid.group_path}/{name}", + zarr_format=3, + consolidated=False, ) array.set_block_selection(chunk_idx, values) @@ -191,7 +208,9 @@ def calculate_cell_statistics( resolved_params[pkey] = cell_data[pval] elif isinstance(pval, str) and any(c in pval for c in cell_data): ns = { - "__builtins__": {}, "np": np, "numpy": np, + "__builtins__": {}, + "np": np, + "numpy": np, **cell_data, } resolved_params[pkey] = eval(pval, ns) # noqa: S307 @@ -204,8 +223,15 @@ def calculate_cell_statistics( return result -def _read_group(h5obj, group: str, data_source: dict, parent_morton: int, grid): - """Read and spatially filter one HDF5 group, returning a DataFrame or None.""" +def _read_group( + h5obj, group: str, data_source: dict, parent_morton: int, grid, arrow: bool = False +): + """Read and spatially filter one HDF5 group. + + Returns a ``pandas.DataFrame`` (default) or, when ``arrow=True``, a + ``pyarrow.Table`` carrying the identical columns. Returns ``None`` when the + group has no observations in this shard. + """ coordinates = data_source["coordinates"] variables = data_source["variables"] quality_filter = data_source.get("quality_filter") @@ -274,6 +300,10 @@ def _read_group(h5obj, group: str, data_source: dict, parent_morton: int, grid): else: data_dict["leaf_id"] = leaf_sliced + if arrow: + import pyarrow as pa + + return pa.table(data_dict) return pd.DataFrame(data_dict) @@ -286,6 +316,7 @@ def process_shard( h5coro_driver=None, config: PipelineConfig | None = None, driver: str | None = None, + handoff: str = "pandas", ) -> Tuple[pd.DataFrame, ProcessingMetadata]: """Process one shard: read granules, filter to this shard, aggregate, return df. @@ -310,6 +341,11 @@ def process_shard( Defaults to ``default_config()``. driver : str, optional ``"s3"`` (default) or ``"https"``. + handoff : str, optional + Per-cell aggregation carrier: ``"pandas"`` (default) or ``"arrow"``. + Both paths share :func:`_group_columns` and the same numpy reductions, so + scalar outputs are byte-for-byte identical; only the read→concat→extract + representation differs. Opt-in while the two are benchmarked (issue #30). Returns ------- @@ -320,6 +356,8 @@ def process_shard( """ if config is None: config = default_config() + if handoff not in ("pandas", "arrow"): + raise ValueError(f"handoff must be 'pandas' or 'arrow', got {handoff!r}") data_source = config.data_source parent_morton = int(shard_key) @@ -332,9 +370,11 @@ def process_shard( driver = config.data_source.get("driver", "s3") if driver == "https": from h5coro import webdriver + h5coro_driver = webdriver.HTTPDriver else: from h5coro import s3driver + h5coro_driver = s3driver.S3Driver # Prepare metadata @@ -373,7 +413,8 @@ def process_shard( # Build URL rewriter for the active driver _rewrite_url = _make_url_rewriter(driver) - all_dataframes = [] + use_arrow = handoff == "arrow" + all_reads = [] files_processed = 0 # Read files and filter spatially @@ -391,9 +432,9 @@ def process_shard( for g in data_source["groups"]: try: - df = _read_group(h5obj, g, data_source, parent_morton, grid) - if df is not None: - all_dataframes.append(df) + chunk = _read_group(h5obj, g, data_source, parent_morton, grid, arrow=use_arrow) + if chunk is not None: + all_reads.append(chunk) except Exception as e: logger.debug(f" Error reading track {g}: {e}") continue @@ -407,14 +448,31 @@ def process_shard( logger.info(f" Processed {files_processed}/{len(granule_urls)} files") metadata["files_processed"] = files_processed - if not all_dataframes: + if not all_reads: logger.info(f" No data after filtering for morton {parent_morton} - skipping") metadata["error"] = "No data after filtering" metadata["duration_s"] = (datetime.now() - start_time).total_seconds() return pd.DataFrame(), metadata - df_all = pd.concat(all_dataframes, ignore_index=True) - logger.info(f" Read {len(df_all):,} observations") + # Concat the per-group reads and split observations by cell. Both carriers + # feed identical numpy arrays into _group_columns, so outputs match exactly. + if use_arrow: + import pyarrow as pa + + table = pa.concat_tables(all_reads).combine_chunks() + n_obs_total = table.num_rows + leaf_ids = table.column("leaf_id").to_numpy(zero_copy_only=False) + cell_col = grid.cells_of(leaf_ids) + col_dict = { + name: table.column(name).to_numpy(zero_copy_only=False) for name in table.column_names + } + col_arrays, cell_to_slice = _group_columns(col_dict, cell_col) + else: + df_all = pd.concat(all_reads, ignore_index=True) + n_obs_total = len(df_all) + cell_col = grid.cells_of(df_all["leaf_id"].values) + col_arrays, cell_to_slice = _build_groups(df_all, cell_col) + logger.info(f" Read {n_obs_total:,} observations") # Calculate statistics for child cells children = grid.children(parent_morton) @@ -433,9 +491,7 @@ def process_shard( else: stats_arrays[name] = np.zeros(n_cells, dtype=zarr_dtype) - # Sort/hash grouping: O(n_obs log n_obs) vs O(n_children x n_obs) - cell_col = grid.cells_of(df_all["leaf_id"].values) - col_arrays, cell_to_slice = _build_groups(df_all, cell_col) + # Per-cell observation slices (grouped above, carrier-agnostic). _empty: dict[str, np.ndarray] = {col: arr[:0] for col, arr in col_arrays.items()} cells_with_data = 0 @@ -487,8 +543,7 @@ def process_morton_cell( Constructs a stateless ``HealpixGrid`` and forwards to ``process_shard``. """ warnings.warn( - "process_morton_cell is deprecated; use process_shard(grid, shard_key, ...) " - "directly.", + "process_morton_cell is deprecated; use process_shard(grid, shard_key, ...) directly.", DeprecationWarning, stacklevel=2, ) diff --git a/tests/test_processing.py b/tests/test_processing.py index 941fc044..11101eea 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -8,6 +8,7 @@ from zagg.grids import HealpixGrid from zagg.processing import ( _build_groups, + _group_columns, calculate_cell_statistics, write_dataframe_to_zarr, ) @@ -44,9 +45,7 @@ def test_write_dataframe_to_zarr(self, mock_dataframe_factory): def test_write_empty_dataframe(self): grid = HealpixGrid(6, 8, layout="fullsphere") store = MemoryStore() - assert write_dataframe_to_zarr( - pd.DataFrame(), store, grid=grid, chunk_idx=(0,) - ) + assert write_dataframe_to_zarr(pd.DataFrame(), store, grid=grid, chunk_idx=(0,)) def test_write_row_count_mismatch(self, mock_dataframe_factory): parent_order = 6 @@ -169,6 +168,67 @@ def test_statistics_match_old_approach(self): ) +class TestArrowHandoff: + """Phase 2 of #30: the Arrow carrier must match the pandas carrier exactly.""" + + def test_group_columns_matches_build_groups(self): + """_group_columns (carrier-agnostic core) == _build_groups (pandas wrapper).""" + cells = np.array([5, 1, 5, 1, 9], dtype=np.int64) + col_dict = { + "h_li": np.arange(5.0, dtype=np.float32), + "s_li": np.ones(5, dtype=np.float32), + "leaf_id": cells, + } + df = pd.DataFrame(col_dict) + arrays_a, slices_a = _build_groups(df, cells) + arrays_b, slices_b = _group_columns(col_dict, cells) + assert slices_a == slices_b + for key in arrays_a: + np.testing.assert_array_equal(arrays_a[key], arrays_b[key]) + + def test_arrow_grouping_matches_pandas(self): + """Arrow-carrier grouping yields byte-for-byte identical stats to pandas.""" + pa = pytest.importorskip("pyarrow") + rng = np.random.default_rng(11) + n = 500 + child_ids = np.array([100, 200, 300, 400, 500], dtype=np.int64) + cells = rng.choice(child_ids, size=n) + h_vals = (rng.standard_normal(n) * 30.0).astype(np.float32) + s_vals = (np.abs(rng.standard_normal(n)) + 0.01).astype(np.float32) + col_dict = {"h_li": h_vals, "s_li": s_vals, "leaf_id": cells} + cfg = default_config() + + # pandas carrier + df = pd.DataFrame(col_dict) + p_arrays, p_slices = _build_groups(df, cells) + + # arrow carrier: read the columns back as numpy and group identically + table = pa.table(col_dict).combine_chunks() + a_carrier = { + name: table.column(name).to_numpy(zero_copy_only=False) for name in table.column_names + } + a_leaf = table.column("leaf_id").to_numpy(zero_copy_only=False) + a_arrays, a_slices = _group_columns(a_carrier, a_leaf) + + assert p_slices == a_slices + for child in child_ids: + child = int(child) + ps, pe = p_slices[child] + as_, ae = a_slices[child] + p_stats = calculate_cell_statistics( + {k: v[ps:pe] for k, v in p_arrays.items()}, config=cfg + ) + a_stats = calculate_cell_statistics( + {k: v[as_:ae] for k, v in a_arrays.items()}, config=cfg + ) + for key in p_stats: + if np.isnan(p_stats[key]) and np.isnan(a_stats[key]): + continue + np.testing.assert_array_equal( + p_stats[key], a_stats[key], err_msg=f"{key} mismatch for cell {child}" + ) + + class TestDataSource: """Test data_source section of default config (replaces old DataSourceConfig tests).""" From 39d24b63c1f3b2e9d2165f7678b1727dbe2afa6f Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 15 Jun 2026 02:34:02 +0000 Subject: [PATCH 4/8] address review: null-free arrow guard + _concat_and_group helper/test (#30) --- src/zagg/processing.py | 73 +++++++++++++++++++++++++++++---------- tests/test_processing.py | 74 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 18 deletions(-) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 35cff70c..659e8393 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -88,6 +88,58 @@ def _build_groups( return _group_columns(col_dict, cell_col) +def _concat_and_group(all_reads, grid, handoff: str): + """Concat the per-group reads and split observations by cell. + + Carrier-agnostic seam shared by :func:`process_shard` and its tests, so the + Arrow path is exercised end-to-end (including multi-table ``concat_tables`` + ordering) rather than re-assembled inline. Both carriers feed identical numpy + arrays into :func:`_group_columns`, so the groupings — and the aggregations + computed from them — are byte-for-byte identical. + + Parameters + ---------- + all_reads : list + Per-group reads from ``_read_group``: ``pandas.DataFrame`` for the pandas + carrier, ``pyarrow.Table`` for the arrow carrier. + grid : OutputGrid + Provides ``cells_of`` to map leaf ids to child cell ids. + handoff : {"pandas", "arrow"} + Which carrier ``all_reads`` holds. + + Returns + ------- + col_arrays : dict[str, np.ndarray] + Column arrays sorted in ascending cell-id order. + cell_to_slice : dict[int, tuple[int, int]] + Maps each observed cell id to ``(start, end)`` into ``col_arrays``. + n_obs_total : int + Total observation count across all reads. + """ + if handoff == "arrow": + import pyarrow as pa + + table = pa.concat_tables(all_reads).combine_chunks() + # The arrow handoff requires dense, null-free columns: ``_read_group`` + # builds tables from raw h5coro reads (no null mask), so + # ``to_numpy(zero_copy_only=False)`` is dtype-exact and matches ``.values`` + # on the pandas side. Guard the precondition so a future nullable source + # can't silently diverge the two carriers instead of failing loudly. + null_cols = [n for n in table.column_names if table.column(n).null_count] + if null_cols: + raise ValueError(f"arrow handoff requires null-free columns; got nulls in {null_cols}") + n_obs_total = table.num_rows + cell_col = grid.cells_of(table.column("leaf_id").to_numpy(zero_copy_only=False)) + col_dict = {n: table.column(n).to_numpy(zero_copy_only=False) for n in table.column_names} + col_arrays, cell_to_slice = _group_columns(col_dict, cell_col) + else: + df_all = pd.concat(all_reads, ignore_index=True) + n_obs_total = len(df_all) + cell_col = grid.cells_of(df_all["leaf_id"].values) + col_arrays, cell_to_slice = _build_groups(df_all, cell_col) + return col_arrays, cell_to_slice, n_obs_total + + def write_dataframe_to_zarr( df_out: pd.DataFrame, store: Store, @@ -454,24 +506,9 @@ def process_shard( metadata["duration_s"] = (datetime.now() - start_time).total_seconds() return pd.DataFrame(), metadata - # Concat the per-group reads and split observations by cell. Both carriers - # feed identical numpy arrays into _group_columns, so outputs match exactly. - if use_arrow: - import pyarrow as pa - - table = pa.concat_tables(all_reads).combine_chunks() - n_obs_total = table.num_rows - leaf_ids = table.column("leaf_id").to_numpy(zero_copy_only=False) - cell_col = grid.cells_of(leaf_ids) - col_dict = { - name: table.column(name).to_numpy(zero_copy_only=False) for name in table.column_names - } - col_arrays, cell_to_slice = _group_columns(col_dict, cell_col) - else: - df_all = pd.concat(all_reads, ignore_index=True) - n_obs_total = len(df_all) - cell_col = grid.cells_of(df_all["leaf_id"].values) - col_arrays, cell_to_slice = _build_groups(df_all, cell_col) + # Concat the per-group reads and split observations by cell (carrier-agnostic; + # both carriers feed identical numpy arrays into _group_columns). + col_arrays, cell_to_slice, n_obs_total = _concat_and_group(all_reads, grid, handoff) logger.info(f" Read {n_obs_total:,} observations") # Calculate statistics for child cells diff --git a/tests/test_processing.py b/tests/test_processing.py index 11101eea..f95a26b2 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -8,6 +8,7 @@ from zagg.grids import HealpixGrid from zagg.processing import ( _build_groups, + _concat_and_group, _group_columns, calculate_cell_statistics, write_dataframe_to_zarr, @@ -228,6 +229,79 @@ def test_arrow_grouping_matches_pandas(self): p_stats[key], a_stats[key], err_msg=f"{key} mismatch for cell {child}" ) + def test_concat_and_group_arrow_matches_pandas(self): + """_concat_and_group drives the real carrier path (incl. multi-table concat).""" + pa = pytest.importorskip("pyarrow") + + class _IdentityGrid: + # leaf_id already is the cell id for this test; isolate the carrier + # path from grid semantics. + @staticmethod + def cells_of(leaf_ids): + return np.asarray(leaf_ids) + + grid = _IdentityGrid() + cfg = default_config() + rng = np.random.default_rng(7) + child_ids = np.array([100, 200, 300, 400, 500], dtype=np.int64) + + # Three reads of differing length -> exercises concat ordering / offsets. + reads = [] + for n in (40, 7, 53): + cells = rng.choice(child_ids, size=n) + reads.append( + { + "h_li": (rng.standard_normal(n) * 30.0).astype(np.float32), + "s_li": (np.abs(rng.standard_normal(n)) + 0.01).astype(np.float32), + "leaf_id": cells, + } + ) + pandas_reads = [pd.DataFrame(r) for r in reads] + arrow_reads = [pa.table(r) for r in reads] + + p_arrays, p_slices, p_n = _concat_and_group(pandas_reads, grid, "pandas") + a_arrays, a_slices, a_n = _concat_and_group(arrow_reads, grid, "arrow") + + assert p_n == a_n == sum(len(r["leaf_id"]) for r in reads) + assert p_slices == a_slices + for child in child_ids: + child = int(child) + if child not in p_slices: + continue + ps, pe = p_slices[child] + as_, ae = a_slices[child] + p_stats = calculate_cell_statistics( + {k: v[ps:pe] for k, v in p_arrays.items()}, config=cfg + ) + a_stats = calculate_cell_statistics( + {k: v[as_:ae] for k, v in a_arrays.items()}, config=cfg + ) + for key in p_stats: + if np.isnan(p_stats[key]) and np.isnan(a_stats[key]): + continue + np.testing.assert_array_equal( + p_stats[key], a_stats[key], err_msg=f"{key} mismatch for cell {child}" + ) + + def test_concat_and_group_arrow_rejects_nulls(self): + """The arrow carrier must fail loudly on null columns, not silently diverge.""" + pa = pytest.importorskip("pyarrow") + + class _IdentityGrid: + @staticmethod + def cells_of(leaf_ids): + return np.asarray(leaf_ids) + + table = pa.table( + { + "h_li": pa.array([1.0, None, 3.0], type=pa.float32()), + "s_li": pa.array([0.1, 0.2, 0.3], type=pa.float32()), + "leaf_id": pa.array([100, 200, 100], type=pa.int64()), + } + ) + with pytest.raises(ValueError, match="null-free"): + _concat_and_group([table], _IdentityGrid(), "arrow") + class TestDataSource: """Test data_source section of default config (replaces old DataSourceConfig tests).""" From d580746c452d693b29c94cd2264b48d0ddd84212 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 15 Jun 2026 02:43:31 +0000 Subject: [PATCH 5/8] fold self-review: dedup grid stub, null-guard test targets leaf_id (#30) --- tests/test_processing.py | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/tests/test_processing.py b/tests/test_processing.py index f95a26b2..6354981f 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -15,6 +15,15 @@ ) +class _IdentityGrid: + """Grid stub whose cell id is the leaf id, isolating the carrier path in + ``_concat_and_group`` tests from real grid semantics.""" + + @staticmethod + def cells_of(leaf_ids): + return np.asarray(leaf_ids) + + class TestWriteDataframeToZarr: def test_write_dataframe_to_zarr(self, mock_dataframe_factory): parent_order = 6 @@ -233,13 +242,6 @@ def test_concat_and_group_arrow_matches_pandas(self): """_concat_and_group drives the real carrier path (incl. multi-table concat).""" pa = pytest.importorskip("pyarrow") - class _IdentityGrid: - # leaf_id already is the cell id for this test; isolate the carrier - # path from grid semantics. - @staticmethod - def cells_of(leaf_ids): - return np.asarray(leaf_ids) - grid = _IdentityGrid() cfg = default_config() rng = np.random.default_rng(7) @@ -284,19 +286,19 @@ def cells_of(leaf_ids): ) def test_concat_and_group_arrow_rejects_nulls(self): - """The arrow carrier must fail loudly on null columns, not silently diverge.""" - pa = pytest.importorskip("pyarrow") + """The arrow carrier must fail loudly on null columns, not silently diverge. - class _IdentityGrid: - @staticmethod - def cells_of(leaf_ids): - return np.asarray(leaf_ids) + The null is in ``leaf_id`` — the grouping key — which is the case the guard + exists to catch: a null there would corrupt the cell assignment under + ``to_numpy``, not just a single stat. + """ + pa = pytest.importorskip("pyarrow") table = pa.table( { - "h_li": pa.array([1.0, None, 3.0], type=pa.float32()), + "h_li": pa.array([1.0, 2.0, 3.0], type=pa.float32()), "s_li": pa.array([0.1, 0.2, 0.3], type=pa.float32()), - "leaf_id": pa.array([100, 200, 100], type=pa.int64()), + "leaf_id": pa.array([100, None, 100], type=pa.int64()), } ) with pytest.raises(ValueError, match="null-free"): From 6194852a57fcaf0f8851ed4e1d6deaa259a1689a Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 15 Jun 2026 10:17:50 +0000 Subject: [PATCH 6/8] phase 2b of issue #30: experimental pyarrow kernel reducer --- benchmarks/handoff_bench.py | 85 ++++++++++-- src/zagg/processing.py | 252 +++++++++++++++++++++++++++++------- tests/test_processing.py | 101 +++++++++++++++ 3 files changed, 383 insertions(+), 55 deletions(-) diff --git a/benchmarks/handoff_bench.py b/benchmarks/handoff_bench.py index 60d5cb40..b5448e59 100644 --- a/benchmarks/handoff_bench.py +++ b/benchmarks/handoff_bench.py @@ -1,20 +1,27 @@ """Synthetic benchmark for the per-cell aggregation handoff (issue #30). -Times the per-shard grouping + aggregation for three approaches on synthetic -in-memory observations and asserts they produce byte-for-byte identical stats: +Times the per-shard grouping + aggregation for four approaches on synthetic +in-memory observations: * ``mask-loop`` -- the pre-#30 O(n_children x n_obs) boolean-mask loop (reference) * ``pandas-group`` -- sort/hash grouping, pandas carrier (current default) - * ``arrow-group`` -- sort/hash grouping, Arrow carrier (opt-in, this phase) + * ``arrow-group`` -- sort/hash grouping, Arrow carrier (opt-in) + * ``arrow-kernel`` -- EXPERIMENTAL pyarrow.compute hash-aggregate reducer (phase 2b) + +The first three feed identical numpy arrays into the reducer, so their stats are +asserted byte-for-byte identical. ``arrow-kernel`` reduces via pyarrow's C++ hash +kernels, whose float mean/variance differ from numpy by ~1 ULP; it is asserted +*close* (``np.allclose`` at ``KERNEL_RTOL``), not identical. This is the CI-runnable half of #30's benchmark: it isolates the grouping -algorithm and the carrier representation cost with no I/O, so it runs anywhere -without credentials. The real-data (ATL03 region) timings — which also exercise -h5coro/S3 reads and real density regimes — land as phase 3 (needs earthaccess/S3). +algorithm, the carrier representation cost, and the kernel reducer with no I/O, +so it runs anywhere without credentials. The real-data (ATL03 region) timings — +which decide whether the experimental kernel path is kept — land as phase 3 +(needs earthaccess/S3). Memory is reported via ``tracemalloc`` (Python-domain peak): it does not capture raw numpy data buffers, but it does capture the pandas BlockManager/Index and -Arrow wrapper overhead, which is where the two carriers actually differ. The +Arrow wrapper overhead, which is where the carriers actually differ. The phase-3 real-shard script reports process RSS instead. Run:: @@ -29,8 +36,22 @@ import numpy as np import pandas as pd -from zagg.config import default_config -from zagg.processing import _build_groups, _group_columns, calculate_cell_statistics +from zagg.config import default_config, get_data_vars +from zagg.processing import ( + KERNEL_RTOL, + _build_groups, + _group_columns, + _kernel_aggregate, + calculate_cell_statistics, +) + + +class _IdentityGrid: + """Grid stub: the synthetic ``leaf_id`` already *is* the destination cell.""" + + @staticmethod + def cells_of(leaf_ids): + return np.asarray(leaf_ids) def make_synthetic(n_obs: int, n_cells: int, seed: int = 0): @@ -88,6 +109,26 @@ def stats_equal(a, b): return True, None +def stats_close(a, b, rtol): + """Like ``stats_equal`` but tolerant — for the kernel path's float divergence.""" + for child in a: + for key in a[child]: + x, y = float(a[child][key]), float(b[child][key]) + if np.isnan(x) and np.isnan(y): + continue + if not np.isclose(x, y, rtol=rtol, equal_nan=True): + return False, (child, key, x, y) + return True, None + + +def kernel_arrays_to_stats(stats_arrays, children): + """Adapt ``_kernel_aggregate``'s ``name -> ndarray`` output to the per-cell dict.""" + return { + int(child): {name: stats_arrays[name][i] for name in stats_arrays} + for i, child in enumerate(children) + } + + def main(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--n-obs", type=int, default=2_000_000) @@ -119,17 +160,39 @@ def run_arrow(): res_ar, dt_ar, mem_ar = timed(run_arrow) + def run_kernel(): + import pyarrow as pa + + table = pa.table(col_dict).combine_chunks() + cell = _IdentityGrid.cells_of(table.column("leaf_id").to_numpy(zero_copy_only=False)) + out = _kernel_aggregate(table, cell, children, "h_li", cfg) + return kernel_arrays_to_stats(out["stats_arrays"], children) + + res_kn, dt_kn, mem_kn = timed(run_kernel) + ok_pd, diff_pd = stats_equal(ref, res_pd) ok_ar, diff_ar = stats_equal(ref, res_ar) + ok_kn, diff_kn = stats_close(ref, res_kn, KERNEL_RTOL) assert ok_pd, f"pandas grouping diverged from mask loop: {diff_pd}" assert ok_ar, f"arrow grouping diverged from mask loop: {diff_ar}" - - print(f"n_obs={args.n_obs:,} n_cells={args.n_cells:,} parity: OK") + assert ok_kn, f"kernel reducer diverged beyond rtol={KERNEL_RTOL}: {diff_kn}" + + # Cross-check that count/min/max (the integral / extremum stats) are *exact* + # under the kernel path, isolating the divergence to the float reductions. + exact = [n for n in get_data_vars(cfg) if n in ("count", "h_min", "h_max")] + for child in ref: + for name in exact: + assert float(ref[child][name]) == float(res_kn[child][name]) or ( + np.isnan(ref[child][name]) and np.isnan(res_kn[child][name]) + ), f"kernel {name} not exact for cell {child}" + + print(f"n_obs={args.n_obs:,} n_cells={args.n_cells:,} parity: OK (kernel rtol={KERNEL_RTOL})") print(f"{'approach':<16}{'wall_s':>10}{'peak_MB':>12}") for name, dt, mem in [ ("mask-loop", dt_mask, mem_mask), ("pandas-group", dt_pd, mem_pd), ("arrow-group", dt_ar, mem_ar), + ("arrow-kernel", dt_kn, mem_kn), ]: print(f"{name:<16}{dt:>10.3f}{mem:>12.1f}") diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 659e8393..6ee76f00 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -275,6 +275,148 @@ def calculate_cell_statistics( return result +# EXPERIMENTAL (phase 2b of #30) ----------------------------------------------- +# Optional pyarrow.compute hash-aggregate ("kernel") reduction path. Unlike the +# pandas/arrow *carriers* — which feed identical numpy arrays into +# ``calculate_cell_statistics`` and are therefore byte-for-byte identical — the +# kernel path computes the kernel-able reductions in a single vectorised C++ +# pass (Acero ``TableGroupBy.aggregate``). pyarrow's float summation differs from +# numpy's, so its ``mean``/``variance`` outputs are NOT byte-identical to the +# numpy path; they agree only within ``KERNEL_RTOL`` (validated in tests and in +# ``benchmarks/handoff_bench.py``). This lever is opt-in via +# ``handoff="arrow-kernel"`` and exists purely so phase 3 can benchmark it on +# real ATL03 data; it is kept gated and clearly experimental, and should be +# dropped if that benchmark shows no material speedup (see PR #33 discussion). + +# Documented tolerance for kernel-vs-numpy float agreement. float32 means/variance +# over millions of obs diverge by ~1 ULP (~1e-6 relative); 1e-5 leaves headroom. +KERNEL_RTOL = 1e-5 + +# numpy/config ``function`` name -> pyarrow hash-aggregate function name. Only +# reductions that are mathematically a pure (unweighted) group reduction appear +# here; weighted ``average``, ``quantile`` (only approximate via tdigest) and any +# ``expression`` field fall back to the per-cell numpy path. +_KERNEL_FUNCS = { + "len": "count", + "count": "count", + "min": "min", + "max": "max", + "var": "variance", + "average": "mean", +} + + +def _kernel_able(meta: dict) -> bool: + """Whether an agg field can be computed by a pyarrow hash-aggregate kernel. + + EXPERIMENTAL. Excludes expression fields, weighted ``average`` (no weighted + hash kernel), quantiles (only approximate tdigest), and anything whose + function has no pure-reduction kernel equivalent. + """ + if meta.get("expression"): + return False + func = meta.get("function") + if func not in _KERNEL_FUNCS: + return False + # Weighted average is not a pure hash reduction. + if func == "average" and "weights" in (meta.get("params") or {}): + return False + return True + + +def _kernel_aggregate( + table, cell_col: np.ndarray, children, value_col: str, config: PipelineConfig +) -> dict: + """EXPERIMENTAL pyarrow hash-aggregate reducer (phase 2b of #30). + + Computes the kernel-able stats (count/min/max/variance/unweighted-mean) for + every child cell in one vectorised ``TableGroupBy.aggregate`` pass, then fills + the remaining (weighted mean, expression, quantile) fields via the per-cell + numpy path so output columns match the default reducer exactly in shape. + + Kernel-computed float stats are NOT byte-identical to numpy — they agree + within :data:`KERNEL_RTOL`. Returns ``stats_arrays`` (``name -> ndarray`` over + ``children``) plus ``cells_with_data``. + + Parameters + ---------- + table : pyarrow.Table + Concatenated, null-free observations (one row per observation). + cell_col : np.ndarray + Child cell id for each row of ``table`` (already ``grid.cells_of`` mapped, + so the group key is the destination cell, not the leaf id). + children : sequence of int + Child cell ids, in canonical chunk order. + value_col : str + Default value column for fields without an explicit ``source``. + config : PipelineConfig + Drives the agg-field metadata. + """ + import pyarrow as pa + + agg_fields = get_agg_fields(config) + data_vars = get_data_vars(config) + n_cells = len(children) + child_index = {int(c): i for i, c in enumerate(children)} + + stats_arrays: dict[str, np.ndarray] = {} + for name in data_vars: + meta = agg_fields[name] + zarr_dtype = np.dtype(meta.get("dtype", "float32")) + fill_value = meta.get("fill_value", "NaN") + if fill_value == "NaN": + stats_arrays[name] = np.full(n_cells, np.nan, dtype=zarr_dtype) + else: + stats_arrays[name] = np.zeros(n_cells, dtype=zarr_dtype) + + kernel_names = [n for n in data_vars if _kernel_able(agg_fields[n])] + fallback_names = [n for n in data_vars if n not in kernel_names] + + # Group by the destination cell id (not the raw leaf id): append cell_col and + # run one vectorised group-by + reduction pass for all kernel-able fields. + keyed = table.append_column("_cell", pa.array(np.asarray(cell_col))) + aggregations = [ + (agg_fields[n].get("source") or value_col, _KERNEL_FUNCS[agg_fields[n]["function"]]) + for n in kernel_names + ] + gd = keyed.group_by("_cell").aggregate(aggregations).to_pydict() + group_cells = gd["_cell"] + # Map each grouped row back to its position in ``children``. + row_to_idx = [child_index.get(int(c)) for c in group_cells] + for n, (src, kfunc) in zip(kernel_names, aggregations): + col = gd[f"{src}_{kfunc}"] + out = stats_arrays[n] + for row, idx in enumerate(row_to_idx): + if idx is not None: + out[idx] = col[row] + + cells_with_data = sum(1 for idx in row_to_idx if idx is not None) + + # Fallback fields (and only those) via the per-cell numpy reducer. Reuse the + # carrier-agnostic grouping so the slices match the default path exactly. + if fallback_names: + col_dict = { + name: table.column(name).to_numpy(zero_copy_only=False) for name in table.column_names + } + col_arrays, cell_to_slice = _group_columns(col_dict, np.asarray(cell_col)) + _empty = {col: arr[:0] for col, arr in col_arrays.items()} + for i, child in enumerate(children): + child = int(child) + if child in cell_to_slice: + s, e = cell_to_slice[child] + cell_data = {col: arr[s:e] for col, arr in col_arrays.items()} + else: + cell_data = _empty + stats = calculate_cell_statistics(cell_data, value_col=value_col, config=config) + for name in fallback_names: + stats_arrays[name][i] = stats[name] + + return {"stats_arrays": stats_arrays, "cells_with_data": cells_with_data} + + +# -- end EXPERIMENTAL kernel path --------------------------------------------- + + def _read_group( h5obj, group: str, data_source: dict, parent_morton: int, grid, arrow: bool = False ): @@ -394,10 +536,14 @@ def process_shard( driver : str, optional ``"s3"`` (default) or ``"https"``. handoff : str, optional - Per-cell aggregation carrier: ``"pandas"`` (default) or ``"arrow"``. - Both paths share :func:`_group_columns` and the same numpy reductions, so - scalar outputs are byte-for-byte identical; only the read→concat→extract - representation differs. Opt-in while the two are benchmarked (issue #30). + Per-cell aggregation carrier: ``"pandas"`` (default), ``"arrow"``, or the + EXPERIMENTAL ``"arrow-kernel"``. ``"pandas"`` and ``"arrow"`` share + :func:`_group_columns` and the same numpy reductions, so scalar outputs + are byte-for-byte identical; only the read→concat→extract representation + differs. ``"arrow-kernel"`` (phase 2b of #30) instead reduces via + ``pyarrow.compute`` hash-aggregate kernels: its float ``mean``/``variance`` + differ from numpy by ~1 ULP (agree within :data:`KERNEL_RTOL`, not byte + identical). All three are opt-in while benchmarked (issue #30). Returns ------- @@ -408,8 +554,8 @@ def process_shard( """ if config is None: config = default_config() - if handoff not in ("pandas", "arrow"): - raise ValueError(f"handoff must be 'pandas' or 'arrow', got {handoff!r}") + if handoff not in ("pandas", "arrow", "arrow-kernel"): + raise ValueError(f"handoff must be 'pandas', 'arrow', or 'arrow-kernel', got {handoff!r}") data_source = config.data_source parent_morton = int(shard_key) @@ -465,7 +611,7 @@ def process_shard( # Build URL rewriter for the active driver _rewrite_url = _make_url_rewriter(driver) - use_arrow = handoff == "arrow" + use_arrow = handoff in ("arrow", "arrow-kernel") all_reads = [] files_processed = 0 @@ -506,46 +652,64 @@ def process_shard( metadata["duration_s"] = (datetime.now() - start_time).total_seconds() return pd.DataFrame(), metadata - # Concat the per-group reads and split observations by cell (carrier-agnostic; - # both carriers feed identical numpy arrays into _group_columns). - col_arrays, cell_to_slice, n_obs_total = _concat_and_group(all_reads, grid, handoff) - logger.info(f" Read {n_obs_total:,} observations") - - # Calculate statistics for child cells children = grid.children(parent_morton) - logger.info(f" Calculating statistics for {len(children)} cells...") - - n_cells = len(children) data_vars = get_data_vars(config) - agg_fields = get_agg_fields(config) - stats_arrays = {} - for name in data_vars: - meta = agg_fields[name] - zarr_dtype = np.dtype(meta.get("dtype", "float32")) - fill_value = meta.get("fill_value", "NaN") - if fill_value == "NaN": - stats_arrays[name] = np.full(n_cells, np.nan, dtype=zarr_dtype) - else: - stats_arrays[name] = np.zeros(n_cells, dtype=zarr_dtype) - # Per-cell observation slices (grouped above, carrier-agnostic). - _empty: dict[str, np.ndarray] = {col: arr[:0] for col, arr in col_arrays.items()} - - cells_with_data = 0 - for i, child_morton in enumerate(children): - if child_morton in cell_to_slice: - start, end = cell_to_slice[child_morton] - cell_data: dict[str, np.ndarray] = { - col: arr[start:end] for col, arr in col_arrays.items() - } - cells_with_data += 1 - else: - cell_data = _empty - stats = calculate_cell_statistics( - cell_data, value_col="h_li", sigma_col="s_li", config=config - ) - for key, value in stats.items(): - stats_arrays[key][i] = value + if handoff == "arrow-kernel": + # EXPERIMENTAL (phase 2b of #30): reduce via pyarrow hash-aggregate kernels + # instead of the per-cell numpy loop. Not byte-identical to the default + # path (float mean/variance diverge by ~1 ULP — see KERNEL_RTOL). + import pyarrow as pa + + table = pa.concat_tables(all_reads).combine_chunks() + null_cols = [n for n in table.column_names if table.column(n).null_count] + if null_cols: + raise ValueError(f"arrow handoff requires null-free columns; got nulls in {null_cols}") + n_obs_total = table.num_rows + cell_col = grid.cells_of(table.column("leaf_id").to_numpy(zero_copy_only=False)) + logger.info(f" Read {n_obs_total:,} observations") + logger.info(f" Calculating statistics for {len(children)} cells (kernel)...") + kernel = _kernel_aggregate(table, cell_col, children, "h_li", config) + stats_arrays = kernel["stats_arrays"] + cells_with_data = kernel["cells_with_data"] + n_cells = len(children) + else: + # Concat the per-group reads and split observations by cell (carrier- + # agnostic; both carriers feed identical numpy arrays into _group_columns). + col_arrays, cell_to_slice, n_obs_total = _concat_and_group(all_reads, grid, handoff) + logger.info(f" Read {n_obs_total:,} observations") + logger.info(f" Calculating statistics for {len(children)} cells...") + + n_cells = len(children) + agg_fields = get_agg_fields(config) + stats_arrays = {} + for name in data_vars: + meta = agg_fields[name] + zarr_dtype = np.dtype(meta.get("dtype", "float32")) + fill_value = meta.get("fill_value", "NaN") + if fill_value == "NaN": + stats_arrays[name] = np.full(n_cells, np.nan, dtype=zarr_dtype) + else: + stats_arrays[name] = np.zeros(n_cells, dtype=zarr_dtype) + + # Per-cell observation slices (grouped above, carrier-agnostic). + _empty: dict[str, np.ndarray] = {col: arr[:0] for col, arr in col_arrays.items()} + + cells_with_data = 0 + for i, child_morton in enumerate(children): + if child_morton in cell_to_slice: + start, end = cell_to_slice[child_morton] + cell_data: dict[str, np.ndarray] = { + col: arr[start:end] for col, arr in col_arrays.items() + } + cells_with_data += 1 + else: + cell_data = _empty + stats = calculate_cell_statistics( + cell_data, value_col="h_li", sigma_col="s_li", config=config + ) + for key, value in stats.items(): + stats_arrays[key][i] = value logger.info(f" Statistics: {cells_with_data}/{n_cells} cells with data") diff --git a/tests/test_processing.py b/tests/test_processing.py index 6354981f..3b648dbc 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -7,9 +7,12 @@ from zagg.config import default_config, get_agg_fields, get_coords, get_data_vars from zagg.grids import HealpixGrid from zagg.processing import ( + KERNEL_RTOL, _build_groups, _concat_and_group, _group_columns, + _kernel_able, + _kernel_aggregate, calculate_cell_statistics, write_dataframe_to_zarr, ) @@ -305,6 +308,104 @@ def test_concat_and_group_arrow_rejects_nulls(self): _concat_and_group([table], _IdentityGrid(), "arrow") +class TestKernelHandoff: + """Phase 2b of #30 (EXPERIMENTAL): the pyarrow hash-aggregate kernel reducer. + + Unlike the pandas<->arrow *carrier* equivalence (byte-for-byte identical), the + kernel path's float mean/variance diverge from numpy by ~1 ULP, so it is + validated within :data:`KERNEL_RTOL`, not by exact equality. + """ + + def _numpy_reference(self, col_dict, cell_col, children, cfg): + """Default per-cell numpy stats, as ``name -> ndarray`` over ``children``.""" + col_arrays, cell_to_slice = _group_columns(col_dict, cell_col) + empty = {c: a[:0] for c, a in col_arrays.items()} + out = {v: np.full(len(children), np.nan, dtype=np.float64) for v in get_data_vars(cfg)} + for i, child in enumerate(children): + child = int(child) + if child in cell_to_slice: + s, e = cell_to_slice[child] + cell_data = {c: a[s:e] for c, a in col_arrays.items()} + else: + cell_data = empty + stats = calculate_cell_statistics(cell_data, config=cfg) + for k, v in stats.items(): + out[k][i] = v + return out + + def test_kernel_able_classification(self): + """count/min/max/var and unweighted average are kernel-able; the rest fall back.""" + cfg = default_config() + fields = get_agg_fields(cfg) + # Default atl06 config: count/h_min/h_max/h_variance are pure reductions; + # h_mean is weighted, h_sigma is an expression, the quantiles are tdigest. + assert _kernel_able(fields["count"]) + assert _kernel_able(fields["h_min"]) + assert _kernel_able(fields["h_max"]) + assert _kernel_able(fields["h_variance"]) + assert not _kernel_able(fields["h_mean"]) # weighted average + assert not _kernel_able(fields["h_sigma"]) # expression + assert not _kernel_able(fields["h_q50"]) # quantile + # Unweighted average would be kernel-able. + assert _kernel_able({"function": "average", "source": "h_li"}) + + def test_kernel_matches_numpy_within_tolerance(self): + """Kernel stats match the numpy reducer within KERNEL_RTOL (exact where integral).""" + pa = pytest.importorskip("pyarrow") + cfg = default_config() + rng = np.random.default_rng(3) + children = np.array([100, 200, 300, 400, 500], dtype=np.int64) + n = 2000 + cells = rng.choice(children, size=n) + h = (rng.standard_normal(n) * 30.0).astype(np.float32) + s = (np.abs(rng.standard_normal(n)) + 0.01).astype(np.float32) + col_dict = {"h_li": h, "s_li": s, "leaf_id": cells} + + ref = self._numpy_reference(col_dict, cells, children, cfg) + table = pa.table(col_dict) + kernel = _kernel_aggregate(table, cells, children, "h_li", cfg) + ks = kernel["stats_arrays"] + + assert kernel["cells_with_data"] == len(children) + # count/min/max are integral or order-independent extrema: exact. + for name in ("count", "h_min", "h_max"): + np.testing.assert_array_equal( + np.asarray(ks[name], dtype=np.float64), ref[name], err_msg=name + ) + # variance is the kernel-reduced float stat: close, not identical. + np.testing.assert_allclose( + np.asarray(ks["h_variance"], dtype=np.float64), + ref["h_variance"], + rtol=KERNEL_RTOL, + equal_nan=True, + ) + # Fallback fields (weighted mean, expression, quantiles) stay byte-identical + # to numpy because the kernel path routes them through the same reducer. + for name in ("h_mean", "h_sigma", "h_q25", "h_q50", "h_q75"): + np.testing.assert_array_equal( + np.asarray(ks[name], dtype=np.float64), ref[name], err_msg=name + ) + + def test_kernel_empty_cells_get_fill_values(self): + """Cells with no observations get count=0 and NaN floats, like the default path.""" + pa = pytest.importorskip("pyarrow") + cfg = default_config() + children = np.array([1, 2, 3], dtype=np.int64) + # Only cell 2 has data. + cells = np.array([2, 2, 2], dtype=np.int64) + col_dict = { + "h_li": np.array([1.0, 2.0, 3.0], dtype=np.float32), + "s_li": np.array([0.1, 0.1, 0.1], dtype=np.float32), + "leaf_id": cells, + } + kernel = _kernel_aggregate(pa.table(col_dict), cells, children, "h_li", cfg) + ks = kernel["stats_arrays"] + assert kernel["cells_with_data"] == 1 + assert list(ks["count"]) == [0, 3, 0] + assert np.isnan(ks["h_min"][0]) and np.isnan(ks["h_min"][2]) + assert ks["h_min"][1] == 1.0 + + class TestDataSource: """Test data_source section of default config (replaces old DataSourceConfig tests).""" From 82f349627e4230ef2edc8c461a38bb7510efd952 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 15 Jun 2026 10:39:23 +0000 Subject: [PATCH 7/8] fold review: kernel NaN semantics + process_shard tests (#30) --- benchmarks/handoff_bench.py | 5 +- src/zagg/processing.py | 63 ++++++++++--- tests/test_processing.py | 177 ++++++++++++++++++++++++++++++++++++ 3 files changed, 232 insertions(+), 13 deletions(-) diff --git a/benchmarks/handoff_bench.py b/benchmarks/handoff_bench.py index b5448e59..345d6b66 100644 --- a/benchmarks/handoff_bench.py +++ b/benchmarks/handoff_bench.py @@ -11,7 +11,10 @@ The first three feed identical numpy arrays into the reducer, so their stats are asserted byte-for-byte identical. ``arrow-kernel`` reduces via pyarrow's C++ hash kernels, whose float mean/variance differ from numpy by ~1 ULP; it is asserted -*close* (``np.allclose`` at ``KERNEL_RTOL``), not identical. +*close* (``np.allclose`` at ``KERNEL_RTOL``), not identical. Its count/min/max are +asserted *exact*; this synthetic data is NaN-free, so the kernel's NaN-propagation +fix (pyarrow min/max skip NaN, numpy propagate) is covered in the unit tests, not +here. This is the CI-runnable half of #30's benchmark: it isolates the grouping algorithm, the carrier representation cost, and the kernel reducer with no I/O, diff --git a/src/zagg/processing.py b/src/zagg/processing.py index 6ee76f00..d0945eed 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -283,9 +283,13 @@ def calculate_cell_statistics( # pass (Acero ``TableGroupBy.aggregate``). pyarrow's float summation differs from # numpy's, so its ``mean``/``variance`` outputs are NOT byte-identical to the # numpy path; they agree only within ``KERNEL_RTOL`` (validated in tests and in -# ``benchmarks/handoff_bench.py``). This lever is opt-in via -# ``handoff="arrow-kernel"`` and exists purely so phase 3 can benchmark it on -# real ATL03 data; it is kept gated and clearly experimental, and should be +# ``benchmarks/handoff_bench.py``). ``count``/``min``/``max`` ARE exact vs numpy, +# including on NaN input: pyarrow's ``min``/``max`` kernels skip NaN by default +# (numpy propagates it), so :func:`_kernel_aggregate` detects NaN per group and +# overwrites those groups' min/max with NaN to restore numpy parity (NaN is a +# value, not an Arrow null, so ``skip_nulls`` does not cover it). This lever is +# opt-in via ``handoff="arrow-kernel"`` and exists purely so phase 3 can benchmark +# it on real ATL03 data; it is kept gated and clearly experimental, and should be # dropped if that benchmark shows no material speedup (see PR #33 discussion). # Documented tolerance for kernel-vs-numpy float agreement. float32 means/variance @@ -296,6 +300,10 @@ def calculate_cell_statistics( # reductions that are mathematically a pure (unweighted) group reduction appear # here; weighted ``average``, ``quantile`` (only approximate via tdigest) and any # ``expression`` field fall back to the per-cell numpy path. +# NOTE: ``"average" -> "mean"`` is currently dead for the shipped atl06 config +# (its ``h_mean`` is a *weighted* average, which ``_kernel_able`` excludes). It is +# kept only so an unweighted ``average`` field — if a future config defines one — +# is kernel-able rather than silently falling back; remove it if that never lands. _KERNEL_FUNCS = { "len": "count", "count": "count", @@ -334,14 +342,21 @@ def _kernel_aggregate( the remaining (weighted mean, expression, quantile) fields via the per-cell numpy path so output columns match the default reducer exactly in shape. - Kernel-computed float stats are NOT byte-identical to numpy — they agree - within :data:`KERNEL_RTOL`. Returns ``stats_arrays`` (``name -> ndarray`` over - ``children``) plus ``cells_with_data``. + ``count``/``min``/``max`` are EXACT vs the numpy reducer, including on NaN + input — pyarrow's min/max kernels skip NaN, so this function detects NaN per + group and propagates it (numpy semantics). The kernel-reduced float stats + (``mean``/``variance``) are NOT byte-identical to numpy; they agree within + :data:`KERNEL_RTOL` (and both yield NaN on a NaN-bearing group). Returns + ``stats_arrays`` (``name -> ndarray`` over ``children``) plus + ``cells_with_data``. Parameters ---------- table : pyarrow.Table - Concatenated, null-free observations (one row per observation). + Concatenated, null-free observations (one row per observation). "Null-free" + is the Arrow-null sense; float NaN values ARE allowed and are handled with + numpy semantics (see above). Callers must enforce the null-free contract + (``process_shard`` does); this function does not re-check it. cell_col : np.ndarray Child cell id for each row of ``table`` (already ``grid.cells_of`` mapped, so the group key is the destination cell, not the leaf id). @@ -379,21 +394,43 @@ def _kernel_aggregate( (agg_fields[n].get("source") or value_col, _KERNEL_FUNCS[agg_fields[n]["function"]]) for n in kernel_names ] - gd = keyed.group_by("_cell").aggregate(aggregations).to_pydict() + # NaN semantics: pyarrow's ``min``/``max`` hash kernels SKIP NaN, whereas + # ``np.min``/``np.max`` PROPAGATE it. To keep count/min/max bit-identical to the + # numpy path on NaN-bearing input (ATL06 ``h_li`` can carry fill/invalid values + # and ``quality_filter`` is a flag check, not a NaN filter), detect NaN per + # group on each min/max source column and overwrite those groups' min/max with + # NaN below. (``count`` already matches: NaN is a value, not a null, so it is + # counted; ``mean``/``variance`` already propagate NaN like numpy.) + extrema_srcs = { + src for n, (src, kfunc) in zip(kernel_names, aggregations) if kfunc in ("min", "max") + } + for src in extrema_srcs: + is_nan = np.isnan(table.column(src).to_numpy(zero_copy_only=False)) + keyed = keyed.append_column(f"_isnan_{src}", pa.array(is_nan)) + aggregations_nan = [(f"_isnan_{src}", "max") for src in extrema_srcs] + gd = keyed.group_by("_cell").aggregate(aggregations + aggregations_nan).to_pydict() group_cells = gd["_cell"] + group_has_nan = {src: gd[f"_isnan_{src}_max"] for src in extrema_srcs} # Map each grouped row back to its position in ``children``. row_to_idx = [child_index.get(int(c)) for c in group_cells] for n, (src, kfunc) in zip(kernel_names, aggregations): col = gd[f"{src}_{kfunc}"] + nan_flags = group_has_nan.get(src) if kfunc in ("min", "max") else None out = stats_arrays[n] for row, idx in enumerate(row_to_idx): if idx is not None: - out[idx] = col[row] + # Propagate NaN for min/max to match numpy (pyarrow skips NaN). + out[idx] = np.nan if (nan_flags is not None and nan_flags[row]) else col[row] cells_with_data = sum(1 for idx in row_to_idx if idx is not None) # Fallback fields (and only those) via the per-cell numpy reducer. Reuse the # carrier-agnostic grouping so the slices match the default path exactly. + # NOTE: ``calculate_cell_statistics`` recomputes the *full* stats dict per cell, + # so the kernel-able stats are computed a second time here and discarded — we + # only read ``fallback_names`` out of it. Acceptable while experimental (the + # fallback set is small: weighted mean, expression, quantiles); revisit if a + # config makes the fallback set dominate. if fallback_names: col_dict = { name: table.column(name).to_numpy(zero_copy_only=False) for name in table.column_names @@ -541,9 +578,11 @@ def process_shard( :func:`_group_columns` and the same numpy reductions, so scalar outputs are byte-for-byte identical; only the read→concat→extract representation differs. ``"arrow-kernel"`` (phase 2b of #30) instead reduces via - ``pyarrow.compute`` hash-aggregate kernels: its float ``mean``/``variance`` - differ from numpy by ~1 ULP (agree within :data:`KERNEL_RTOL`, not byte - identical). All three are opt-in while benchmarked (issue #30). + ``pyarrow.compute`` hash-aggregate kernels: ``count``/``min``/``max`` stay + exact vs numpy (NaN included — see :func:`_kernel_aggregate`), while its + float ``mean``/``variance`` differ by ~1 ULP (agree within + :data:`KERNEL_RTOL`, not byte identical). All three are opt-in while + benchmarked (issue #30). Returns ------- diff --git a/tests/test_processing.py b/tests/test_processing.py index 3b648dbc..f625f201 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -14,6 +14,7 @@ _kernel_able, _kernel_aggregate, calculate_cell_statistics, + process_shard, write_dataframe_to_zarr, ) @@ -405,6 +406,182 @@ def test_kernel_empty_cells_get_fill_values(self): assert np.isnan(ks["h_min"][0]) and np.isnan(ks["h_min"][2]) assert ks["h_min"][1] == 1.0 + def test_kernel_nan_matches_numpy_semantics(self): + """NaN-bearing cells: count/min/max stay EXACT vs numpy (NaN-propagating). + + pyarrow's min/max kernels skip NaN; numpy's propagate it. _kernel_aggregate + must restore numpy semantics so the "count/min/max exact" contract holds on + the NaN-bearing ``h_li`` values ATL06 can carry (the quality_filter is a + flag check, not a NaN/fill filter). count is unaffected (NaN is a value, not + a null) and mean/variance already propagate NaN like numpy. + """ + pa = pytest.importorskip("pyarrow") + cfg = default_config() + children = np.array([10, 20, 30], dtype=np.int64) + # cell 10: clean; cell 20: one NaN; cell 30: all NaN. + cells = np.array([10, 10, 10, 20, 20, 20, 30, 30], dtype=np.int64) + h = np.array([1.0, 2.0, 4.0, 1.0, np.nan, 3.0, np.nan, np.nan], dtype=np.float32) + s = np.full(len(cells), 0.1, dtype=np.float32) + col_dict = {"h_li": h, "s_li": s, "leaf_id": cells} + + ref = self._numpy_reference(col_dict, cells, children, cfg) + ks = _kernel_aggregate(pa.table(col_dict), cells, children, "h_li", cfg)["stats_arrays"] + + # count: exact everywhere (NaN counts as a value). + np.testing.assert_array_equal(np.asarray(ks["count"], dtype=np.float64), ref["count"]) + # min/max: bit-identical to numpy, including the NaN cells (10 clean, 20/30 + # propagate NaN). assert_array_equal treats NaN==NaN here. + for name in ("h_min", "h_max"): + np.testing.assert_array_equal( + np.asarray(ks[name], dtype=np.float64), ref[name], err_msg=name + ) + # Clean cell 10 is finite; NaN cells 20/30 propagate to NaN. + assert ks["h_min"][0] == 1.0 and ks["h_max"][0] == 4.0 + assert np.isnan(ks["h_min"][1]) and np.isnan(ks["h_max"][1]) + assert np.isnan(ks["h_min"][2]) and np.isnan(ks["h_max"][2]) + # mean/variance already propagate NaN in both paths -> NaN on cells 20/30. + for name in ("h_variance",): + assert np.isnan(ks[name][1]) and np.isnan(ref[name][1]) + assert np.isnan(ks[name][2]) and np.isnan(ref[name][2]) + + +class _KernelShardGrid: + """Minimal grid stub driving the ``process_shard`` kernel branch. + + Exposes only what the ``handoff="arrow-kernel"`` path post-read needs: + ``children``/``cells_of``/``chunk_coords`` (and ``chunk_shape`` is unused by + process_shard itself). Spatial read methods are bypassed because the test + monkeypatches ``_read_group`` to return canned tables. + """ + + def __init__(self, children, leaf_to_cell): + self._children = np.asarray(children, dtype=np.int64) + self._leaf_to_cell = leaf_to_cell + + def children(self, shard_key): + return self._children + + def cells_of(self, leaf_ids): + return np.array([self._leaf_to_cell[int(x)] for x in leaf_ids], dtype=np.int64) + + def chunk_coords(self, shard_key): + return { + "cell_lat": np.zeros(len(self._children)), + "cell_lon": np.zeros(len(self._children)), + } + + +class TestProcessShardKernelBranch: + """HIGH-2 of PR #33 review: exercise the production ``process_shard`` kernel + branch (null guard, ``cells_of``, ``concat_tables().combine_chunks()``, and the + ``handoff`` validation), including NaN-bearing input so the NaN-semantics fix is + covered end-to-end, not only in ``_kernel_aggregate``.""" + + def _patch_reads(self, monkeypatch, tables): + """Make ``_read_group`` yield the canned tables once, then None. + + Also stubs ``h5coro.H5Coro`` so the read loop never touches the network; + the canned tables stand in for the spatially filtered group reads. + """ + it = iter(tables) + + def fake_read_group(*args, **kwargs): + return next(it, None) + + monkeypatch.setattr("zagg.processing._read_group", fake_read_group) + monkeypatch.setattr("zagg.processing.h5coro.H5Coro", lambda *a, **k: object()) + # Avoid resolving a real h5coro driver (s3driver import / creds plumbing). + monkeypatch.setattr("zagg.processing._make_url_rewriter", lambda driver: (lambda u: u)) + + def test_kernel_branch_matches_default_path(self, monkeypatch): + """process_shard(handoff="arrow-kernel") agrees with the default path on the + kernel-able stats (count/min/max exact, variance within KERNEL_RTOL), + running the real concat + null guard + cells_of.""" + pa = pytest.importorskip("pyarrow") + + cfg = default_config() + leaf_to_cell = {1: 10, 2: 10, 3: 20, 4: 30} + children = [10, 20, 30] + grid = _KernelShardGrid(children, leaf_to_cell) + + rng = np.random.default_rng(5) + + def make_table(n): + leaf = rng.choice([1, 2, 3, 4], size=n).astype(np.int64) + h = (rng.standard_normal(n) * 10.0).astype(np.float32) + s = (np.abs(rng.standard_normal(n)) + 0.01).astype(np.float32) + return pa.table({"h_li": h, "s_li": s, "leaf_id": leaf}) + + # Two reads -> exercises pa.concat_tables(...).combine_chunks(). + tables = [make_table(60), make_table(25)] + # Reuse the same data for the default path via a copy of the iterator. + kernel_tables = [t for t in tables] + default_tables = [t for t in tables] + + self._patch_reads(monkeypatch, kernel_tables) + df_k, meta_k = process_shard( + grid, 0, ["s3://x"], s3_credentials={}, config=cfg, handoff="arrow-kernel" + ) + + self._patch_reads(monkeypatch, default_tables) + df_d, meta_d = process_shard( + grid, 0, ["s3://x"], s3_credentials={}, config=cfg, handoff="arrow" + ) + + assert meta_k["cells_with_data"] == meta_d["cells_with_data"] + assert meta_k["total_obs"] == meta_d["total_obs"] == 85 + for name in ("count", "h_min", "h_max"): + np.testing.assert_array_equal( + df_k[name].to_numpy(), df_d[name].to_numpy(), err_msg=name + ) + np.testing.assert_allclose( + df_k["h_variance"].to_numpy(), + df_d["h_variance"].to_numpy(), + rtol=KERNEL_RTOL, + equal_nan=True, + ) + + def test_kernel_branch_nan_input(self, monkeypatch): + """End-to-end NaN handling through process_shard's kernel branch: a NaN in + ``h_li`` propagates to that cell's min/max (numpy semantics), count is + unaffected, and the null guard does NOT trip (NaN is not an Arrow null).""" + pa = pytest.importorskip("pyarrow") + + cfg = default_config() + leaf_to_cell = {1: 10, 2: 20} + children = [10, 20] + grid = _KernelShardGrid(children, leaf_to_cell) + + # cell 10 clean, cell 20 has a NaN. + table = pa.table( + { + "h_li": pa.array([1.0, 2.0, 4.0, 5.0, np.nan], type=pa.float32()), + "s_li": pa.array([0.1, 0.1, 0.1, 0.1, 0.1], type=pa.float32()), + "leaf_id": pa.array([1, 1, 1, 2, 2], type=pa.int64()), + } + ) + self._patch_reads(monkeypatch, [table]) + df, meta = process_shard( + grid, 0, ["s3://x"], s3_credentials={}, config=cfg, handoff="arrow-kernel" + ) + + idx = {c: i for i, c in enumerate(children)} + # Clean cell 10: finite extrema. + assert df["h_min"].to_numpy()[idx[10]] == 1.0 + assert df["h_max"].to_numpy()[idx[10]] == 4.0 + # NaN cell 20: min/max propagate NaN (numpy semantics), count still 2. + assert np.isnan(df["h_min"].to_numpy()[idx[20]]) + assert np.isnan(df["h_max"].to_numpy()[idx[20]]) + assert df["count"].to_numpy()[idx[20]] == 2 + assert meta["total_obs"] == 5 + + def test_invalid_handoff_rejected(self): + """The ``handoff`` validation rejects unknown carriers before any read.""" + + grid = _KernelShardGrid([10], {1: 10}) + with pytest.raises(ValueError, match="handoff must be"): + process_shard(grid, 0, ["s3://x"], s3_credentials={}, handoff="bogus") + class TestDataSource: """Test data_source section of default config (replaces old DataSourceConfig tests).""" From f6b24f112d5e8571a2db97df954100ad10786805 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 16 Jun 2026 00:06:54 +0000 Subject: [PATCH 8/8] document numpy nan-op contract; test nan-aware aggregation (#30) --- src/zagg/processing.py | 32 ++++++++++++++++++++++++++++++++ tests/test_processing.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/src/zagg/processing.py b/src/zagg/processing.py index d0945eed..7b30f610 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -204,6 +204,21 @@ def calculate_cell_statistics( """ Calculate summary statistics for a cell, driven by pipeline config metadata. + User contract + ------------- + The supported aggregation surface is *anything expressible in numpy*. Each + agg field names a ``function`` that :func:`zagg.config.resolve_function` + turns into a callable: a bare name (``"min"``, ``"nanmean"``) resolves to + ``np.`` via ``getattr(np, ...)``, an ``"np."``-prefixed name the same + way, and a dotted path (``"numpy.quantile"``) via import. This means the full + numpy **NaN-aware family** — ``np.nanmean``, ``np.nanvar``, ``np.nanmax``, + ``np.nanmin``, ``np.nansum``, ``np.nanstd``, ``np.nanmedian``, … — is + usable directly from the config template with no special-casing, and is + reduced with numpy's own NaN semantics (see ``test_numpy_nan_aware_functions``). + The experimental ``handoff="arrow-kernel"`` path is an *opt-in* acceleration + for the kernel-able subset only; it does not change or narrow this contract + (see the EXPERIMENTAL block below). + Parameters ---------- cell_data : dict[str, np.ndarray] @@ -276,6 +291,23 @@ def calculate_cell_statistics( # EXPERIMENTAL (phase 2b of #30) ----------------------------------------------- +# Dual aggregation contract +# ------------------------- +# The DEFAULT, fully-supported contract is "any aggregation expressible in numpy", +# including the NaN-aware family (``np.nanmean``/``np.nanvar``/``np.nanmax``/…); +# the user picks the function in the agg template and it runs through +# ``calculate_cell_statistics`` with numpy's own semantics (see that docstring and +# ``test_numpy_nan_aware_functions``). Arrow kernels do NOT replace or narrow that +# contract — they are an OPT-IN acceleration for the kernel-able subset, and the +# user chooses numpy vs arrow per run via the ``handoff`` flag. +# +# Why arrow kernels aren't drop-in nan-operators: pyarrow compute has +# ``mean``/``min_max``/``variance`` with ``skip_nulls``, but an Arrow NULL is a +# distinct missing-value bit, NOT a float NaN — ``skip_nulls`` does not skip NaN. +# So there is no arrow "nanmean" kernel equivalent; the kernel path instead +# replicates numpy's NaN behaviour by hand (NaN-propagating min/max, see below) +# rather than pretending arrow nulls and float NaN are the same thing. +# # Optional pyarrow.compute hash-aggregate ("kernel") reduction path. Unlike the # pandas/arrow *carriers* — which feed identical numpy arrays into # ``calculate_cell_statistics`` and are therefore byte-for-byte identical — the diff --git a/tests/test_processing.py b/tests/test_processing.py index f625f201..88fe881d 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -111,6 +111,39 @@ def test_with_explicit_config(self): np.average([10, 20, 30], weights=1.0 / np.array([0.1, 0.2, 0.1]) ** 2), ) + def test_numpy_nan_aware_functions(self): + """The user contract — any aggregation expressible in numpy, including the + ``nan*`` family — resolves and runs end-to-end through the default numpy + path (``resolve_function`` does ``getattr(np, name)``). NaN-bearing input + must be reduced NaN-skipping, matching the bare ``np.nan*`` operators.""" + from zagg.config import PipelineConfig + + cfg = PipelineConfig( + aggregation={ + "variables": { + "h_nanmean": {"function": "np.nanmean", "source": "h_li", "dtype": "float32"}, + "h_nanmax": {"function": "nanmax", "source": "h_li", "dtype": "float32"}, + "h_nanmin": {"function": "np.nanmin", "source": "h_li", "dtype": "float32"}, + "h_nanvar": {"function": "nanvar", "source": "h_li", "dtype": "float32"}, + "h_nansum": {"function": "np.nansum", "source": "h_li", "dtype": "float32"}, + "h_nanstd": {"function": "np.nanstd", "source": "h_li", "dtype": "float32"}, + } + } + ) + vals = np.array([1.0, np.nan, 3.0, 5.0], dtype=np.float32) + result = calculate_cell_statistics({"h_li": vals}, config=cfg) + + np.testing.assert_allclose(result["h_nanmean"], np.nanmean(vals)) + np.testing.assert_allclose(result["h_nanmax"], np.nanmax(vals)) + np.testing.assert_allclose(result["h_nanmin"], np.nanmin(vals)) + np.testing.assert_allclose(result["h_nanvar"], np.nanvar(vals)) + np.testing.assert_allclose(result["h_nansum"], np.nansum(vals)) + np.testing.assert_allclose(result["h_nanstd"], np.nanstd(vals)) + # NaN-skipping really happened: a plain np.mean/np.max would propagate NaN. + assert not np.isnan(result["h_nanmean"]) + assert result["h_nanmax"] == 5.0 + assert result["h_nanmin"] == 1.0 + class TestBuildGroups: def test_slice_counts_match_per_cell_mask(self):