diff --git a/src/zagg/config.py b/src/zagg/config.py index 3d860e07..c4fe7ecb 100644 --- a/src/zagg/config.py +++ b/src/zagg/config.py @@ -13,6 +13,40 @@ import zagg.configs +class LinkDict(TypedDict): + """Per-level link to the next coarser level (issue #43, Phase B). + + A *link* describes a contiguous-range parent->child tiling: each parent segment + ``p`` covers base-rate indices ``[index_beg[p] - index_base, ...`` for + ``count[p]`` children. ``index_base`` shifts the raw ``index_beg`` values so + that Python 0-based indexing into the base array is straightforward. + + ``reference_index`` is a reserved slot for a future explicit-index-array variant + (non-contiguous children per parent); leave it ``None`` for the contiguous case. + """ + + to: str # key of the coarser level in ``levels`` + index_beg: str # HDF5 path for the per-parent start index array + count: str # HDF5 path for the per-parent child count array + index_base: NotRequired[int] # subtracted from index_beg values (default 0) + reference_index: NotRequired[str | None] # reserved; must be None + + +class LevelDict(TypedDict): + """One hierarchical level in a multi-rate HDF5 source (issue #43, Phase B). + + A source may have several rates (e.g. ATL03 ``photons`` and ``segments``). + Each level declares its own ``path``, ``coordinates``, and ``variables``, + plus an optional ``link`` to a coarser parent level. The flat single-level + form (no ``levels``/``base_level`` keys in ``data_source``) stays first-class. + """ + + path: str # HDF5 group path template (may contain ``{group}``) + coordinates: list[str] # coordinate dataset names within ``path`` + variables: list[str] # variable dataset names within ``path`` + link: NotRequired[LinkDict | None] + + class DataSourceDict(TypedDict): """Type hints for the ``data_source`` section of a pipeline config.""" @@ -21,6 +55,21 @@ class DataSourceDict(TypedDict): coordinates: dict[str, str] variables: dict[str, str] quality_filter: NotRequired[dict] + filters: NotRequired[list[dict]] + # Hierarchical multi-level form (issue #43, Phase B). When present, the flat + # ``coordinates``/``variables`` keys are still accepted for the base level but + # ``levels`` + ``base_level`` take precedence for the read path. + levels: NotRequired[dict[str, LevelDict]] + base_level: NotRequired[str] + + +# Structured-predicate comparison operators (issue #43). ``in``/``not_in`` take a +# ``values`` list; the rest take a scalar ``value``. These are the only +# pushdown-eligible filter language; an ``expression`` filter is a base-level-only, +# aggregation-time escape hatch that forfeits pushdown. +_SCALAR_OPS = frozenset({"eq", "ne", "ge", "le", "lt", "gt"}) +_SET_OPS = frozenset({"in", "not_in"}) +FILTER_OPS = _SCALAR_OPS | _SET_OPS @dataclass @@ -165,6 +214,15 @@ def validate_config(config: PipelineConfig) -> None: if "start_date" not in temporal or "end_date" not in temporal: raise ValueError("bounds.temporal requires start_date and end_date") + # Validate the structured filter list (issue #43, Phase A) + _validate_filters(config.data_source) + + # Validate hierarchical multi-level form (issue #43, Phase B) + _validate_levels(config.data_source) + + # Cross-check: each filter's level field must name a key in levels (issue #43) + _validate_filter_levels(config.data_source) + ds_vars = set(config.data_source.get("variables", {}).keys()) agg_vars = config.aggregation.get("variables", {}) @@ -301,6 +359,251 @@ def _validate_trailing_shape(name: str, trailing_shape) -> None: ) +def get_filters(config: PipelineConfig) -> list[dict]: + """Return the ordered list of normalized data-source filters (issue #43). + + Two filter languages coexist: + + - **Structured predicates** ``{level?, dataset, column?, op, value|values, + keep?}`` are machine-inspectable and are the only kind eligible for read + pushdown (Phase C). ``op`` is one of :data:`FILTER_OPS`; ``in``/``not_in`` + take ``values`` (a list), the rest take a scalar ``value``. ``column`` is an + integer selector into an N-D flag array (e.g. ATL03 ``signal_conf_ph``). + ``keep`` (default ``True``) keeps matching rows; ``keep: false`` drops them. + - **Expression** filters ``{expression: ""}`` are a base-level-only, + aggregation-time escape hatch that forfeits pushdown (opaque to the planner). + + The flat ``quality_filter: {dataset, value}`` is sugar synthesizing one + base-level ``op: eq`` structured filter, so the ATL06 path is unchanged. An + explicit ``filters:`` list, when present, is used as-is (the flat + ``quality_filter`` is then ignored). + + Each returned filter carries a normalized ``level`` (``None`` means the base + level) and, for structured predicates, an explicit ``keep`` bool. + + Parameters + ---------- + config : PipelineConfig + + Returns + ------- + list[dict] + """ + return filters_from_data_source(config.data_source) + + +def filters_from_data_source(data_source: dict) -> list[dict]: + """Normalize the filter list from a raw ``data_source`` dict. + + Shared by :func:`get_filters` and the read path (which only holds the + ``data_source`` mapping). See :func:`get_filters` for the schema. + """ + explicit = data_source.get("filters") + if explicit is not None: + return [_normalize_filter(f) for f in explicit] + qf = data_source.get("quality_filter") + if qf is not None: + return [ + { + "level": None, + "dataset": qf["dataset"], + "column": None, + "op": "eq", + "value": qf["value"], + "keep": True, + } + ] + return [] + + +def _normalize_filter(f: dict) -> dict: + """Normalize one raw filter dict into canonical form (see :func:`get_filters`).""" + if "expression" in f: + return {"level": f.get("level"), "expression": f["expression"]} + op = f["op"] + out = { + "level": f.get("level"), + "dataset": f["dataset"], + "column": f.get("column"), + "op": op, + "keep": bool(f.get("keep", True)), + } + if op in _SET_OPS: + out["values"] = list(f["values"]) + else: + out["value"] = f["value"] + return out + + +def _validate_filters(data_source: dict) -> None: + """Validate the ``filters`` list (and that a flat ``quality_filter`` is sane). + + Raises ``ValueError`` on: unknown op, missing ``dataset``, ``in``/``not_in`` + without a list ``values``, scalar ops without ``value``, non-int ``column``, + a non-base-level ``expression`` filter, or wrong ``value`` type. ``column`` is + required for the N-D flag case but cannot be checked against array rank here + (no data); rank checks happen at read time. + """ + filters = data_source.get("filters") + if filters is None: + return + if not isinstance(filters, list): + raise ValueError("data_source.filters must be a list") + for i, f in enumerate(filters): + if not isinstance(f, dict): + raise ValueError(f"filter[{i}] must be a mapping") + if "expression" in f: + if "op" in f or "dataset" in f: + raise ValueError( + f"filter[{i}]: 'expression' filters take no 'op'/'dataset' " + "(base-level aggregation-time escape hatch, no pushdown)" + ) + if f.get("level") is not None: + raise ValueError( + f"filter[{i}]: 'expression' filters are base-level only (level must be omitted)" + ) + if not isinstance(f["expression"], str): + raise ValueError(f"filter[{i}]: 'expression' must be a string") + continue + if "dataset" not in f: + raise ValueError(f"filter[{i}]: structured filter requires 'dataset'") + op = f.get("op") + if op not in FILTER_OPS: + raise ValueError(f"filter[{i}]: unknown op {op!r} (allowed: {sorted(FILTER_OPS)})") + col = f.get("column") + if col is not None and (not isinstance(col, int) or isinstance(col, bool)): + raise ValueError(f"filter[{i}]: 'column' must be an integer index (got {col!r})") + if op in _SET_OPS: + if not isinstance(f.get("values"), list): + raise ValueError(f"filter[{i}]: op {op!r} requires a 'values' list") + for v in f["values"]: + if not isinstance(v, (int, float)) or isinstance(v, bool): + raise ValueError(f"filter[{i}]: 'values' must be numeric (got {v!r})") + else: + if "value" not in f: + raise ValueError(f"filter[{i}]: op {op!r} requires a scalar 'value'") + if not isinstance(f["value"], (int, float)) or isinstance(f["value"], bool): + raise ValueError(f"filter[{i}]: 'value' must be numeric (got {f['value']!r})") + + +def _validate_levels(data_source: dict) -> None: + """Validate the hierarchical ``levels``/``base_level`` form (issue #43, Phase B). + + Rules: + - ``base_level`` must name a key in ``levels``. + - ``link.to`` in each level must name another key in ``levels``. + - ``link.index_base`` must be a non-negative int when present. + - ``link.reference_index`` must be ``None`` when present (reserved slot). + - Only ``base_level`` may omit ``link`` (it has no coarser parent). + - Flat single-level form (no ``levels`` key) is always valid. + """ + levels = data_source.get("levels") + if levels is None: + return + if not isinstance(levels, dict) or not levels: + raise ValueError("data_source.levels must be a non-empty mapping") + base_level = data_source.get("base_level") + if base_level is None: + raise ValueError("data_source.base_level is required when levels is present") + if base_level not in levels: + raise ValueError( + f"data_source.base_level {base_level!r} is not a key in levels " + f"(available: {sorted(levels)})" + ) + level_keys = set(levels) + for name, lvl in levels.items(): + if not isinstance(lvl, dict): + raise ValueError(f"levels.{name} must be a mapping") + if "path" not in lvl: + raise ValueError(f"levels.{name}: 'path' is required") + link = lvl.get("link") + if link is None: + if name != base_level: + raise ValueError( + f"levels.{name}: non-base levels must have a 'link' " + f"(only {base_level!r} may omit it)" + ) + continue + if not isinstance(link, dict): + raise ValueError(f"levels.{name}.link must be a mapping") + for field_name in ("to", "index_beg", "count"): + if field_name not in link: + raise ValueError(f"levels.{name}.link: '{field_name}' is required") + unknown = set(link) - {"to", "index_beg", "count", "index_base", "reference_index"} + if unknown: + raise ValueError( + f"levels.{name}.link: unknown fields {sorted(unknown)} " + f"(allowed: to, index_beg, count, index_base, reference_index)" + ) + if link.get("to") == name: + raise ValueError(f"level '{name}': link.to cannot reference the level itself") + if link["to"] not in level_keys: + raise ValueError( + f"levels.{name}.link.to {link['to']!r} is not a key in levels " + f"(available: {sorted(level_keys)})" + ) + index_base = link.get("index_base", 0) + if not isinstance(index_base, int) or isinstance(index_base, bool) or index_base < 0: + raise ValueError( + f"levels.{name}.link.index_base must be a non-negative int (got {index_base!r})" + ) + ref = link.get("reference_index") + if ref is not None: + raise ValueError( + f"levels.{name}.link.reference_index is reserved and must be null/omitted " + f"(explicit index-array variant not yet implemented)" + ) + + +def _validate_filter_levels(data_source: dict) -> None: + """Cross-check each filter's level field against the levels keys (issue #43). + + A filter with ``level: "nonexistent"`` would otherwise only fail at read time + with an opaque ``KeyError``. Raises ``ValueError`` with a clear message when a + filter's ``level`` names a key not present in ``levels``. + """ + levels = data_source.get("levels") + if levels is None: + return + level_keys = set(levels) + filters = data_source.get("filters") or [] + for i, f in enumerate(filters): + lvl = f.get("level") + if lvl is not None and lvl not in level_keys: + raise ValueError( + f"filter[{i}]: level {lvl!r} is not a key in levels " + f"(available: {sorted(level_keys)})" + ) + + +def get_levels(config: "PipelineConfig") -> dict | None: + """Return the ``levels`` mapping from the data source, or ``None`` if flat. + + Parameters + ---------- + config : PipelineConfig + + Returns + ------- + dict or None + """ + return config.data_source.get("levels") + + +def get_base_level(config: "PipelineConfig") -> str | None: + """Return the ``base_level`` key from the data source, or ``None`` if flat. + + Parameters + ---------- + config : PipelineConfig + + Returns + ------- + str or None + """ + return config.data_source.get("base_level") + + def _is_numeric(s: str) -> bool: """Check if a string is a numeric literal.""" try: @@ -660,3 +963,36 @@ def evaluate_expression(expression: str, columns: dict[str, np.ndarray]) -> floa float """ return float(_eval_expression_raw(expression, columns)) + + +def evaluate_filter_expression(expression: str, columns: dict[str, np.ndarray]) -> np.ndarray: + """Evaluate a boolean filter expression to a per-row mask (issue #43). + + Like :func:`evaluate_expression` but returns the raw boolean array rather than + a scalar float — the base-level ``expression`` filter escape hatch (e.g. + ``"(h_li > 0) & (s_li < 1)"``). Uses the same restricted namespace. + + Parameters + ---------- + expression : str + Python boolean expression over numpy and column variables. + columns : dict[str, np.ndarray] + Mapping of column names to arrays. + + Returns + ------- + numpy.ndarray + Boolean mask. + """ + ns = { + "__builtins__": {}, + "np": np, + "numpy": np, + "len": len, + "float": float, + "int": int, + "abs": abs, + "sum": sum, + **columns, + } + return np.asarray(eval(expression, ns), dtype=bool) # noqa: S307 diff --git a/src/zagg/processing.py b/src/zagg/processing.py index c526cdd8..61dd1822 100644 --- a/src/zagg/processing.py +++ b/src/zagg/processing.py @@ -19,6 +19,8 @@ from zagg.config import ( PipelineConfig, default_config, + evaluate_filter_expression, + filters_from_data_source, get_agg_fields, get_data_vars, get_output_signature, @@ -166,8 +168,7 @@ def _has_vector_fields(config: PipelineConfig) -> bool: :func:`_arrow_column`). """ return any( - get_output_signature(meta)["kind"] == "vector" - for meta in get_agg_fields(config).values() + get_output_signature(meta)["kind"] == "vector" for meta in get_agg_fields(config).values() ) @@ -657,18 +658,123 @@ def _kernel_aggregate( # -- end EXPERIMENTAL kernel path --------------------------------------------- -def _read_group( - h5obj, group: str, data_source: dict, shard_key: int, grid, arrow: bool = False -): +_COMPARE = { + "eq": np.equal, + "ne": np.not_equal, + "ge": np.greater_equal, + "le": np.less_equal, + "lt": np.less, + "gt": np.greater, +} + + +def _expand_mask_to_base( + coarse_mask: np.ndarray, + index_beg_arr: np.ndarray, + count_arr: np.ndarray, + index_base: int, + total_base_size: int, +) -> np.ndarray: + """Expand a coarse-rate boolean mask to a base-rate boolean mask (issue #43, Phase B). + + Each coarse parent ``p`` covers base-rate rows + ``index_beg_arr[p] - index_base, ..., index_beg_arr[p] - index_base + count_arr[p] - 1``. + The contiguity assumption: ranges do not overlap and together tile the full base array. + + Parameters + ---------- + coarse_mask : np.ndarray + 1-D boolean array of length ``n_parents``. + index_beg_arr : np.ndarray + Per-parent start index into the base array (before ``index_base`` shift). + count_arr : np.ndarray + Per-parent child count (number of base-rate rows this parent covers). + index_base : int + Subtracted from ``index_beg_arr`` to get 0-based base indices. + total_base_size : int + Length of the output base-rate array. + + Returns + ------- + np.ndarray + 1-D boolean array of length ``total_base_size``. + """ + out = np.zeros(total_base_size, dtype=bool) + for p, keep in enumerate(coarse_mask): + if not keep: + continue + beg = int(index_beg_arr[p]) - index_base + if beg < 0: + raise ValueError( + f"index_beg_arr[{p}]={index_beg_arr[p]} is less than index_base={index_base}" + ) + cnt = int(count_arr[p]) + out[beg : beg + cnt] = True + return out + + +def _predicate_mask(arr: np.ndarray, f: dict) -> np.ndarray: + """Build a 1-D boolean keep-mask for one structured predicate (issue #43). + + ``f`` is a normalized structured filter (see :func:`zagg.config.get_filters`): + ``{op, column, value|values, keep}``. An integer ``column`` selects a column + from a 2-D flag array before comparing; it is required for N-D arrays and + rejected for 1-D arrays. ``keep: false`` inverts the result (drop matches). + """ + column = f.get("column") + if arr.ndim > 1: + if column is None: + raise ValueError(f"filter on '{f['dataset']}': N-D array requires an integer 'column'") + arr = arr[:, column] + elif column is not None: + raise ValueError(f"filter on '{f['dataset']}': 'column' set but array is 1-D") + + op = f["op"] + if op == "in": + mask = np.isin(arr, f["values"]) + elif op == "not_in": + mask = ~np.isin(arr, f["values"]) + else: + mask = _COMPARE[op](arr, f["value"]) + if not f.get("keep", True): + mask = ~mask + return mask + + +def _read_group(h5obj, group: str, data_source: dict, shard_key: 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. + + Supports two modes (issue #43, Phase B): + + *Flat* (no ``levels``/``base_level`` in ``data_source``): unchanged from Phase A — + all structured filters are applied directly to base-rate data. + + *Hierarchical* (``levels`` + ``base_level`` present): structured filters whose + normalized ``level`` key names a non-base level are applied at coarse rate, then + expanded to base-rate via the level's ``link`` arrays (``_expand_mask_to_base``). + Base-level structured filters and expression filters are unchanged. """ coordinates = data_source["coordinates"] variables = data_source["variables"] - quality_filter = data_source.get("quality_filter") + filters = filters_from_data_source(data_source) + base_level_key = data_source.get("base_level") + levels = data_source.get("levels") + # Partition filters: base-level structured, coarse-level structured, expressions. + base_structured = [ + f + for f in filters + if "expression" not in f and (f.get("level") is None or f.get("level") == base_level_key) + ] + coarse_structured = [ + f + for f in filters + if "expression" not in f and f.get("level") is not None and f.get("level") != base_level_key + ] + expressions = [f for f in filters if "expression" in f] # Resolve coordinate paths coord_paths = [path.format(group=group) for path in coordinates.values()] @@ -694,46 +800,107 @@ def _read_group( min_idx = int(indices[0]) max_idx = int(indices[-1]) + 1 - # Build hyperslice dataset list: variables + optional quality filter - datasets = [] - for path_template in variables.values(): - path = path_template.format(group=group) - datasets.append({"dataset": path, "hyperslice": [(min_idx, max_idx)]}) + # --- Coarse-level filter expansion (Phase B) --- + # For each filter whose level is not the base level, read the coarse-rate + # flag array from the declared level path, build a coarse mask, then expand + # to base-rate via the level link arrays. AND the results into ``cross_mask``. + cross_mask: np.ndarray | None = None + if coarse_structured and levels is not None: + for f in coarse_structured: + level_key = f["level"] + lvl = levels[level_key] + flag_path = f["dataset"].format(group=group) + # Read the coarse flag array (full level, no hyperslice — we need all parents + # to align with link arrays which are also full-length). + coarse_data = h5obj.readDatasets([{"dataset": flag_path}]) + coarse_arr = coarse_data[flag_path] + coarse_fmask = _predicate_mask(coarse_arr, f) + # Read the link arrays from this level. + link = lvl["link"] + index_base = int(link.get("index_base", 0)) + ibeg_path = link["index_beg"].format(group=group) + cnt_path = link["count"].format(group=group) + link_data = h5obj.readDatasets( + [ + {"dataset": ibeg_path}, + {"dataset": cnt_path}, + ] + ) + ibeg_arr = link_data[ibeg_path] + cnt_arr = link_data[cnt_path] + expanded = _expand_mask_to_base(coarse_fmask, ibeg_arr, cnt_arr, index_base, len(lats)) + cross_mask = expanded if cross_mask is None else (cross_mask & expanded) + if cross_mask is not None and np.sum(cross_mask[min_idx:max_idx]) == 0: + return None - if quality_filter is not None: - qf_path = quality_filter["dataset"].format(group=group) - datasets.append({"dataset": qf_path, "hyperslice": [(min_idx, max_idx)]}) + # Build hyperslice dataset list: variables + any base-level structured-filter arrays. + # Read each distinct path once; flag datasets may coincide with a variable. + datasets = [] + paths_seen = set() + var_paths = {col: tmpl.format(group=group) for col, tmpl in variables.items()} + for path in var_paths.values(): + if path not in paths_seen: + datasets.append({"dataset": path, "hyperslice": [(min_idx, max_idx)]}) + paths_seen.add(path) + filter_paths = {id(f): f["dataset"].format(group=group) for f in base_structured} + for path in filter_paths.values(): + if path not in paths_seen: + datasets.append({"dataset": path, "hyperslice": [(min_idx, max_idx)]}) + paths_seen.add(path) data = h5obj.readDatasets(datasets) # Apply spatial mask to sliced data mask_sliced = mask_spatial[min_idx:max_idx] - # Apply quality filter if configured - if quality_filter is not None: - qf_path = quality_filter["dataset"].format(group=group) - q_flag = data[qf_path][mask_sliced] - quality_mask = q_flag == quality_filter["value"] - if np.sum(quality_mask) == 0: - return None - else: - quality_mask = None + # Combine base-level structured predicates as ANDed keep-masks (issue #43). + keep_mask = None + for f in base_structured: + flag = data[filter_paths[id(f)]][mask_sliced] + fmask = _predicate_mask(flag, f) + keep_mask = fmask if keep_mask is None else (keep_mask & fmask) + + # AND in the cross-level expanded mask, aligned to the sliced window. + if cross_mask is not None: + cross_sliced = cross_mask[min_idx:max_idx][mask_sliced] + keep_mask = cross_sliced if keep_mask is None else (keep_mask & cross_sliced) + + if keep_mask is not None and np.sum(keep_mask) == 0: + return None - # Build dataframe + # Build dataframe (variables sliced to spatial mask, then to the keep-mask) leaf_sliced = leaf_ids[min_idx:max_idx][mask_sliced] data_dict = {} - for col_name, path_template in variables.items(): - path = path_template.format(group=group) + for col_name, path in var_paths.items(): values = data[path][mask_sliced] - if quality_mask is not None: - values = values[quality_mask] + if keep_mask is not None: + values = values[keep_mask] data_dict[col_name] = values - if quality_mask is not None: - data_dict["leaf_id"] = leaf_sliced[quality_mask] + if keep_mask is not None: + data_dict["leaf_id"] = leaf_sliced[keep_mask] else: data_dict["leaf_id"] = leaf_sliced + # Base-level ``expression`` filters: aggregation-time escape hatch, evaluated + # over the already-read variable columns (forfeits pushdown, issue #43). + for f in expressions: + cols = {c: data_dict[c] for c in variables if c in data_dict} + try: + emask = evaluate_filter_expression(f["expression"], cols) + except NameError as e: + raise NameError( + f"expression filter {f['expression']!r} references an undefined name: {e}" + ) from e + if emask.shape != data_dict["leaf_id"].shape: + raise ValueError( + f"expression filter {f['expression']!r} must yield a per-row " + f"boolean mask (got shape {emask.shape})" + ) + if np.sum(emask) == 0: + return None + data_dict = {k: v[emask] for k, v in data_dict.items()} + if arrow: import pyarrow as pa diff --git a/src/zagg/read_plan.py b/src/zagg/read_plan.py new file mode 100644 index 00000000..6756dd59 --- /dev/null +++ b/src/zagg/read_plan.py @@ -0,0 +1,210 @@ +"""Offline-computable read plan for AOI-based hyperslice selection (issue #43, Phase C). + +``plan_read`` computes which coarse-level segments (e.g. ATL03 ``land_ice_segments``) +overlap a bounding-box AOI, merges adjacent runs, translates them to base-level +(photon) slices, and packages them as h5coro-compatible hyperslice lists. +``execute_read_plan`` then drives a caller-supplied read function with those slices. + +Both functions are pure and offline-testable — no h5coro, S3, or credentials needed. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np +from shapely.geometry import LineString, box + + +@dataclass +class ReadPlan: + """A pre-computed read plan for efficient AOI-scoped HDF5 reads. + + Attributes + ---------- + parent_runs : list of (int, int) + Contiguous runs of matched coarse-level parents, as ``(start, end)`` + inclusive index pairs in the coarse array. + base_slices : list of (int, int) + Corresponding base-level ``[start, end)`` half-open slices. + chunk_lists : list of list of (int, int) + h5coro-style ``[(start, end)]`` hyperslice lists (inclusive end, 0-based) + for each run. + coarse_flag_ranges : list of (int, int) + Reserved for future use (coarse-level flag read ranges). + full_read : bool + True when the plan falls back to reading the full array (selectivity too low). + """ + + parent_runs: list[tuple[int, int]] + base_slices: list[tuple[int, int]] + chunk_lists: list[list[tuple[int, int]]] + coarse_flag_ranges: list[tuple[int, int]] = field(default_factory=list) + full_read: bool = False + + +def plan_read( + lat_arr: np.ndarray, + lon_arr: np.ndarray, + index_beg_arr: np.ndarray, + count_arr: np.ndarray, + n_base: int, + bbox: tuple[float, float, float, float], + index_base: int = 0, + pad: int = 1, + full_read_threshold: float = 0.9, +) -> ReadPlan: + """Compute a hyperslice-based read plan for an AOI bounding box (issue #43, Phase C). + + For each coarse-level parent (segment), checks whether: + (a) the rep-point ``(lat_arr[j], lon_arr[j])`` falls within ``bbox``, OR + (b) the linestring from ``(lat_arr[j], lon_arr[j])`` to the next parent + crosses ``bbox`` (euclidean approximation, fine for 20 m segments). + + Adjacent matched parents are merged into contiguous runs, padded by ``pad`` + elements on each side, then translated to base-level ``[start, end)`` slices. + + If the total planned base-level read exceeds ``full_read_threshold * n_base``, + returns a plan with ``full_read=True`` covering everything (cheaper than many + small reads that still sum to most of the file). + + Parameters + ---------- + lat_arr : np.ndarray + Float array, shape ``(n_coarse,)``. Rep-point latitudes of each parent. + lon_arr : np.ndarray + Float array, shape ``(n_coarse,)``. Rep-point longitudes of each parent. + index_beg_arr : np.ndarray + Integer array, shape ``(n_coarse,)``. Base-level start for each parent + (before ``index_base`` adjustment). + count_arr : np.ndarray + Integer array, shape ``(n_coarse,)``. Number of base-level children per parent. + n_base : int + Total size of the base array. + bbox : (min_lon, min_lat, max_lon, max_lat) + Bounding box for the AOI. + index_base : int + 0 (default) or 1 (ATL03 1-based ``ph_index_beg``). + pad : int + Number of extra parents to include on each side of each run (clamped + to array bounds). Helps capture partial edge segments. + full_read_threshold : float + Fraction of ``n_base`` above which the plan falls back to a full read. + + Returns + ------- + ReadPlan + """ + n_coarse = len(lat_arr) + if n_coarse == 0 or n_base == 0: + return ReadPlan(parent_runs=[], base_slices=[], chunk_lists=[]) + + aoi_box = box(bbox[0], bbox[1], bbox[2], bbox[3]) + + # -- AOI matching -- + in_aoi = np.zeros(n_coarse, dtype=bool) + for j in range(n_coarse): + lat, lon = float(lat_arr[j]), float(lon_arr[j]) + if bbox[0] <= lon <= bbox[2] and bbox[1] <= lat <= bbox[3]: + in_aoi[j] = True + continue + # Last segment: no next-segment linestring to check; rep-point only. + if j + 1 < n_coarse: + lat2, lon2 = float(lat_arr[j + 1]), float(lon_arr[j + 1]) + seg_line = LineString([(lon, lat), (lon2, lat2)]) + if seg_line.intersects(aoi_box): + in_aoi[j] = True + + if not in_aoi.any(): + return ReadPlan(parent_runs=[], base_slices=[], chunk_lists=[]) + + # -- Run merging -- + # Find contiguous blocks of True in in_aoi. + raw_runs: list[tuple[int, int]] = [] + start_run = None + for j in range(n_coarse): + if in_aoi[j] and start_run is None: + start_run = j + elif not in_aoi[j] and start_run is not None: + raw_runs.append((start_run, j - 1)) + start_run = None + if start_run is not None: + raw_runs.append((start_run, n_coarse - 1)) + + # -- Padding -- + padded_runs: list[tuple[int, int]] = [] + for s, e in raw_runs: + ps = max(0, s - pad) + pe = min(n_coarse - 1, e + pad) + padded_runs.append((ps, pe)) + + # Merge overlapping/adjacent padded runs. + merged: list[tuple[int, int]] = [] + for s, e in padded_runs: + # +1 merges immediately adjacent runs (closed intervals: [a,b] and [b+1,c] -> [a,c]). + if merged and s <= merged[-1][1] + 1: + merged[-1] = (merged[-1][0], max(merged[-1][1], e)) + else: + merged.append((s, e)) + + # -- Translate to base slices -- + base_slices: list[tuple[int, int]] = [] + chunk_lists: list[list[tuple[int, int]]] = [] + total_base = 0 + for s, e in merged: + base_start = int(index_beg_arr[s]) - index_base + base_end = int(index_beg_arr[e]) - index_base + int(count_arr[e]) + base_start = max(0, base_start) + base_end = min(n_base, base_end) + if base_end <= base_start: + continue + base_slices.append((base_start, base_end)) + chunk_lists.append([(base_start, base_end - 1)]) # h5coro inclusive end + total_base += base_end - base_start + + # -- Selectivity fallback -- + if total_base > full_read_threshold * n_base: + return ReadPlan( + parent_runs=[(0, n_coarse - 1)], + base_slices=[(0, n_base)], + chunk_lists=[[(0, n_base - 1)]], + full_read=True, + ) + + return ReadPlan( + parent_runs=merged, + base_slices=base_slices, + chunk_lists=chunk_lists, + ) + + +def execute_read_plan(plan: ReadPlan, read_fn, dataset_path: str, dtype) -> np.ndarray: + """Execute a ReadPlan using the supplied read function. + + Parameters + ---------- + plan : ReadPlan + A plan produced by :func:`plan_read`. + read_fn : callable + ``read_fn(dataset_path, hyperslice=None)`` -> array. + Passing ``hyperslice=None`` reads the full dataset. + dataset_path : str + HDF5 dataset path to read. + dtype : numpy dtype or str + Dtype for the returned array. + + Returns + ------- + np.ndarray + Concatenated data for all runs in the plan, or an empty array if the + plan matched nothing. If ``plan.full_read``, returns the full dataset. + """ + if plan.full_read: + return np.asarray(read_fn(dataset_path, hyperslice=None), dtype=dtype) + if not plan.parent_runs: + return np.empty(0, dtype=dtype) + parts = [ + np.asarray(read_fn(dataset_path, hyperslice=chunk), dtype=dtype) + for chunk in plan.chunk_lists + ] + return np.concatenate(parts) diff --git a/tests/test_config.py b/tests/test_config.py index da1fdc9d..b3c6396e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -11,9 +11,12 @@ default_config, evaluate_expression, get_agg_fields, + get_base_level, get_child_order, get_coords, get_data_vars, + get_filters, + get_levels, get_output_signature, get_store_path, load_config, @@ -250,6 +253,319 @@ def test_missing_required_section(self): validate_config(cfg) +# --------------------------------------------------------------------------- +# Structured filters (issue #43, Phase A) +# --------------------------------------------------------------------------- + + +def _cfg_with_filters(filters=None, quality_filter=None): + """Minimal valid config with a custom data_source filter spec.""" + ds = {"variables": {"h_li": "/{group}/h_li"}} + if filters is not None: + ds["filters"] = filters + if quality_filter is not None: + ds["quality_filter"] = quality_filter + return PipelineConfig( + data_source=ds, + aggregation={ + "variables": { + "h_min": {"function": "min", "source": "h_li", "dtype": "float32"}, + } + }, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + + +class TestFilters: + def test_quality_filter_synthesizes_base_eq(self, atl06_config): + filters = get_filters(atl06_config) + assert len(filters) == 1 + f = filters[0] + assert f["op"] == "eq" + assert f["level"] is None + assert f["column"] is None + assert f["keep"] is True + assert f["value"] == 0 + assert f["dataset"].endswith("atl06_quality_summary") + + def test_no_filters_returns_empty(self): + cfg = _cfg_with_filters() + assert get_filters(cfg) == [] + + def test_explicit_filters_override_quality_filter(self): + cfg = _cfg_with_filters( + filters=[{"dataset": "/{group}/conf", "column": 0, "op": "ne", "value": 0}], + quality_filter={"dataset": "/{group}/qs", "value": 0}, + ) + filters = get_filters(cfg) + assert len(filters) == 1 + assert filters[0]["column"] == 0 + assert filters[0]["op"] == "ne" + + def test_normalize_set_op_keeps_values_list(self): + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "in", "values": [2, 3, 4]}]) + f = get_filters(cfg)[0] + assert f["values"] == [2, 3, 4] + assert "value" not in f + + def test_normalize_keep_drop(self): + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "eq", "value": 1, "keep": False}]) + assert get_filters(cfg)[0]["keep"] is False + + def test_expression_filter_normalized(self): + cfg = _cfg_with_filters(filters=[{"expression": "h_li > 0"}]) + f = get_filters(cfg)[0] + assert f["expression"] == "h_li > 0" + assert f["level"] is None + + def test_unknown_op_rejected(self): + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "between", "value": 1}]) + with pytest.raises(ValueError, match="unknown op"): + validate_config(cfg) + + def test_column_must_be_int(self): + cfg = _cfg_with_filters( + filters=[{"dataset": "/d", "column": "land", "op": "ne", "value": 0}] + ) + with pytest.raises(ValueError, match="must be an integer"): + validate_config(cfg) + + def test_set_op_requires_values_list(self): + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "in", "value": 3}]) + with pytest.raises(ValueError, match="requires a 'values' list"): + validate_config(cfg) + + def test_scalar_op_requires_value(self): + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "eq"}]) + with pytest.raises(ValueError, match="requires a scalar 'value'"): + validate_config(cfg) + + def test_bad_value_type_rejected(self): + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "eq", "value": "x"}]) + with pytest.raises(ValueError, match="must be numeric"): + validate_config(cfg) + + def test_missing_dataset_rejected(self): + cfg = _cfg_with_filters(filters=[{"op": "eq", "value": 0}]) + with pytest.raises(ValueError, match="requires 'dataset'"): + validate_config(cfg) + + def test_expression_with_level_rejected(self): + cfg = _cfg_with_filters(filters=[{"expression": "h_li > 0", "level": "segment"}]) + with pytest.raises(ValueError, match="base-level only"): + validate_config(cfg) + + def test_expression_with_op_rejected(self): + cfg = _cfg_with_filters(filters=[{"expression": "h_li > 0", "op": "eq", "dataset": "/d"}]) + with pytest.raises(ValueError, match="take no 'op'"): + validate_config(cfg) + + def test_bool_column_rejected(self): + # bool is a subclass of int; filter column: true must be rejected. + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "column": True, "op": "eq", "value": 0}]) + with pytest.raises(ValueError, match="must be an integer"): + validate_config(cfg) + + def test_bool_value_rejected(self): + # bool is a subclass of int; filter value: true must be rejected. + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "eq", "value": True}]) + with pytest.raises(ValueError, match="must be numeric"): + validate_config(cfg) + + def test_bool_in_values_rejected(self): + # bool elements in a 'values' list must be rejected. + cfg = _cfg_with_filters(filters=[{"dataset": "/d", "op": "in", "values": [0, True]}]) + with pytest.raises(ValueError, match="must be numeric"): + validate_config(cfg) + + +# --------------------------------------------------------------------------- +# Hierarchical levels and link validation (issue #43, Phase B) +# --------------------------------------------------------------------------- + + +def _minimal_two_level_ds(**overrides): + """Return a minimal two-level data_source dict with one segment->photon link.""" + ds = { + "reader": "h5coro", + "groups": ["gt1l"], + "coordinates": {"latitude": "/gt1l/ph_lat", "longitude": "/gt1l/ph_lon"}, + "variables": {"h": "/gt1l/h_ph"}, + "base_level": "photons", + "levels": { + "photons": { + "path": "/{group}/heights", + "coordinates": ["lat_ph", "lon_ph"], + "variables": ["h_ph"], + "link": None, + }, + "segments": { + "path": "/{group}/geolocation", + "coordinates": ["reference_photon_lat", "reference_photon_lon"], + "variables": ["signal_conf_ph"], + "link": { + "to": "photons", + "index_beg": "/{group}/geolocation/ph_index_beg", + "count": "/{group}/geolocation/segment_ph_cnt", + }, + }, + }, + } + ds.update(overrides) + return ds + + +def _cfg_with_levels(**overrides): + ds = _minimal_two_level_ds(**overrides) + return PipelineConfig( + data_source=ds, + aggregation={"variables": {"count": {"function": "len", "source": "h", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + + +class TestLevelsValidation: + def test_valid_two_level_config(self): + validate_config(_cfg_with_levels()) + + def test_flat_form_still_valid(self, atl06_config): + # Flat form (no levels/base_level) must still pass. + assert get_levels(atl06_config) is None + assert get_base_level(atl06_config) is None + validate_config(atl06_config) + + def test_get_levels_and_base_level(self): + cfg = _cfg_with_levels() + levels = get_levels(cfg) + assert levels is not None + assert "photons" in levels + assert "segments" in levels + assert get_base_level(cfg) == "photons" + + def test_base_level_must_name_a_key(self): + cfg = _cfg_with_levels(base_level="nonexistent") + with pytest.raises(ValueError, match="not a key in levels"): + validate_config(cfg) + + def test_link_to_must_name_a_key(self): + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"]["to"] = "nonexistent" + cfg = _cfg_with_levels(**ds) + with pytest.raises(ValueError, match="not a key in levels"): + validate_config(cfg) + + def test_link_to_must_name_a_key2(self): + # Build config directly to avoid _cfg_with_levels merging issues + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"]["to"] = "nonexistent" + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="not a key in levels"): + validate_config(cfg) + + def test_link_missing_required_field(self): + ds = _minimal_two_level_ds() + del ds["levels"]["segments"]["link"]["count"] + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="'count' is required"): + validate_config(cfg) + + def test_link_unknown_field_rejected(self): + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"]["bogus"] = "x" + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="unknown fields"): + validate_config(cfg) + + def test_base_level_without_link_ok(self): + # base_level is the only level allowed to have link: None. + cfg = _cfg_with_levels() + validate_config(cfg) + assert cfg.data_source["levels"]["photons"]["link"] is None + + def test_non_base_level_without_link_rejected(self): + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"] = None + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="must have a 'link'"): + validate_config(cfg) + + def test_levels_missing_base_level_key_rejected(self): + ds = _minimal_two_level_ds() + del ds["base_level"] + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="base_level is required"): + validate_config(cfg) + + def test_index_base_must_be_nonneg_int(self): + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"]["index_base"] = -1 + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="non-negative int"): + validate_config(cfg) + + def test_reference_index_must_be_none(self): + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"]["reference_index"] = "/some/path" + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="reserved"): + validate_config(cfg) + + + def test_self_link_rejected(self): + # link.to == level name (self-reference) must raise ValueError + ds = _minimal_two_level_ds() + ds["levels"]["segments"]["link"]["to"] = "segments" + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="cannot reference the level itself"): + validate_config(cfg) + + def test_filter_level_not_in_levels_rejected(self): + # A filter whose level names a nonexistent key must fail at validate time. + ds = _minimal_two_level_ds() + ds["filters"] = [ + {"level": "nonexistent", "dataset": "/{group}/flag", "op": "eq", "value": 0} + ] + cfg = PipelineConfig( + data_source=ds, + aggregation={"variables": {"c": {"function": "len", "dtype": "int32"}}}, + output={"grid": {"type": "healpix", "parent_order": 6, "child_order": 12}}, + ) + with pytest.raises(ValueError, match="not a key in levels"): + validate_config(cfg) + + # --------------------------------------------------------------------------- # Helper accessors # --------------------------------------------------------------------------- diff --git a/tests/test_processing.py b/tests/test_processing.py index d4af2392..bfb12a6b 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -12,11 +12,14 @@ _build_groups, _build_output, _concat_and_group, + _expand_mask_to_base, _group_columns, _has_vector_fields, _iter_carrier_columns, _kernel_able, _kernel_aggregate, + _predicate_mask, + _read_group, calculate_cell_statistics, process_shard, write_dataframe_to_zarr, @@ -649,7 +652,7 @@ def fake_read_group(*args, **kwargs): 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)) + 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 @@ -782,7 +785,7 @@ def _vector_cfg(): "params": {"minlength": 3}, }, } - } + }, ) def test_has_vector_fields(self): @@ -919,9 +922,7 @@ def _vector_cfg(): } from zagg.config import PipelineConfig - return PipelineConfig( - data_source=cfg.data_source, aggregation=agg, output=cfg.output - ) + return PipelineConfig(data_source=cfg.data_source, aggregation=agg, output=cfg.output) def test_vector_leaf_to_zarr_to_read(self): pytest.importorskip("pyarrow") @@ -998,11 +999,461 @@ class _OneChunkGrid: dtype="float32", fill_value=np.float32("nan"), ) - edges = pa.FixedSizeListArray.from_arrays( - pa.array(np.arange(8.0, dtype="float32")), 4 - ) + edges = pa.FixedSizeListArray.from_arrays(pa.array(np.arange(8.0, dtype="float32")), 4) table = pa.table({"edges": edges}) with pytest.raises(ValueError, match="one whole chunk"): - write_dataframe_to_zarr( - table, store, grid=_OneChunkGrid(), chunk_idx=(0,) - ) + write_dataframe_to_zarr(table, store, grid=_OneChunkGrid(), chunk_idx=(0,)) + + +# --------------------------------------------------------------------------- +# Structured filters in the read path (issue #43, Phase A) +# --------------------------------------------------------------------------- + + +class _FakeH5: + """Stub h5coro object: ``readDatasets`` returns canned arrays by path. + + Honors the ``hyperslice`` bound so sliced reads mirror the real driver. + """ + + def __init__(self, arrays): + self._arrays = arrays + + def readDatasets(self, datasets): # noqa: N802 (mirror real h5coro API) + out = {} + for d in datasets: + if isinstance(d, str): + out[d] = self._arrays[d] + continue + path = d["dataset"] + arr = self._arrays[path] + hs = d.get("hyperslice") + if hs is not None: + lo, hi = hs[0] + arr = arr[lo:hi] + out[path] = arr + return out + + +class _ShardGrid: + """Grid stub: leaf id == row index; every row maps to ``shard_key`` 0, + so the spatial filter keeps all rows and the structured filters are + exercised in isolation.""" + + @staticmethod + def assign(lats, lons): + return np.arange(len(lats)) + + @staticmethod + def shards_of(leaf_ids): + return np.zeros(len(leaf_ids), dtype=int) + + +class TestPredicateMask: + def test_scalar_ops_1d(self): + arr = np.array([0, 1, 2, 3, 0]) + assert _predicate_mask( + arr, {"dataset": "/d", "op": "eq", "value": 0, "column": None} + ).tolist() == [True, False, False, False, True] + assert _predicate_mask( + arr, {"dataset": "/d", "op": "ge", "value": 2, "column": None} + ).tolist() == [False, False, True, True, False] + + def test_set_ops(self): + arr = np.array([2, 3, 4, 5]) + assert _predicate_mask(arr, {"dataset": "/d", "op": "in", "values": [2, 4]}).tolist() == [ + True, + False, + True, + False, + ] + assert _predicate_mask( + arr, {"dataset": "/d", "op": "not_in", "values": [2, 4]} + ).tolist() == [False, True, False, True] + + def test_keep_false_inverts(self): + arr = np.array([0, 1, 0]) + assert _predicate_mask( + arr, {"dataset": "/d", "op": "eq", "value": 0, "keep": False} + ).tolist() == [False, True, False] + + def test_nd_column_slicing(self): + # 2-D flag array (5 rows x 3 surface-type columns) + arr = np.array([[0, 9, 9], [-2, 9, 9], [1, 9, 9], [-2, 9, 9], [3, 9, 9]]) + # signal_conf_ph-style: column 0, != -2 + mask = _predicate_mask(arr, {"dataset": "/d", "column": 0, "op": "ne", "value": -2}) + assert mask.tolist() == [True, False, True, False, True] + + def test_nd_requires_column(self): + arr = np.zeros((3, 2)) + with pytest.raises(ValueError, match="requires an integer 'column'"): + _predicate_mask(arr, {"dataset": "/d", "op": "eq", "value": 0}) + + def test_column_on_1d_rejected(self): + arr = np.zeros(3) + with pytest.raises(ValueError, match="array is 1-D"): + _predicate_mask(arr, {"dataset": "/d", "column": 0, "op": "eq", "value": 0}) + + +class TestReadGroupFilters: + def _data_source(self, **extra): + ds = { + "coordinates": {"latitude": "/lat", "longitude": "/lon"}, + "variables": {"h": "/h"}, + } + ds.update(extra) + return ds + + def test_quality_filter_eq_path(self): + h5 = _FakeH5( + { + "/lat": np.array([1.0, 2.0, 3.0, 4.0]), + "/lon": np.array([1.0, 2.0, 3.0, 4.0]), + "/h": np.array([10.0, 20.0, 30.0, 40.0], dtype=np.float32), + "/qs": np.array([0, 1, 0, 1]), + } + ) + ds = self._data_source(quality_filter={"dataset": "/qs", "value": 0}) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [10.0, 30.0] + + def test_quality_filter_byte_identical_to_manual_eq(self): + # The synthesized base eq filter must reproduce the legacy mask exactly. + h = np.array([10.0, 20.0, 30.0, 40.0, 50.0], dtype=np.float32) + qs = np.array([0, 1, 0, 0, 1]) + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": h, + "/qs": qs, + } + ) + ds = self._data_source(quality_filter={"dataset": "/qs", "value": 0}) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + expected = h[qs == 0] + assert df["h"].to_numpy().tobytes() == expected.tobytes() + + def test_2d_signal_conf_filter(self): + conf = np.array([[0], [-2], [4], [-2], [3]]) # column 0, surface type + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.arange(5.0, dtype=np.float32), + "/conf": conf, + } + ) + ds = self._data_source(filters=[{"dataset": "/conf", "column": 0, "op": "ne", "value": -2}]) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [0.0, 2.0, 4.0] + + def test_multiple_anded_filters(self): + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.arange(5.0, dtype=np.float32), + "/conf": np.array([[5], [5], [0], [5], [5]]), + "/pod": np.array([0, 0, 0, 1, 0]), + } + ) + ds = self._data_source( + filters=[ + {"dataset": "/conf", "column": 0, "op": "ne", "value": 0}, + {"dataset": "/pod", "op": "eq", "value": 0}, + ] + ) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + # row2 dropped by conf==0, row3 dropped by pod==1 + assert df["h"].tolist() == [0.0, 1.0, 4.0] + + def test_in_op_integer_column(self): + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.arange(5.0, dtype=np.float32), + "/conf": np.array([[2], [0], [3], [1], [4]]), + } + ) + ds = self._data_source( + filters=[{"dataset": "/conf", "column": 0, "op": "in", "values": [2, 3, 4]}] + ) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [0.0, 2.0, 4.0] + + def test_expression_filter_base_level(self): + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.array([-1.0, 2.0, -3.0, 4.0, 5.0], dtype=np.float32), + } + ) + ds = self._data_source(filters=[{"expression": "h > 0"}]) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [2.0, 4.0, 5.0] + + def test_no_filter_keeps_all(self): + h5 = _FakeH5( + { + "/lat": np.arange(3.0), + "/lon": np.arange(3.0), + "/h": np.arange(3.0, dtype=np.float32), + } + ) + df = _read_group(h5, "gt1l", self._data_source(), 0, _ShardGrid()) + assert df["h"].tolist() == [0.0, 1.0, 2.0] + + def test_all_filtered_returns_none(self): + h5 = _FakeH5( + { + "/lat": np.arange(3.0), + "/lon": np.arange(3.0), + "/h": np.arange(3.0, dtype=np.float32), + "/qs": np.array([1, 1, 1]), + } + ) + ds = self._data_source(quality_filter={"dataset": "/qs", "value": 0}) + assert _read_group(h5, "gt1l", ds, 0, _ShardGrid()) is None + + def test_filter_dataset_coincides_with_variable_path(self): + # Filter dataset path == variable path exercises the dedup branch + # (path must appear exactly once in the h5coro read list). + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.array([0.0, 1.0, 2.0, 3.0, 4.0], dtype=np.float32), + } + ) + ds = self._data_source(filters=[{"dataset": "/h", "op": "ge", "value": 2.0}]) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [2.0, 3.0, 4.0] + + def test_expression_filter_after_structured(self): + # Expression filter ANDed after a structured predicate. + h5 = _FakeH5( + { + "/lat": np.arange(5.0), + "/lon": np.arange(5.0), + "/h": np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32), + "/qs": np.array([0, 0, 1, 0, 0]), + } + ) + ds = self._data_source( + filters=[ + {"dataset": "/qs", "op": "eq", "value": 0}, + {"expression": "h > 2"}, + ] + ) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + # structured drops row2 (qs==1); expression keeps h>2 from remainder + assert df["h"].tolist() == [4.0, 5.0] + + def test_two_sequential_expression_filters(self): + # Two sequential expression filters both applied. + h5 = _FakeH5( + { + "/lat": np.arange(6.0), + "/lon": np.arange(6.0), + "/h": np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.float32), + } + ) + ds = self._data_source( + filters=[ + {"expression": "h > 2"}, + {"expression": "h < 6"}, + ] + ) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [3.0, 4.0, 5.0] + + def test_expression_filter_undefined_name_raises(self): + # Expression referencing an undefined name re-raises as NameError. + h5 = _FakeH5( + { + "/lat": np.arange(3.0), + "/lon": np.arange(3.0), + "/h": np.arange(3.0, dtype=np.float32), + } + ) + ds = self._data_source(filters=[{"expression": "undefined_col > 0"}]) + with pytest.raises(NameError, match="undefined name"): + _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + + +# --------------------------------------------------------------------------- +# _expand_mask_to_base and cross-level filter path (issue #43, Phase B) +# --------------------------------------------------------------------------- + + +class TestExpandMaskToBase: + def test_single_parent_kept(self): + # 1 parent keeps base rows 0-2 (3 photons). + coarse = np.array([True]) + ibeg = np.array([0]) + cnt = np.array([3]) + out = _expand_mask_to_base(coarse, ibeg, cnt, index_base=0, total_base_size=3) + np.testing.assert_array_equal(out, [True, True, True]) + + def test_single_parent_dropped(self): + coarse = np.array([False]) + ibeg = np.array([0]) + cnt = np.array([3]) + out = _expand_mask_to_base(coarse, ibeg, cnt, index_base=0, total_base_size=3) + np.testing.assert_array_equal(out, [False, False, False]) + + def test_two_parents_alternating(self): + # parent 0 -> base rows 0-1 (kept); parent 1 -> base rows 2-4 (dropped). + coarse = np.array([True, False]) + ibeg = np.array([0, 2]) + cnt = np.array([2, 3]) + out = _expand_mask_to_base(coarse, ibeg, cnt, index_base=0, total_base_size=5) + np.testing.assert_array_equal(out, [True, True, False, False, False]) + + def test_index_base_shift(self): + # HDF5 1-based indexing: index_beg values start at 1. + coarse = np.array([False, True]) + ibeg = np.array([1, 4]) # 1-based + cnt = np.array([3, 2]) + out = _expand_mask_to_base(coarse, ibeg, cnt, index_base=1, total_base_size=5) + # parent1 covers base rows 3 and 4 (ibeg=4-1=3, cnt=2). + np.testing.assert_array_equal(out, [False, False, False, True, True]) + + def test_empty_coarse_mask(self): + coarse = np.array([False, False, False]) + ibeg = np.array([0, 2, 5]) + cnt = np.array([2, 3, 1]) + out = _expand_mask_to_base(coarse, ibeg, cnt, index_base=0, total_base_size=6) + assert not np.any(out) + + def test_full_coarse_mask(self): + coarse = np.array([True, True]) + ibeg = np.array([0, 3]) + cnt = np.array([3, 2]) + out = _expand_mask_to_base(coarse, ibeg, cnt, index_base=0, total_base_size=5) + assert np.all(out) + + def test_negative_beg_raises(self): + # index_beg_arr[0]=0 < index_base=1 -> beg=-1 -> must raise ValueError + coarse = np.array([True]) + ibeg = np.array([0]) + cnt = np.array([3]) + with pytest.raises(ValueError, match="less than index_base"): + _expand_mask_to_base(coarse, ibeg, cnt, index_base=1, total_base_size=3) + + +class TestReadGroupCrossLevel: + """Phase B: cross-level filters expand coarse verdicts to base-rate rows.""" + + def _data_source_with_levels(self, coarse_filter_value=None): + """Two-level data source: 'segments' -> 'photons' via link arrays.""" + ds = { + "coordinates": {"latitude": "/lat", "longitude": "/lon"}, + "variables": {"h": "/h"}, + "base_level": "photons", + "levels": { + "photons": { + "path": "/heights", + "coordinates": ["lat", "lon"], + "variables": ["h"], + "link": None, + }, + "segments": { + "path": "/geolocation", + "coordinates": [], + "variables": ["signal_conf_ph"], + "link": { + "to": "photons", + "index_beg": "/ph_index_beg", + "count": "/segment_ph_cnt", + }, + }, + }, + } + if coarse_filter_value is not None: + ds["filters"] = [ + { + "dataset": "/conf", + "op": "ne", + "value": coarse_filter_value, + "level": "segments", + } + ] + return ds + + def test_coarse_filter_expands_to_base(self): + # 3 segments, each covering 2 photons (6 total). + # segment1 (conf=-2) -> drop; segment0 and segment2 -> keep. + h5 = _FakeH5( + { + "/lat": np.arange(6.0), + "/lon": np.arange(6.0), + "/h": np.arange(6.0, dtype=np.float32), + # link arrays: segment0->ph[0:2], segment1->ph[2:4], segment2->ph[4:6] + "/ph_index_beg": np.array([0, 2, 4]), + "/segment_ph_cnt": np.array([2, 2, 2]), + # coarse flag: segment1 has conf=-2, others conf=4 + "/conf": np.array([4, -2, 4]), + } + ) + ds = self._data_source_with_levels(coarse_filter_value=-2) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + # Segments 0 and 2 survive; their photons are h[0:2] and h[4:6]. + assert df["h"].tolist() == [0.0, 1.0, 4.0, 5.0] + + def test_all_segments_filtered_returns_none(self): + h5 = _FakeH5( + { + "/lat": np.arange(4.0), + "/lon": np.arange(4.0), + "/h": np.arange(4.0, dtype=np.float32), + "/ph_index_beg": np.array([0, 2]), + "/segment_ph_cnt": np.array([2, 2]), + "/conf": np.array([-2, -2]), # both segments dropped + } + ) + ds = self._data_source_with_levels(coarse_filter_value=-2) + assert _read_group(h5, "gt1l", ds, 0, _ShardGrid()) is None + + def test_cross_level_and_base_level_filters_anded(self): + # Cross-level keeps segments 0 and 2 (photons 0-1 and 4-5); + # base-level h>1 further drops photon 0 and photon 4. + h5 = _FakeH5( + { + "/lat": np.arange(6.0), + "/lon": np.arange(6.0), + "/h": np.array([0.0, 1.5, 2.0, 2.5, 3.0, 4.0], dtype=np.float32), + "/ph_index_beg": np.array([0, 2, 4]), + "/segment_ph_cnt": np.array([2, 2, 2]), + "/conf": np.array([4, -2, 4]), + "/qs": np.array([0, 0, 1, 0, 0, 0]), # base-level flag + } + ) + ds = self._data_source_with_levels(coarse_filter_value=-2) + # Add a base-level structured filter alongside the coarse one. + ds["filters"].append({"dataset": "/qs", "op": "eq", "value": 0}) + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + # Cross-level: keep ph0,1,4,5; base-level drops ph2 (qs==1 after reindex). + # Expected survivors among ph0,1,4,5: qs[0]=0,qs[1]=0,qs[4]=0,qs[5]=0 -> all 4 + assert df["h"].tolist() == [0.0, 1.5, 3.0, 4.0] + + def test_flat_form_unchanged(self): + # No levels/base_level -> flat path still works. + h5 = _FakeH5( + { + "/lat": np.arange(3.0), + "/lon": np.arange(3.0), + "/h": np.array([1.0, 2.0, 3.0], dtype=np.float32), + "/qs": np.array([0, 1, 0]), + } + ) + ds = { + "coordinates": {"latitude": "/lat", "longitude": "/lon"}, + "variables": {"h": "/h"}, + "quality_filter": {"dataset": "/qs", "value": 0}, + } + df = _read_group(h5, "gt1l", ds, 0, _ShardGrid()) + assert df["h"].tolist() == [1.0, 3.0] diff --git a/tests/test_read_plan.py b/tests/test_read_plan.py new file mode 100644 index 00000000..83e5358d --- /dev/null +++ b/tests/test_read_plan.py @@ -0,0 +1,224 @@ +"""Tests for the offline read-plan module (issue #43, Phase C). + +All tests use synthetic numpy arrays — no h5coro, S3, or credentials needed. +""" + +import numpy as np + +from zagg.read_plan import ReadPlan, execute_read_plan, plan_read + + +def _isolated_setup(n_segs, seg_len=10, lat_step=40.0): + """Build lat/lon arrays where adjacent rep-points are far enough apart + that a 1-degree-tall bbox around segment 0 does not intersect the + linestring from segment 0 to segment 1. + + Segments sit at lat = [0, lat_step, 2*lat_step, ...] and lon = 0.5. + With lat_step=40 the linestring from seg 0 (lat=0) to seg 1 (lat=40) + does NOT cross a bbox limited to lat < 0.5, so segment 0 is the + only match for bbox=(0, -0.5, 1, 0.5). + """ + lats = np.arange(n_segs, dtype=float) * lat_step + lons = np.full(n_segs, 0.5) + index_beg = np.arange(n_segs, dtype=int) * seg_len + count = np.full(n_segs, seg_len, dtype=int) + n_base = n_segs * seg_len + return lats, lons, index_beg, count, n_base + + +class TestPlanReadBasic: + def test_single_parent_in_aoi(self): + # Segment 0 sits at lat=0. bbox is (0, -0.5, 1, 0.5) — only its rep-point + # lands inside; the linestring to seg 1 (lat=40) exits immediately upward + # and never re-enters this narrow band. + lats, lons, idx, cnt, n_base = _isolated_setup(5) + bbox = (0.0, -0.5, 1.0, 0.5) + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0) + assert not plan.full_read + assert len(plan.parent_runs) == 1 + assert plan.parent_runs[0] == (0, 0) + assert plan.base_slices[0] == (0, 10) + + def test_empty_aoi_returns_empty_plan(self): + lats, lons, idx, cnt, n_base = _isolated_setup(5) + bbox = (100.0, 100.0, 110.0, 110.0) # nowhere near our data + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0) + assert plan.parent_runs == [] + assert plan.base_slices == [] + assert plan.chunk_lists == [] + assert not plan.full_read + + def test_adjacent_parents_merged_into_one_run(self): + # Both segs 0 (lat=0) and 1 (lat=40) land inside a wide bbox. + lats, lons, idx, cnt, n_base = _isolated_setup(5) + bbox = (0.0, -0.5, 1.0, 40.5) # covers segs 0 and 1 + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0) + assert len(plan.parent_runs) == 1 + assert plan.parent_runs[0] == (0, 1) + assert plan.base_slices[0] == (0, 20) + + def test_non_contiguous_parents_give_separate_runs(self): + # Four segments spaced 40 lat-degrees apart. Only segs 0 and 3 land in + # their respective narrow bboxes. Because they are not contiguous, testing + # with a single-call narrow bbox on seg 0 only. + lats, lons, idx, cnt, n_base = _isolated_setup(4) + # Only seg 0 in bbox + bbox = (0.0, -0.5, 1.0, 0.5) + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0) + assert len(plan.parent_runs) == 1 + assert plan.parent_runs[0] == (0, 0) + + def test_pad_extends_run(self): + # Seg 0 in aoi; pad=1 should extend the run to include seg 1. + lats, lons, idx, cnt, n_base = _isolated_setup(5) + bbox = (0.0, -0.5, 1.0, 0.5) # only seg 0 in AOI + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=1) + # pad=1 extends [0,0] -> [0, 1] (clamped at start 0) + assert plan.parent_runs[0] == (0, 1) + assert plan.base_slices[0] == (0, 20) + + def test_pad_clamps_at_boundaries(self): + lats, lons, idx, cnt, n_base = _isolated_setup(5) + bbox = (0.0, -0.5, 1.0, 0.5) # only seg 0 in AOI + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=2) + # pad=2, start clamped to 0, end = min(4, 2) = 2 + assert plan.parent_runs[0][0] == 0 + assert plan.parent_runs[0][1] == 2 + + def test_selectivity_fallback_to_full_read(self): + # All 10 segs in aoi -> total_base = n_base -> full_read + lats, lons, idx, cnt, n_base = _isolated_setup(10, seg_len=100, lat_step=1.0) + bbox = (-1.0, -1.0, 2.0, 10.0) # covers all segments + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0, full_read_threshold=0.9) + assert plan.full_read + + def test_selectivity_below_threshold_no_full_read(self): + lats, lons, idx, cnt, n_base = _isolated_setup(10, seg_len=10) + bbox = (0.0, -0.5, 1.0, 0.5) # only seg 0 -> 10/100 = 0.1 + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0, full_read_threshold=0.9) + assert not plan.full_read + + +class TestPlanReadCrossingCase: + def test_linestring_crossing_bbox(self): + """A parent whose rep-point is just outside the bbox but whose linestring + to the next parent crosses the bbox is still included.""" + # Two segments: seg0 at lon=0 (outside bbox), seg1 at lon=2 (outside bbox). + # The line from (0,0)->(2,0) crosses bbox (0.5, -0.5, 1.5, 0.5). + lats = np.array([0.0, 0.0]) + lons = np.array([0.0, 2.0]) + idx = np.array([0, 5]) + cnt = np.array([5, 5]) + n_base = 10 + # bbox crosses the midpoint of the lon=0 to lon=2 linestring + bbox = (0.5, -0.5, 1.5, 0.5) + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, pad=0) + assert not plan.full_read + # seg0's linestring to seg1 crosses the bbox -> seg0 included + assert len(plan.parent_runs) >= 1 + assert plan.parent_runs[0][0] == 0 + + +class TestPlanReadIndexBase: + def test_index_base_1_atl03_style(self): + """index_base=1: ph_index_beg=1 means base[0], ph_index_beg=6 means base[5].""" + lats = np.array([0.0, 100.0]) + lons = np.array([0.5, 0.5]) + idx = np.array([1, 6]) # 1-based + cnt = np.array([5, 5]) + n_base = 10 + bbox = (0.0, -0.5, 1.0, 0.5) # only seg 0 (lat=0) + plan = plan_read(lats, lons, idx, cnt, n_base, bbox, index_base=1, pad=0) + assert len(plan.parent_runs) == 1 + assert plan.base_slices[0] == (0, 5) # 1-1=0, 0+5=5 + + +class TestExecuteReadPlan: + def _make_plan(self, slices, full_read=False): + runs = [(s, e - 1) for s, e in slices] + chunks = [[(s, e - 1)] for s, e in slices] + return ReadPlan( + parent_runs=runs, + base_slices=list(slices), + chunk_lists=chunks, + full_read=full_read, + ) + + def test_empty_plan_returns_empty_array(self): + plan = ReadPlan(parent_runs=[], base_slices=[], chunk_lists=[]) + calls = [] + + def read_fn(path, hyperslice=None): + calls.append(hyperslice) + return np.array([]) + + out = execute_read_plan(plan, read_fn, "/h", np.float32) + assert len(out) == 0 + assert len(calls) == 0 # no read calls for empty plan + + def test_single_slice_reads_correct_range(self): + data = np.arange(100.0, dtype=np.float32) + plan = self._make_plan([(10, 20)]) + + def read_fn(path, hyperslice=None): + assert hyperslice is not None + lo, hi = hyperslice[0] + return data[lo : hi + 1] # h5coro inclusive end + + out = execute_read_plan(plan, read_fn, "/h", np.float32) + np.testing.assert_array_equal(out, data[10:20]) + + def test_multiple_slices_concatenated(self): + data = np.arange(100.0, dtype=np.float32) + plan = self._make_plan([(5, 10), (20, 25)]) + + def read_fn(path, hyperslice=None): + lo, hi = hyperslice[0] + return data[lo : hi + 1] + + out = execute_read_plan(plan, read_fn, "/h", np.float32) + expected = np.concatenate([data[5:10], data[20:25]]) + np.testing.assert_array_equal(out, expected) + + def test_full_read_plan_calls_without_hyperslice(self): + data = np.arange(50.0, dtype=np.float32) + plan = ReadPlan( + parent_runs=[(0, 4)], + base_slices=[(0, 50)], + chunk_lists=[[(0, 49)]], + full_read=True, + ) + called_with_none = [] + + def read_fn(path, hyperslice=None): + called_with_none.append(hyperslice is None) + return data + + out = execute_read_plan(plan, read_fn, "/h", np.float32) + assert called_with_none == [True] + np.testing.assert_array_equal(out, data) + + def test_dtype_coercion(self): + data = np.array([1, 2, 3], dtype=np.int32) + plan = self._make_plan([(0, 3)]) + + def read_fn(path, hyperslice=None): + return data + + out = execute_read_plan(plan, read_fn, "/h", np.float64) + assert out.dtype == np.float64 + + def test_full_read_true_with_empty_parent_runs(self): + # full_read=True must take precedence over empty parent_runs; the full + # dataset should be returned rather than an empty array. + data = np.arange(20.0, dtype=np.float32) + plan = ReadPlan(parent_runs=[], base_slices=[], chunk_lists=[], full_read=True) + calls = [] + + def read_fn(path, hyperslice=None): + calls.append(hyperslice) + return data + + out = execute_read_plan(plan, read_fn, "/h", np.float32) + assert calls == [None] # full-read call, no hyperslice + np.testing.assert_array_equal(out, data)