Generic hierarchical filtering & targeted reads (#43)

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: "<py expr>"}`` 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