From 552e54a8c8833acb461591b0ec40f6ce0d65adeb Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Wed, 14 Oct 2026 14:47:08 +0200 Subject: [PATCH 01/16] feat(query-engine): add opt-in Spark pushdown optimizer for resample Introduce a rule-based query optimizer that computes supported expression subtrees natively in Spark before the pandas UDF, starting with `channel.resample(rate)` on physical channels. - Add `QueryOptimizer`, `PushdownPlan`, `PushdownRule`, and `PushdownContext` in `analyze.query.optimizer`. - Implement `ResamplePushdownRule` to push `SampleSeries.resample` down to Spark, reproducing the pandas behavior block-by-block. - Wire `enable_pushdown` through `DefaultSolver`, `QueryEngine`, and report config; default is `false`. - Make `TimeSeriesOp.build` consult `SeriesCache.materialized` so pushed-down results can be injected into the evaluation tree. - Add optimizer API reference pages, update configuration docs, and document the Spark pushdown opt-in behavior in query solver and TSAL guides. --- docs/impulse/docs/config/configuration.md | 1 + .../metadata/time_series_expression.md | 4 + .../analyze/query/optimizer/plan.md | 99 ++++++ .../analyze/query/optimizer/pushdown.md | 156 +++++++++ .../query/optimizer/query_optimizer.md | 54 +++ .../analyze/query/optimizer/resample_rule.md | 28 ++ .../analyze/query/solvers/default_solver.md | 5 + .../impulse_reporting/config/config_parser.md | 5 + docs/impulse/docs/references/api/sidebar.json | 10 + .../references/query_engine/query_solvers.md | 47 +++ .../query_engine/tsal/defining_expressions.md | 4 + .../query_engine/tsal/evaluation.md | 5 + docs/impulse/pydoc-markdown.yml | 4 + .../metadata/time_series_expression.py | 8 + .../analyze/query/optimizer/__init__.py | 16 + .../analyze/query/optimizer/plan.py | 143 ++++++++ .../analyze/query/optimizer/pushdown.py | 148 +++++++++ .../query/optimizer/query_optimizer.py | 103 ++++++ .../analyze/query/optimizer/resample_rule.py | 221 +++++++++++++ .../analyze/query/optimizer/traversal.py | 82 +++++ .../analyze/query/solvers/default_solver.py | 204 ++++++++++-- .../analyze/query/solvers/series_cache.py | 20 ++ src/impulse_reporting/config/config_parser.py | 7 + src/impulse_reporting/core/report.py | 1 + .../unit/analyze/query/optimizer/__init__.py | 0 .../query/optimizer/query_optimizer_test.py | 235 +++++++++++++ .../query/optimizer/resample_rule_test.py | 181 ++++++++++ .../solvers/default_solver_pushdown_test.py | 308 ++++++++++++++++++ .../unit/config/config_parser_test.py | 12 + .../unit/core/report_test.py | 19 ++ 30 files changed, 2104 insertions(+), 26 deletions(-) create mode 100644 docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md create mode 100644 docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md create mode 100644 docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md create mode 100644 docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md create mode 100644 src/impulse_query_engine/analyze/query/optimizer/__init__.py create mode 100644 src/impulse_query_engine/analyze/query/optimizer/plan.py create mode 100644 src/impulse_query_engine/analyze/query/optimizer/pushdown.py create mode 100644 src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py create mode 100644 src/impulse_query_engine/analyze/query/optimizer/resample_rule.py create mode 100644 src/impulse_query_engine/analyze/query/optimizer/traversal.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/optimizer/__init__.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py diff --git a/docs/impulse/docs/config/configuration.md b/docs/impulse/docs/config/configuration.md index d8199953..e9894c24 100644 --- a/docs/impulse/docs/config/configuration.md +++ b/docs/impulse/docs/config/configuration.md @@ -141,6 +141,7 @@ Two independent filter families: | `drop_implausible_data` | `bool` | `false` | When `true`, drops `channels` rows where `is_plausible = false`. Requires `data_type = "RAW"`; combining with `"RLE"` raises a validation error. | | `max_channels_per_batch` | `int` | `500` | Caps how many distinct channels are read in one solve batch, to bound per-solve memory. Counts channel selections (each distinct `channel(...)` / `poi_channel(...)` in your expressions). Applies to events, aggregations, and calculated channels; for calculated channels it counts the *input* channels they read, not the number of calculated channels. The cap is a packing target, not a hard limit: an expression that on its own reads more channels than the cap still runs, in a batch of its own. Must be `>= 1`. The old name `batch_size` is a **deprecated alias** (still accepted, warns, removed in a future release); setting both raises a validation error. | | `max_containers_per_batch` | `int` | `null` | Caps how many containers are solved at once, to bound per-solve memory. `null` disables the cap; when set it must be `>= 1`. On an **incremental** run, `Report.run()` commits at most this many upserted containers per iteration and loops until the population is drained (see [incremental](#incremental-optional)). On a **full** run, and when a changed definition recomputes over all containers, the population is split into chunks of about this size and solved chunk by chunk. Chunk sizes are approximate (the cap is a memory heuristic, not an exact bound). | +| `enable_pushdown` | `bool` | `false` | When `true`, supported expression subtrees (currently `channel.resample(rate)` on a physical channel) are computed natively in Spark before the pandas evaluation, so fewer rows are shipped to it. A channel is only pushed down when every use of it in a solve batch is; results are unchanged. Meant for downsampling: upsampling produces more rows than the raw channel. See [Spark pushdown](../references/query_engine/query_solvers.md#spark-pushdown-opt-in). | | `solver_config` | `SolverConfig` | `null` | Per-table column mappings, per-table equality filters, and project scoping. Set `project_id` to scope reads by project — it is applied to `container_tags` (if configured), `container_metrics`, and `channel_mapping` (if configured), so it works in both narrow EAV and wide-only data models. Omit it when you don't need project scoping. See [Solver column mappings and filters](#solver-column-mappings-and-filters). | If `query_engine` is omitted, the default is `DefaultSolver` with diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/metadata/time_series_expression.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/metadata/time_series_expression.md index e20fc5c0..5a3c9a7e 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/metadata/time_series_expression.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/metadata/time_series_expression.md @@ -353,6 +353,10 @@ def build(cache: SeriesCache) Build the time series from cache. +A result the cache already holds for this node (see +:meth:`SeriesCache.materialized`) is returned as-is, without building the +operands. + **Arguments**: - `cache` (`SeriesCache`): Cache containing time series data. diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md new file mode 100644 index 00000000..9bf7bb41 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md @@ -0,0 +1,99 @@ +--- +sidebar_label: plan +title: impulse_query_engine.analyze.query.optimizer.plan +--- + +The outcome of query optimization: which subtrees Spark computes ahead of the UDF. + + +## PushdownPlan + +```python +class PushdownPlan() +``` + +Pushdowns selected for one solve call, shipped to the solver and the UDF. + +Variant ``0`` is the raw channel data; variants ``1..n`` are derived series, +one per distinct :attr:`Pushdown.variant_key`. + +**Arguments**: + +- `rules` (`tuple of PushdownRule`): Rules used to recognize pushed-down nodes during evaluation. +- `variant_ids` (`Mapping[tuple, int]`): Variant key to variant id. +- `pushdowns_by_variant` (`Mapping[int, Pushdown]`): Variant id to a representative pushdown (rule and parameters). +- `variants_by_selector` (`Mapping[int, tuple of int]`): Selector id to the variant ids derived from it. +- `raw_selector_ids` (`frozenset of int`): Selectors whose raw channel data is still required. + +#### build + +```python +def build(rules: Iterable[PushdownRule], pushdowns: Iterable[Pushdown], + selector_ids: Iterable[int]) -> PushdownPlan +``` + +Assign variant ids to *pushdowns*, in discovery order. + +**Arguments**: + +- `rules` (`Iterable[PushdownRule]`): The optimizer's rules. +- `pushdowns` (`Iterable[Pushdown]`): Accepted pushdowns; each one's selector is pushed down entirely. +- `selector_ids` (`Iterable[int]`): Ids of all selectors in the solve. + +#### is\_empty + +```python +def is_empty() -> bool +``` + +``True`` when nothing is pushed down. + + +#### pushed\_selector\_ids + +```python +def pushed_selector_ids() -> frozenset[int] +``` + +Selectors whose raw channel data is replaced by derived variants. + + +#### variants\_by\_rule + +```python +def variants_by_rule() -> list[tuple[PushdownRule, dict[int, Pushdown]]] +``` + +Group the variants by the rule that computes them. + + +#### lookup + +```python +def lookup(node: TimeSeriesExpression) -> tuple[Pushdown, int] | None +``` + +Return ``(pushdown, variant)`` when *node* is a pushed-down subtree. + + +#### channel\_variants + +```python +def channel_variants(selector_ids: Column) -> Column +``` + +Variant ids a matched channel ships, given its ``selector_ids`` array. + +A channel ships only its raw data (``[0]``) unless every one of its +selectors is pushed down; then it ships the derived variants instead. + + +#### describe + +```python +def describe() -> str +``` + +One-line summary of the plan for logging. + + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md new file mode 100644 index 00000000..04fa9774 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -0,0 +1,156 @@ +--- +sidebar_label: pushdown +title: impulse_query_engine.analyze.query.optimizer.pushdown +--- + +Building blocks of Spark pushdown rules. + + +#### VARIANT\_COL + +Column tagging each channel-data row with the series variant it belongs to. + + +#### RAW\_VARIANT + +Variant id of the unmodified (raw) channel data. + + +## PushdownContext + +```python +class PushdownContext() +``` + +Plan-time facts about the channel data a pushdown would run on. + +**Arguments**: + +- `tstart_type, tend_type` (`pyspark.sql.types.DataType`): Spark types of the channel-data ``tstart`` / ``tend`` columns. + +#### from\_schema + +```python +def from_schema(schema: T.StructType, config: SolverConfig) -> PushdownContext +``` + +Build the context from the column-mapped channel-data schema. + + +## Pushdown + +```python +class Pushdown() +``` + +A subtree that a :class:`PushdownRule` computes natively in Spark. + +Expression nodes overload ``==``, so pushdowns compare by identity; use +:attr:`variant_key` to identify equivalent pushdowns. + +**Arguments**: + +- `rule` (`PushdownRule`): The rule that matched the subtree. +- `selector` (`TimeSeriesSelector`): The physical channel selector the subtree reads. +- `params` (`tuple of (str, Any)`): Canonical, hashable operation parameters (e.g. ``(("sample_rate", 1e6),)``). +- `consumed` (`tuple of TimeSeriesSelector`): Every selector occurrence inside the subtree (its ``get_selectors()``). + +#### variant\_key + +```python +def variant_key() -> tuple +``` + +Identity of the derived series, independent of the input channel. + + +#### describe + +```python +def describe() -> str +``` + +Human-readable ``rule(param=value, ...)`` label. + + +## PushdownRule + +```python +class PushdownRule(abc.ABC) +``` + +Recognizes a subtree and computes it in Spark ahead of the pandas UDF. + +A rule's derived rows use the channel-data shape +``(container_id, channel_id, variant, tstart, tend, value)``; :meth:`decode` +turns one container/channel/variant slice of them back into the object the +original subtree evaluates to. The result must equal evaluating the subtree +in pandas on the same channel data. + + +#### name + +Unique rule name, part of every :attr:`Pushdown.variant_key`. + + +#### match + +```python +def match(node: TimeSeriesExpression) -> Pushdown | None +``` + +Return the pushdown for *node*, or ``None`` if the rule does not apply. + +Purely syntactic: it runs on the driver while planning and again inside +the UDF to recognize the node, so it must not depend on data or schemas. + + +#### supports + +```python +def supports(pushdown: Pushdown, context: PushdownContext) -> bool +``` + +Whether *pushdown* can run on channel data described by *context*. + + +#### result\_type + +```python +def result_type(pushdown: Pushdown) -> type +``` + +Core-model class :meth:`decode` returns for *pushdown*. + + +#### to\_spark + +```python +def to_spark(rows: DataFrame, variants: Mapping[int, Pushdown], + config: SolverConfig) -> DataFrame +``` + +Compute all *variants* of this rule in one pass. + +**Arguments**: + +- `rows` (`pyspark.sql.DataFrame`): Channel-data rows tagged with :data:`VARIANT_COL`, one copy per +variant a channel needs; only ids in *variants* occur. +- `variants` (`Mapping[int, Pushdown]`): Variant id to a representative pushdown carrying its parameters. +- `config` (`SolverConfig`): Provides the internal column names. + +**Returns**: + +`pyspark.sql.DataFrame`: Derived rows with the container id, channel id, :data:`VARIANT_COL`, +``tstart``, ``tend`` and ``value`` columns of *rows*, same types. + +#### decode + +```python +def decode(rows: pd.DataFrame, pushdown: Pushdown, + cache: TimeSeriesCache) -> Any +``` + +Build the subtree's result from one container/channel/variant slice. + + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md new file mode 100644 index 00000000..aceaecb0 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md @@ -0,0 +1,54 @@ +--- +sidebar_label: query_optimizer +title: impulse_query_engine.analyze.query.optimizer.query_optimizer +--- + +Rule-based optimizer choosing the subtrees Spark computes ahead of the pandas UDF. + + +#### default\_rules + +```python +def default_rules() -> tuple[PushdownRule, ...] +``` + +All built-in pushdown rules. + + +## QueryOptimizer + +```python +class QueryOptimizer() +``` + +Plans which subtrees of a solve's selections are computed natively in Spark. + +All selections of a solve share the raw data of their physical channels, so +a subtree is pushed down only when **every** occurrence of its input selector, +across all selections, sits inside a pushed-down subtree. Then the raw rows +of that selector are no longer needed and are replaced by the (smaller) +derived rows. Selectors used anywhere else keep their raw data, and their +subtrees are evaluated in pandas as before. + +**Arguments**: + +- `rules` (`Sequence[PushdownRule]`): Rules to apply, in priority order. Defaults to :func:`default_rules`. + +#### optimize + +```python +def optimize(selections: Iterable[Any], + context: PushdownContext) -> PushdownPlan +``` + +Build the pushdown plan for the *selections* of one solve call. + +**Arguments**: + +- `selections` (`Iterable`): The query's selections. +- `context` (`PushdownContext`): Plan-time facts about the channel data. + +**Returns**: + +`PushdownPlan`: The plan; empty when nothing can be pushed down. + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md new file mode 100644 index 00000000..60f36085 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md @@ -0,0 +1,28 @@ +--- +sidebar_label: resample_rule +title: impulse_query_engine.analyze.query.optimizer.resample_rule +--- + +Pushdown rule computing ``SampleSeries.resample`` natively in Spark. + + +## ResamplePushdownRule + +```python +class ResamplePushdownRule(PushdownRule) +``` + +Pushes ``channel.resample(sample_rate)`` on a physical channel down to Spark. + +Reproduces :meth:`SampleSeries.resample`: per continuous block of samples, a +regular grid starting at the block's first ``tstart`` (built like +``numpy.arange``), each grid point holding the value of the last sample that +starts at or before it. Meant for downsampling: a rate finer than the native +sample spacing yields more rows than the raw channel. + +Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. +numpy builds that fuse this into a multiply-add can differ by one ulp for +fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral +grids are exact. + + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md index 9df87eef..463e0f3c 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md @@ -48,6 +48,11 @@ silver layer. ``RawEncoder.RLE`` (default) run-length encodes equal-valued runs; ``RawEncoder.INTERVAL`` only derives ``tend`` and drops exact duplicates. Only consulted when ``is_raw_data`` is ``True``. +- `enable_pushdown` (`bool`): Whether to compute supported expression subtrees (e.g. +``channel.resample(rate)``) natively in Spark before the pandas UDF, so +fewer rows are shipped to it (see :class:`QueryOptimizer`). A channel +is only pushed down when every use of it in the solve is. Default +``False``. #### filter\_container\_tags diff --git a/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md b/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md index dcaca28e..23e27e02 100644 --- a/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md +++ b/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md @@ -208,6 +208,11 @@ scoping. Key sub-fields: When omitted, all default column names are used and no project/toolbox filtering is applied. +- `enable_pushdown` (`bool, default=False`): Compute supported expression subtrees (currently +``channel.resample(rate)`` on a physical channel) natively in Spark +before the pandas UDF, so fewer rows are shipped to it. A channel is +only pushed down when every use of it in a solve batch is. Meant for +downsampling. #### validate\_drop\_implausible\_data\_requires\_raw diff --git a/docs/impulse/docs/references/api/sidebar.json b/docs/impulse/docs/references/api/sidebar.json index 3f08df0f..90c19dfd 100644 --- a/docs/impulse/docs/references/api/sidebar.json +++ b/docs/impulse/docs/references/api/sidebar.json @@ -30,6 +30,16 @@ "label": "impulse_query_engine.analyze.query.channels", "type": "category" }, + { + "items": [ + "references/api/impulse_query_engine/analyze/query/optimizer/plan", + "references/api/impulse_query_engine/analyze/query/optimizer/pushdown", + "references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer", + "references/api/impulse_query_engine/analyze/query/optimizer/resample_rule" + ], + "label": "impulse_query_engine.analyze.query.optimizer", + "type": "category" + }, { "items": [ "references/api/impulse_query_engine/analyze/query/solvers/default_solver", diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index c71cd297..30ecb684 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -126,6 +126,9 @@ Solver tuning lives under the `query_engine` section of your report config: — adapt the solver to a silver layer whose physical column names diverge from Impulse's internal names, scope reads by `project_id`, or apply per-table equality filters. +- [`query_engine.enable_pushdown`](../../config/configuration.md#query_engine-optional) + — compute supported expressions natively in Spark before the pandas UDF + (see [Spark pushdown](#spark-pushdown-opt-in)). Off by default. If `query_engine` is omitted from your config entirely, the engine runs with `DefaultSolver` and `data_type = "RLE"`. @@ -158,6 +161,47 @@ df = db.query.select(cc).solve_calculated_channels(spark, solver=DefaultSolver(s The reporting layer builds on this stage to persist derived channels to the gold layer — see [Channels](../report/channel.md). +## Spark pushdown (opt-in) + +Every selection of a solve is evaluated in one pandas UDF per container, and by default the raw +rows of every channel the selections reference are shipped into it. With pushdown enabled, a +rule-based optimizer finds expression subtrees that Spark can compute natively, computes them +**before** the UDF, and ships their (smaller) result instead of the raw rows. + +Supported today: + +| Expression | Computed in Spark as | +|------------------------------------|---------------------------------------------------------------------------| +| `channel.resample(sample_rate)` | The resampled series, identical to `SampleSeries.resample`, on a physical `channel()` or `channel_with_alias()` selection | + +The subtree can sit anywhere in an expression, e.g. `veh_spd.resample(1e6).cumtrapz()`, inside a +`CalculatedChannel`, an aggregation, or an event condition; only the pushed-down part runs in Spark. + +```python +solver = DefaultSolver(spark, enable_pushdown=True) +``` + +In a report config, set `query_engine.enable_pushdown` to `true`. + +:::note When a channel is pushed down +All selections of a solve share the raw data of their channels, so a channel is pushed down only +if **every** use of it in the solve is a pushed-down subtree. If the same channel is also used in +any other way, e.g. `rpm.resample(1e6)` next to `rpm.max()`, its raw rows are still needed: they are +shipped as before and the resample runs in pandas. Results are the same either way. +::: + +:::caution Intended for downsampling +`resample` is pushed down regardless of the rate. A rate finer than the channel's native sample +spacing (upsampling) produces **more** rows than the raw channel, so enable pushdown for workloads +that downsample. +::: + +Results match the pandas evaluation exactly for integer timestamp columns with an integral +`sample_rate` (e.g. µs or ns epochs), and for double timestamps whose grid is exact in binary +floating point. For double timestamps with a fractional rate such as `0.1`, the grid points can +differ in the last bit from numpy builds that fuse multiply-add operations. The order of samples +that share the same `tstart` is unspecified, as in the pandas evaluation. + ## API reference Auto-generated symbol-level docs: @@ -167,3 +211,6 @@ Auto-generated symbol-level docs: — abstract base class defining the six-stage solver pipeline. - [`SolverConfig`](../api/impulse_query_engine/analyze/query/solvers/solver_config.md) — per-table column mappings, filters, and project scoping. +- [`QueryOptimizer`](../api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md) + and [`PushdownRule`](../api/impulse_query_engine/analyze/query/optimizer/pushdown.md) + — the Spark pushdown optimizer and the interface for its rules. diff --git a/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md b/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md index f57b9756..a8bd2647 100644 --- a/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md +++ b/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md @@ -177,6 +177,10 @@ type) at execution time. distance_km = veh_spd.resample(1e6).cumtrapz() / 3600 / 1e6 ``` +When downsampling large channels, `.resample(...)` on a physical channel can be computed in Spark +before the data reaches the pandas evaluation — see +[Spark pushdown](../query_solvers.md#spark-pushdown-opt-in). + ### Filtering | Method | Signature | Description | diff --git a/docs/impulse/docs/references/query_engine/tsal/evaluation.md b/docs/impulse/docs/references/query_engine/tsal/evaluation.md index a8734e14..36becb3d 100644 --- a/docs/impulse/docs/references/query_engine/tsal/evaluation.md +++ b/docs/impulse/docs/references/query_engine/tsal/evaluation.md @@ -63,6 +63,11 @@ following per container: | `PointsInTimeSeries` | `ArrayType(ArrayType(DoubleType))` | `[[tstart, value], ...]`| | scalar | `DoubleType` | the value | +With [Spark pushdown](../query_solvers.md#spark-pushdown-opt-in) enabled, some subtrees (e.g. +`channel.resample(rate)`) are computed natively in Spark before step 1. The solver ships their +results instead of the raw channel rows, and step 2 uses those results in place of evaluating the +subtrees in pandas. The tree itself is unchanged, and so is every result. + `toPandas()` deserializes the binary `SampleSeries` columns back into objects; the array-backed types are returned as nested lists. See the [Core Data Model](core_data_model.md) for the semantics of each result class, and [Query Solvers](../query_solvers.md) for the full solver pipeline and the diff --git a/docs/impulse/pydoc-markdown.yml b/docs/impulse/pydoc-markdown.yml index f52dda2b..515a8dfe 100644 --- a/docs/impulse/pydoc-markdown.yml +++ b/docs/impulse/pydoc-markdown.yml @@ -24,6 +24,10 @@ loaders: - impulse_query_engine.analyze.query.solvers.default_solver - impulse_query_engine.analyze.query.solvers.solver_config - impulse_query_engine.analyze.query.solvers.query_solver + - impulse_query_engine.analyze.query.optimizer.pushdown + - impulse_query_engine.analyze.query.optimizer.plan + - impulse_query_engine.analyze.query.optimizer.query_optimizer + - impulse_query_engine.analyze.query.optimizer.resample_rule - impulse_query_engine.model.series.sample_series - impulse_query_engine.model.series.intervals - impulse_query_engine.model.series.points_in_time diff --git a/src/impulse_query_engine/analyze/metadata/time_series_expression.py b/src/impulse_query_engine/analyze/metadata/time_series_expression.py index 6a2b75bd..8b335011 100644 --- a/src/impulse_query_engine/analyze/metadata/time_series_expression.py +++ b/src/impulse_query_engine/analyze/metadata/time_series_expression.py @@ -1226,6 +1226,10 @@ def build(self, cache: SeriesCache): """ Build the time series from cache. + A result the cache already holds for this node (see + :meth:`SeriesCache.materialized`) is returned as-is, without building the + operands. + Parameters ---------- cache : SeriesCache @@ -1236,6 +1240,10 @@ def build(self, cache: SeriesCache): Any Built time series object. """ + # Looked up on the type so caches without the hook (or test doubles) are skipped. + materialized = getattr(type(cache), "materialized", None) + if materialized is not None and (result := materialized(cache, self)) is not None: + return result argsb = [a.build(cache) if isinstance(a, TimeSeriesExpression) else a for a in self.args] kwargsb = { k: a.build(cache) if isinstance(a, TimeSeriesExpression) else a diff --git a/src/impulse_query_engine/analyze/query/optimizer/__init__.py b/src/impulse_query_engine/analyze/query/optimizer/__init__.py new file mode 100644 index 00000000..4dad882a --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/__init__.py @@ -0,0 +1,16 @@ +from .plan import PushdownPlan +from .pushdown import RAW_VARIANT, VARIANT_COL, Pushdown, PushdownContext, PushdownRule +from .query_optimizer import QueryOptimizer, default_rules +from .resample_rule import ResamplePushdownRule + +__all__ = [ + "RAW_VARIANT", + "VARIANT_COL", + "Pushdown", + "PushdownContext", + "PushdownPlan", + "PushdownRule", + "QueryOptimizer", + "ResamplePushdownRule", + "default_rules", +] diff --git a/src/impulse_query_engine/analyze/query/optimizer/plan.py b/src/impulse_query_engine/analyze/query/optimizer/plan.py new file mode 100644 index 00000000..5e30299e --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/plan.py @@ -0,0 +1,143 @@ +"""The outcome of query optimization: which subtrees Spark computes ahead of the UDF.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field + +import pyspark.sql.functions as F +from pyspark.sql import Column + +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesExpression + +from .pushdown import RAW_VARIANT, Pushdown, PushdownRule + + +@dataclass(frozen=True, eq=False) +class PushdownPlan: + """Pushdowns selected for one solve call, shipped to the solver and the UDF. + + Variant ``0`` is the raw channel data; variants ``1..n`` are derived series, + one per distinct :attr:`Pushdown.variant_key`. + + Parameters + ---------- + rules : tuple of PushdownRule + Rules used to recognize pushed-down nodes during evaluation. + variant_ids : Mapping[tuple, int] + Variant key to variant id. + pushdowns_by_variant : Mapping[int, Pushdown] + Variant id to a representative pushdown (rule and parameters). + variants_by_selector : Mapping[int, tuple of int] + Selector id to the variant ids derived from it. + raw_selector_ids : frozenset of int + Selectors whose raw channel data is still required. + """ + + rules: tuple[PushdownRule, ...] = () + variant_ids: Mapping[tuple, int] = field(default_factory=dict) + pushdowns_by_variant: Mapping[int, Pushdown] = field(default_factory=dict) + variants_by_selector: Mapping[int, tuple[int, ...]] = field(default_factory=dict) + raw_selector_ids: frozenset[int] = frozenset() + + @staticmethod + def build( + rules: Iterable[PushdownRule], + pushdowns: Iterable[Pushdown], + selector_ids: Iterable[int], + ) -> PushdownPlan: + """Assign variant ids to *pushdowns*, in discovery order. + + Parameters + ---------- + rules : Iterable[PushdownRule] + The optimizer's rules. + pushdowns : Iterable[Pushdown] + Accepted pushdowns; each one's selector is pushed down entirely. + selector_ids : Iterable[int] + Ids of all selectors in the solve. + """ + variant_ids: dict[tuple, int] = {} + pushdowns_by_variant: dict[int, Pushdown] = {} + variants_by_selector: dict[int, list[int]] = {} + for pushdown in pushdowns: + key = pushdown.variant_key + if key not in variant_ids: + variant_ids[key] = len(variant_ids) + 1 + pushdowns_by_variant[variant_ids[key]] = pushdown + variants = variants_by_selector.setdefault(pushdown.selector.selector_id, []) + if variant_ids[key] not in variants: + variants.append(variant_ids[key]) + return PushdownPlan( + rules=tuple(rules), + variant_ids=variant_ids, + pushdowns_by_variant=pushdowns_by_variant, + variants_by_selector={sid: tuple(v) for sid, v in variants_by_selector.items()}, + raw_selector_ids=frozenset(selector_ids) - variants_by_selector.keys(), + ) + + @property + def is_empty(self) -> bool: + """``True`` when nothing is pushed down.""" + return not self.variant_ids + + @property + def pushed_selector_ids(self) -> frozenset[int]: + """Selectors whose raw channel data is replaced by derived variants.""" + return frozenset(self.variants_by_selector) + + def variants_by_rule(self) -> list[tuple[PushdownRule, dict[int, Pushdown]]]: + """Group the variants by the rule that computes them.""" + groups: dict[str, tuple[PushdownRule, dict[int, Pushdown]]] = {} + for variant, pushdown in self.pushdowns_by_variant.items(): + rule = pushdown.rule + groups.setdefault(rule.name, (rule, {}))[1][variant] = pushdown + return list(groups.values()) + + def lookup(self, node: TimeSeriesExpression) -> tuple[Pushdown, int] | None: + """Return ``(pushdown, variant)`` when *node* is a pushed-down subtree.""" + for rule in self.rules: + pushdown = rule.match(node) + if pushdown is None: + continue + variant = self.variant_ids.get(pushdown.variant_key) + if variant in self.variants_by_selector.get(pushdown.selector.selector_id, ()): + return pushdown, variant + return None + return None + + def channel_variants(self, selector_ids: Column) -> Column: + """Variant ids a matched channel ships, given its ``selector_ids`` array. + + A channel ships only its raw data (``[0]``) unless every one of its + selectors is pushed down; then it ships the derived variants instead. + """ + ids = F.transform(selector_ids, lambda sid: sid.cast("long")) + pushed = sorted(self.variants_by_selector) + needs_raw = F.exists(ids, lambda sid: ~sid.isin(pushed)) | (F.size(ids) == 0) + mapping = F.create_map( + *[ + item + for sid in pushed + for item in ( + F.lit(sid).cast("long"), + F.array(*[F.lit(v) for v in self.variants_by_selector[sid]]), + ) + ] + ) + derived = F.array_distinct( + F.flatten(F.transform(ids, lambda sid: F.try_element_at(mapping, sid))) + ) + return F.when(F.coalesce(needs_raw, F.lit(True)), F.array(F.lit(RAW_VARIANT))).otherwise( + derived + ) + + def describe(self) -> str: + """One-line summary of the plan for logging.""" + variants = ", ".join( + f"{vid}={pushdown.describe()}" for vid, pushdown in self.pushdowns_by_variant.items() + ) + return ( + f"variants [{variants}]; pushed selectors {sorted(self.pushed_selector_ids)}; " + f"raw selectors {sorted(self.raw_selector_ids)}" + ) diff --git a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py new file mode 100644 index 00000000..7ac8bb16 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -0,0 +1,148 @@ +"""Building blocks of Spark pushdown rules.""" + +from __future__ import annotations + +import abc +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar + +import pyspark.sql.types as T +from pyspark.sql import DataFrame + +from impulse_query_engine.analyze.metadata.time_series_expression import ( + TimeSeriesExpression, + TimeSeriesSelector, +) +from impulse_query_engine.model.series.sample_series import SampleSeries + +if TYPE_CHECKING: + import pandas as pd + + from impulse_query_engine.analyze.query.solvers.default_solver import TimeSeriesCache + from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig + +#: Column tagging each channel-data row with the series variant it belongs to. +VARIANT_COL = "__impulse_variant" +#: Variant id of the unmodified (raw) channel data. +RAW_VARIANT = 0 + + +@dataclass(frozen=True) +class PushdownContext: + """Plan-time facts about the channel data a pushdown would run on. + + Parameters + ---------- + tstart_type, tend_type : pyspark.sql.types.DataType + Spark types of the channel-data ``tstart`` / ``tend`` columns. + """ + + tstart_type: T.DataType + tend_type: T.DataType + + @staticmethod + def from_schema(schema: T.StructType, config: SolverConfig) -> PushdownContext: + """Build the context from the column-mapped channel-data schema.""" + return PushdownContext( + tstart_type=schema[config.tstart_col].dataType, + tend_type=schema[config.tend_col].dataType, + ) + + +@dataclass(frozen=True, eq=False) +class Pushdown: + """A subtree that a :class:`PushdownRule` computes natively in Spark. + + Expression nodes overload ``==``, so pushdowns compare by identity; use + :attr:`variant_key` to identify equivalent pushdowns. + + Parameters + ---------- + rule : PushdownRule + The rule that matched the subtree. + selector : TimeSeriesSelector + The physical channel selector the subtree reads. + params : tuple of (str, Any) + Canonical, hashable operation parameters (e.g. ``(("sample_rate", 1e6),)``). + consumed : tuple of TimeSeriesSelector + Every selector occurrence inside the subtree (its ``get_selectors()``). + """ + + rule: PushdownRule + selector: TimeSeriesSelector + params: tuple[tuple[str, Any], ...] + consumed: tuple[TimeSeriesSelector, ...] + + @property + def variant_key(self) -> tuple: + """Identity of the derived series, independent of the input channel.""" + return (self.rule.name, self.params) + + def describe(self) -> str: + """Human-readable ``rule(param=value, ...)`` label.""" + params = ", ".join(f"{k}={v}" for k, v in self.params) + return f"{self.rule.name}({params})" + + +class PushdownRule(abc.ABC): + """Recognizes a subtree and computes it in Spark ahead of the pandas UDF. + + A rule's derived rows use the channel-data shape + ``(container_id, channel_id, variant, tstart, tend, value)``; :meth:`decode` + turns one container/channel/variant slice of them back into the object the + original subtree evaluates to. The result must equal evaluating the subtree + in pandas on the same channel data. + """ + + #: Unique rule name, part of every :attr:`Pushdown.variant_key`. + name: ClassVar[str] + + @abc.abstractmethod + def match(self, node: TimeSeriesExpression) -> Pushdown | None: + """Return the pushdown for *node*, or ``None`` if the rule does not apply. + + Purely syntactic: it runs on the driver while planning and again inside + the UDF to recognize the node, so it must not depend on data or schemas. + """ + + def supports(self, pushdown: Pushdown, context: PushdownContext) -> bool: + """Whether *pushdown* can run on channel data described by *context*.""" + return True + + def result_type(self, pushdown: Pushdown) -> type: + """Core-model class :meth:`decode` returns for *pushdown*.""" + return SampleSeries + + @abc.abstractmethod + def to_spark( + self, rows: DataFrame, variants: Mapping[int, Pushdown], config: SolverConfig + ) -> DataFrame: + """Replace the rows of this rule's *variants* with their derived rows. + + *rows* is read once and shared by all rules, so implementations must + stay a single lineage (no self-joins or unions of *rows*) and keep its + container partitioning (windows partitioned by container id, channel id + and :data:`VARIANT_COL` reuse it). + + Parameters + ---------- + rows : pyspark.sql.DataFrame + All channel-data rows of the solve, tagged with :data:`VARIANT_COL` + (one copy per variant a channel ships) and partitioned by container. + variants : Mapping[int, Pushdown] + This rule's variant ids, each with a representative pushdown carrying + its parameters. + config : SolverConfig + Provides the internal column names. + + Returns + ------- + pyspark.sql.DataFrame + *rows* with the same columns and types, the rows of *variants* + replaced by derived rows and all other rows unchanged. + """ + + def decode(self, rows: pd.DataFrame, pushdown: Pushdown, cache: TimeSeriesCache) -> Any: + """Build the subtree's result from one container/channel/variant slice.""" + return cache.to_sample_series(rows, uses_alias=pushdown.selector.uses_alias) diff --git a/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py new file mode 100644 index 00000000..810153e3 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py @@ -0,0 +1,103 @@ +"""Rule-based optimizer choosing the subtrees Spark computes ahead of the pandas UDF.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Iterable, Sequence +from typing import Any + +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesExpression + +from .plan import PushdownPlan +from .pushdown import Pushdown, PushdownContext, PushdownRule +from .resample_rule import ResamplePushdownRule +from .traversal import discover + + +def default_rules() -> tuple[PushdownRule, ...]: + """All built-in pushdown rules.""" + return (ResamplePushdownRule(),) + + +class QueryOptimizer: + """Plans which subtrees of a solve's selections are computed natively in Spark. + + All selections of a solve share the raw data of their physical channels, so + a subtree is pushed down only when **every** occurrence of its input selector, + across all selections, sits inside a pushed-down subtree. Then the raw rows + of that selector are no longer needed and are replaced by the (smaller) + derived rows. Selectors used anywhere else keep their raw data, and their + subtrees are evaluated in pandas as before. + + Parameters + ---------- + rules : Sequence[PushdownRule], optional + Rules to apply, in priority order. Defaults to :func:`default_rules`. + """ + + def __init__(self, rules: Sequence[PushdownRule] | None = None): + self.rules: tuple[PushdownRule, ...] = ( + tuple(rules) if rules is not None else default_rules() + ) + + def optimize(self, selections: Iterable[Any], context: PushdownContext) -> PushdownPlan: + """Build the pushdown plan for the *selections* of one solve call. + + Parameters + ---------- + selections : Iterable + The query's selections. + context : PushdownContext + Plan-time facts about the channel data. + + Returns + ------- + PushdownPlan + The plan; empty when nothing can be pushed down. + """ + selections = list(selections) + if not all(isinstance(s, TimeSeriesExpression) for s in selections): + return PushdownPlan(rules=self.rules) + + total: Counter[int] = Counter() + candidates: list[tuple[TimeSeriesExpression, Pushdown]] = [] + for selection in selections: + expected = Counter(s.selector_id for s in selection.get_selectors()) + total.update(expected) + found, leaves = discover(selection, self.rules) + walked = Counter(s.selector_id for s in leaves) + walked.update(s.selector_id for _, p in found for s in p.consumed) + # A walk that disagrees with get_selectors() missed or double-counted part + # of the tree; its candidates can't be trusted to cover all occurrences. + if walked == expected: + candidates.extend(found) + + candidates = [ + (node, pushdown) + for node, pushdown in candidates + if pushdown.rule.supports(pushdown, context) and _has_result_type(node, pushdown) + ] + + while True: + consumed = Counter(s.selector_id for _, p in candidates for s in p.consumed) + fully_pushed = {sid for sid, count in consumed.items() if count == total[sid]} + kept = [ + (node, pushdown) + for node, pushdown in candidates + if all(s.selector_id in fully_pushed for s in pushdown.consumed) + ] + if len(kept) == len(candidates): + break + candidates = kept + + return PushdownPlan.build(self.rules, [p for _, p in candidates], total) + + +def _has_result_type(node: TimeSeriesExpression, pushdown: Pushdown) -> bool: + """Whether *node* evaluates to the type the rule decodes into.""" + try: + actual = node.evaluation_type() + except Exception: + # Only disqualifies the candidate; the query's own type inference reports errors. + return False + return actual is pushdown.rule.result_type(pushdown) diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py new file mode 100644 index 00000000..7c398b92 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -0,0 +1,221 @@ +"""Pushdown rule computing ``SampleSeries.resample`` natively in Spark.""" + +from __future__ import annotations + +import math +import numbers +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import pyspark.sql.functions as F +import pyspark.sql.types as T +from pyspark.sql import Column, DataFrame, Window + +from impulse_query_engine.analyze.metadata.time_series_expression import ( + SeriesType, + TimeSeriesExpression, + TimeSeriesOp, + TimeSeriesSelector, +) + +from .pushdown import VARIANT_COL, Pushdown, PushdownContext, PushdownRule + +if TYPE_CHECKING: + from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig + +_RATE = "sample_rate" +_DEFAULT_RATE = 1.0 +_TIME_TYPES = (T.ByteType, T.ShortType, T.IntegerType, T.LongType, T.DoubleType) + + +class ResamplePushdownRule(PushdownRule): + """Pushes ``channel.resample(sample_rate)`` on a physical channel down to Spark. + + Reproduces :meth:`SampleSeries.resample`: per continuous block of samples, a + regular grid starting at the block's first ``tstart`` (built like + ``numpy.arange``), each grid point holding the value of the last sample that + starts at or before it. Meant for downsampling: a rate finer than the native + sample spacing yields more rows than the raw channel. + + Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. + numpy builds that fuse this into a multiply-add can differ by one ulp for + fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral + grids are exact. + """ + + name = "resample" + + def match(self, node: TimeSeriesExpression) -> Pushdown | None: + if ( + type(node) is not TimeSeriesOp + or node.optype != "cls" + or not isinstance(node.operation, str) + or node.operation != "resample" + ): + return None + if not node.args or not isinstance(node.args[0], TimeSeriesSelector): + return None + selector = node.args[0] + if selector.series_type is not SeriesType.CONTINUOUS: + return None + rate = _sample_rate(node.args[1:], node.kwargs) + if rate is None: + return None + return Pushdown( + rule=self, selector=selector, params=((_RATE, rate),), consumed=(selector,) + ) + + def supports(self, pushdown: Pushdown, context: PushdownContext) -> bool: + time_types = (context.tstart_type, context.tend_type) + if not all(isinstance(t, _TIME_TYPES) for t in time_types): + return False + # Integral time columns receive the grid cast back from double; an integral + # rate keeps every grid point integral, so that cast is exact. + rate = dict(pushdown.params)[_RATE] + return rate.is_integer() or not any(isinstance(t, T.IntegralType) for t in time_types) + + def to_spark( + self, rows: DataFrame, variants: Mapping[int, Pushdown], config: SolverConfig + ) -> DataFrame: + cid, ch = config.container_id_col, config.channel_id_col + ts, te, val = config.tstart_col, config.tend_col, config.value_col + key = [cid, ch, VARIANT_COL] + rate = F.try_element_at( + F.create_map( + *[ + item + for vid, pushdown in sorted(variants.items()) + for item in (F.lit(vid), F.lit(dict(pushdown.params)[_RATE])) + ] + ), + F.col(VARIANT_COL), + ) + + # All per-channel state comes from windows sharing one partitioning and order + # (the pandas cache's tstart order), so the rows are sorted once and the + # incoming container partitioning is kept. + ordered = Window.partitionBy(*key).orderBy(ts, te) + running = ordered.rowsBetween(Window.unboundedPreceding, Window.currentRow) + whole = ordered.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) + ts_d, te_d = F.col("_rs_ts"), F.col("_rs_te") + prev_te = F.lag(te_d).over(ordered) + df = ( + rows.withColumn("_rs_ts", F.col(ts).cast("double")) + .withColumn("_rs_te", F.col(te).cast("double")) + .withColumn("_rs_break", (prev_te.isNull() | ((ts_d - prev_te) != 0)).cast("long")) + .withColumn("_rs_n", F.count(F.lit(1)).over(whole)) + .withColumn("_rs_last_ts", F.last(ts_d).over(whole)) + .withColumn("_rs_last_te", F.last(te_d).over(whole)) + .withColumn("_rs_block", F.sum("_rs_break").over(running)) + .withColumn( + "_rs_block_end", F.coalesce(F.lead("_rs_break").over(ordered), F.lit(1)) == 1 + ) + ) + # Block bounds as seen on the block's last row: its first tstart, and the + # largest tend so far within the block (blocks are numbered in order). + df = ( + df.withColumn( + "_rs_t_min", + F.last(F.when(F.col("_rs_break") == 1, ts_d), ignorenulls=True).over(running), + ) + .withColumn( + "_rs_t_max", + F.max(F.struct(F.col("_rs_block"), _ignore_nan(te_d).alias("te"))) + .over(running) + .getField("te"), + ) + .withColumn("_rs_nblk", F.max("_rs_block").over(whole)) + ) + + # numpy.arange(t_min, t_max + rate, rate): length ceil(((t_max + rate) - t_min) / rate), + # filled as t_min, t_min + rate, then t_min + k * ((t_min + rate) - t_min). + t_min, t_max = F.col("_rs_t_min"), F.col("_rs_t_max") + length = F.ceil(((t_max + rate) - t_min) / rate) + delta = (t_min + rate) - t_min + # The last block keeps a grid point at t_max when the series ends with a + # zero-duration sample. + inclusive = (F.col("_rs_block") == F.col("_rs_nblk")) & ( + F.col("_rs_last_ts") == F.col("_rs_last_te") + ) + resampled = F.col(VARIANT_COL).isin(sorted(variants)) & (F.col("_rs_n") >= 2) + grid_points = F.filter( + F.transform( + F.sequence(F.lit(0).cast("long"), length - 1), + lambda k: F.when(k == 0, t_min) + .when(k == 1, t_min + rate) + .otherwise(t_min + k.cast("double") * delta), + ), + lambda g: (g < t_max) | (inclusive & (g == t_max)), + ) + df = df.withColumn( + "_rs_grid", + F.when(resampled & F.col("_rs_block_end") & (length > 0), grid_points), + ) + + # One output element per row: the row itself for series this rule leaves + # alone (kind 2), otherwise the sample (kind 0) plus, on a block's last row, + # the block's grid points (kind 1). Fewer than two samples: returned as-is. + def element(kind: int, t: Column, tend: Column) -> Column: + return F.struct( + F.lit(kind).alias("kind"), t.alias("t"), tend.cast("double").alias("tend") + ) + + grid = F.transform( + F.col("_rs_grid"), + lambda g, i: element( + 1, g, F.coalesce(F.try_element_at(F.col("_rs_grid"), i + 2), t_max) + ), + ) + no_tend = F.lit(None).cast("double") + elements = ( + F.when(~resampled, F.array(element(2, ts_d, no_tend))) + .when( + F.col("_rs_grid").isNotNull(), F.concat(F.array(element(0, ts_d, no_tend)), grid) + ) + .otherwise(F.array(element(0, ts_d, no_tend))) + ) + df = df.withColumn("_rs_e", F.explode(elements)) + kind, t = F.col("_rs_e.kind"), F.col("_rs_e.t") + + # Previous-value lookup: each grid point takes the value of the last sample (in + # channel order) starting at or before it. Samples sort before grid points at + # equal times; the struct keeps null values distinguishable from "no sample". + as_of = ( + Window.partitionBy(*key) + .orderBy(t, kind, ts, te) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + previous = F.last(F.when(kind == 0, F.struct(F.col(val).alias("value"))), ignorenulls=True) + df = df.withColumn("_rs_previous", previous.over(as_of)).where(kind != 0) + + is_grid = kind == 1 + derived = { + ts: t.cast(rows.schema[ts].dataType), + te: F.col("_rs_e.tend").cast(rows.schema[te].dataType), + # Beyond the last sample's end the pandas interpolation yields NaN. + val: F.when(t > F.col("_rs_last_te"), F.lit(None)) + .otherwise(F.col("_rs_previous.value")) + .cast(rows.schema[val].dataType), + } + return df.select( + *[ + F.when(is_grid, derived[c]).otherwise(F.col(c)).alias(c) if c in derived else c + for c in rows.columns + ] + ) + + +def _sample_rate(args: tuple, kwargs: dict[str, Any]) -> float | None: + """The resample rate of a ``resample(...)`` call, or ``None`` if unsupported.""" + if len(args) > 1 or set(kwargs) - {_RATE} or (args and kwargs): + return None + rate = args[0] if args else kwargs.get(_RATE, _DEFAULT_RATE) + if isinstance(rate, bool) or not isinstance(rate, numbers.Real): + return None + rate = float(rate) + return rate if math.isfinite(rate) and rate > 0 else None + + +def _ignore_nan(col: Column) -> Column: + """Map NaN to null so ``min`` / ``max`` behave like ``numpy.nanmin`` / ``nanmax``.""" + return F.when(~F.isnan(col), col) diff --git a/src/impulse_query_engine/analyze/query/optimizer/traversal.py b/src/impulse_query_engine/analyze/query/optimizer/traversal.py new file mode 100644 index 00000000..cf6d2f72 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/traversal.py @@ -0,0 +1,82 @@ +"""Generic walk over TSAL expression trees.""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator, Sequence + +from impulse_query_engine.analyze.metadata.time_series_expression import ( + TimeSeriesAliasSelector, + TimeSeriesExpression, + TimeSeriesOp, + TimeSeriesSelector, +) + +from .pushdown import Pushdown, PushdownRule + + +def child_expressions(node: TimeSeriesExpression) -> Iterator[TimeSeriesExpression]: + """Yield the direct sub-expressions of *node*. + + Operations yield their operands; other node types (aggregations, events, + calculated channels, …) yield every expression held in their attributes, + including inside lists, tuples and dicts. The walk is best-effort: the + optimizer checks the leaves it finds against ``get_selectors()``. + """ + if isinstance(node, TimeSeriesSelector): + return + if isinstance(node, TimeSeriesOp): + # Covers TimeSeriesUDF, whose ``optype`` attribute holds its first operand. + values: Iterable = (*node.args, *node.kwargs.values()) + elif isinstance(node, TimeSeriesAliasSelector): + values = node._aliases + else: + values = (value for key, value in vars(node).items() if key != "_alias") + yield from _expressions_in(values) + + +def _expressions_in(values: Iterable) -> Iterator[TimeSeriesExpression]: + for value in values: + if isinstance(value, TimeSeriesExpression): + yield value + elif isinstance(value, list | tuple): + yield from _expressions_in(value) + elif isinstance(value, dict): + yield from _expressions_in(value.values()) + + +def match_node(node: TimeSeriesExpression, rules: Sequence[PushdownRule]) -> Pushdown | None: + """Return the first rule's pushdown for *node*, or ``None``.""" + for rule in rules: + pushdown = rule.match(node) + if pushdown is not None: + return pushdown + return None + + +def discover( + root: TimeSeriesExpression, rules: Sequence[PushdownRule] +) -> tuple[list[tuple[TimeSeriesExpression, Pushdown]], list[TimeSeriesSelector]]: + """Find the maximal subtrees of *root* that a rule can push down. + + The walk does not descend into a matched subtree and keeps no visited set, so + a node referenced twice is counted twice, like ``get_selectors()`` does. + + Returns + ------- + tuple + ``(candidates, leaves)``: the matched ``(node, pushdown)`` pairs and the + selectors reached outside of them. + """ + candidates: list[tuple[TimeSeriesExpression, Pushdown]] = [] + leaves: list[TimeSeriesSelector] = [] + stack: list[TimeSeriesExpression] = [root] + while stack: + node = stack.pop() + pushdown = match_node(node, rules) + if pushdown is not None: + candidates.append((node, pushdown)) + elif isinstance(node, TimeSeriesSelector): + leaves.append(node) + else: + stack.extend(child_expressions(node)) + return candidates, leaves diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index f9e6d7b0..bc5aa214 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from collections.abc import Iterable from functools import partial from typing import TYPE_CHECKING, Any @@ -16,6 +17,13 @@ SeriesValueType, TimeSeriesExpression, ) +from impulse_query_engine.analyze.query.optimizer import ( + RAW_VARIANT, + VARIANT_COL, + PushdownContext, + PushdownPlan, + QueryOptimizer, +) from impulse_query_engine.model.series.points_in_time_series import PointsInTimeSeries from impulse_query_engine.model.series.sample_series import SampleSeries @@ -29,6 +37,8 @@ if TYPE_CHECKING: from impulse_query_engine.measurement_db import MeasurementDB +logger = logging.getLogger(__name__) + class TimeSeriesCache(SeriesCache): def __init__( @@ -37,6 +47,7 @@ def __init__( col_map: dict[str, str], container_tag_cols=None, container_metric_cols=None, + plan: PushdownPlan | None = None, ): """ Initialize the TimeSeriesCache. @@ -52,7 +63,9 @@ def __init__( the caller's frame is reordered by ``(cid, ch, ts)`` and its index reset. The sort makes each channel's rows contiguous, which the constructor exploits to build a ``(cid, ch)`` → - ``(start, stop)`` range index for :meth:`load_blob`. + ``(start, stop)`` range index for :meth:`load_blob`. When *pdf* + carries the pushdown variant column, rows are additionally keyed + by variant (raw data is variant ``0``). col_map : dict[str, str] Mapping with keys ``"cid"``, ``"ch"``, ``"ts"``, ``"te"``, ``"val"``, ``"conv"`` to the actual column names in *pdf*. The @@ -65,6 +78,9 @@ def __init__( container_metric_cols : list of str, optional Columns on *pdf* holding container-level metric values, exposed via :attr:`container_metrics`. These are constant per container. + plan : PushdownPlan, optional + Pushdowns whose derived rows *pdf* carries; :meth:`materialized` + serves the matching subtrees from them. """ self._cid_col = col_map["cid"] self._ch_col = col_map["ch"] @@ -78,28 +94,32 @@ def __init__( self._value_string_col = col_map.get("value_string") self._container_tag_cols = list(container_tag_cols or []) self._container_metric_cols = list(container_metric_cols or []) + self._plan = plan + has_variants = VARIANT_COL in pdf.columns # *pdf* holds channel data for a whole container, so avoid creating unnecessary copies of the data. meta_cols = [ - c for c in pdf.columns if c not in (self._ts_col, self._te_col, self._val_col) + c + for c in pdf.columns + if c not in (self._ts_col, self._te_col, self._val_col, VARIANT_COL) ] meta_source = pdf if "selector_ids" in pdf.columns: meta_source = pdf.loc[pdf["selector_ids"].notna()] first_rows = ~meta_source.duplicated(subset=[self._cid_col, self._ch_col]) self.mdf = meta_source.loc[first_rows, meta_cols].reset_index(drop=True) - pdf.sort_values( - [self._cid_col, self._ch_col, self._ts_col], inplace=True, ignore_index=True - ) + group_cols = [self._cid_col, self._ch_col] + ([VARIANT_COL] if has_variants else []) + pdf.sort_values([*group_cols, self._ts_col], inplace=True, ignore_index=True) self.pdf = pdf - # Each (cid, ch) group is a contiguous block after the sort, so its + # Each (cid, ch, variant) group is a contiguous block after the sort, so its # row positions collapse to a (start, stop) range self._ranges = { - key: (int(positions[0]), int(positions[-1]) + 1) - for key, positions in pdf.groupby( - [self._cid_col, self._ch_col], sort=False - ).indices.items() + (key if has_variants else (*key, RAW_VARIANT)): ( + int(positions[0]), + int(positions[-1]) + 1, + ) + for key, positions in pdf.groupby(group_cols, sort=False).indices.items() } def resolve(self, selection): @@ -169,8 +189,7 @@ def load_blob( ------- SampleSeries or PointsInTimeSeries """ - lo, hi = self._ranges.get((mid, cid), (0, 0)) - s = self.pdf.iloc[lo:hi] + s = self._slice(mid, cid, RAW_VARIANT) if series_type == SeriesType.POINTS_IN_TIME: # The selector's declared value type picks the column; the constructor @@ -185,12 +204,72 @@ def load_blob( ) return PointsInTimeSeries(s[self._ts_col], s[value_col], value_type=value_type) - values = s[self._val_col] - if self._has_conversion and len(s) > 0 and uses_alias: - factor = s[self._conv_col].iloc[0] + return self.to_sample_series(s, uses_alias=uses_alias) + + def materialized(self, expression): + """Serve a pushed-down subtree from its derived rows. + + Returns ``None`` (evaluate normally) when *expression* is not pushed down + or its channel carries no derived rows in this container. + + Parameters + ---------- + expression : TimeSeriesExpression + The node about to be built. + + Returns + ------- + Any or None + The rule-decoded result, or ``None``. + """ + if self._plan is None: + return None + hit = self._plan.lookup(expression) + if hit is None: + return None + pushdown, variant = hit + candidates = self.resolve(pushdown.selector) + if len(candidates) == 0: + return None + rows = self._slice( + candidates[self._cid_col].iloc[0], candidates[self._ch_col].iloc[0], variant + ) + if len(rows) == 0: + return None + return pushdown.rule.decode(rows, pushdown, self) + + def to_sample_series(self, rows: pd.DataFrame, uses_alias: bool = False) -> SampleSeries: + """Build a :class:`SampleSeries` from channel-data *rows* of one channel. + + Applies the per-channel conversion factor when *uses_alias* is ``True`` + (see :meth:`load_blob`). + """ + values = rows[self._val_col] + if self._has_conversion and len(rows) > 0 and uses_alias: + factor = rows[self._conv_col].iloc[0] if pd.notna(factor): values = values * factor - return SampleSeries(s[self._ts_col], s[self._te_col], values) + return SampleSeries(rows[self._ts_col], rows[self._te_col], values) + + @property + def tstart_col(self) -> str: + """Name of the ``tstart`` column of the cached rows.""" + return self._ts_col + + @property + def tend_col(self) -> str: + """Name of the ``tend`` column of the cached rows.""" + return self._te_col + + @property + def value_col(self) -> str: + """Name of the numeric value column of the cached rows.""" + return self._val_col + + def _slice(self, mid, cid, variant: int) -> pd.DataFrame: + """Rows of one container/channel/variant (empty when absent).""" + lo, hi = self._ranges.get((mid, cid, variant), (0, 0)) + return self.pdf.iloc[lo:hi] @property def container_tags(self) -> dict: @@ -256,6 +335,12 @@ class DefaultSolver(QuerySolver): ``RawEncoder.RLE`` (default) run-length encodes equal-valued runs; ``RawEncoder.INTERVAL`` only derives ``tend`` and drops exact duplicates. Only consulted when ``is_raw_data`` is ``True``. + enable_pushdown : bool, optional + Whether to compute supported expression subtrees (e.g. + ``channel.resample(rate)``) natively in Spark before the pandas UDF, so + fewer rows are shipped to it (see :class:`QueryOptimizer`). A channel + is only pushed down when every use of it in the solve is. Default + ``False``. """ def __init__( @@ -265,12 +350,14 @@ def __init__( is_raw_data: bool = False, drop_implausible_data: bool = False, raw_encoder: RawEncoder = RawEncoder.RLE, + enable_pushdown: bool = False, ): super().__init__(config=config) self.spark = spark self.is_raw_data = is_raw_data self.drop_implausible_data: bool = drop_implausible_data self.raw_encoder: RawEncoder = raw_encoder + self.optimizer: QueryOptimizer | None = QueryOptimizer() if enable_pushdown else None # In RAW mode an is_plausible channels filter drops samples before raw # encoding, bridging intervals instead of splitting them. Reject it. self.config.reject_implausible_channels_filter_in_raw(is_raw=self.is_raw_data) @@ -1007,6 +1094,7 @@ def _solve_udf( col_map: dict[str, str], container_tag_cols=None, container_metric_cols=None, + plan: PushdownPlan | None = None, ) -> pd.DataFrame: """ UDF to solve for a single container by applying selections. @@ -1024,6 +1112,8 @@ def _solve_udf( container_metric_cols : list of str, optional Columns on *pdf* holding container-level metric values, exposed to expressions via ``cache.container_metrics``. + plan : PushdownPlan, optional + Pushdowns whose derived rows *pdf* carries. Returns ------- @@ -1037,6 +1127,7 @@ def _solve_udf( col_map=col_map, container_tag_cols=container_tag_cols, container_metric_cols=container_metric_cols, + plan=plan, ) for s in selections: res = s.build(cache) @@ -1075,7 +1166,7 @@ def solve(self, query, channels_df, selections, dtypes) -> DataFrame: DataFrame containing results for each container. """ col_map = self.config.col_map - q, joined_df, container_count = self._prepare_channels_join(query, channels_df) + q, joined_df, container_count, plan = self._prepare_channels_join(query, channels_df) container_tag_cols, container_metric_cols = self._container_meta_cols(query) schema = self._build_solve_output_schema(q, selections, dtypes) @@ -1086,6 +1177,7 @@ def solve(self, query, channels_df, selections, dtypes) -> DataFrame: col_map=col_map, container_tag_cols=container_tag_cols, container_metric_cols=container_metric_cols, + plan=plan, ), schema, F.PandasUDFType.GROUPED_MAP, @@ -1098,7 +1190,9 @@ def solve(self, query, channels_df, selections, dtypes) -> DataFrame: container_meta_cols=[*container_tag_cols, *container_metric_cols], ) - def _prepare_channels_join(self, query, channels_df) -> tuple[DataFrame, DataFrame, int]: + def _prepare_channels_join( + self, query, channels_df + ) -> tuple[DataFrame, DataFrame, int, PushdownPlan | None]: """Shared prelude for :meth:`solve` and :meth:`solve_calculated_channels`. Applies optional per-channel unit conversion, reads and column-maps the @@ -1110,6 +1204,11 @@ def _prepare_channels_join(self, query, channels_df) -> tuple[DataFrame, DataFra *channels_df* (attached by :meth:`attach_container_metadata`) ride through the broadcast join into ``joined_df``. + With pushdown enabled, subtrees chosen by the optimizer are computed on + the channel data before the join (see :meth:`_apply_pushdowns`), and + ``joined_df`` carries their derived rows instead of the raw rows they + replace. + Parameters ---------- query : QueryBuilder @@ -1119,11 +1218,12 @@ def _prepare_channels_join(self, query, channels_df) -> tuple[DataFrame, DataFra Returns ------- - tuple[DataFrame, DataFrame, int] - ``(channels_q, joined_df, container_count)`` — the column-mapped + tuple[DataFrame, DataFrame, int, PushdownPlan or None] + ``(channels_q, joined_df, container_count, plan)`` — the column-mapped channel-data DataFrame (authoritative for output ``container_id`` / - ``channel_id`` types), the join fed to the grouped-map UDF, and the - number of distinct containers. + ``channel_id`` types), the join fed to the grouped-map UDF, the + number of distinct containers, and the pushdown plan the UDF must + receive (``None`` when nothing is pushed down). """ source_unit_col = self.config.source_unit_col target_unit_col = self.config.target_unit_col @@ -1168,12 +1268,59 @@ def _prepare_channels_join(self, query, channels_df) -> tuple[DataFrame, DataFra if DefaultSolver._query_contains_poi_selections(query.selections): q = self._union_poi_channel_data(query, q) - joined_df = q.join( + container_count = channels_df.select(self.config.container_id_col).distinct().count() + plan = self._plan_pushdowns(query, q) if container_count > 0 else None + channel_data = q + if plan is not None: + channel_data = self._apply_pushdowns(q, channels_df, plan, container_count) + + joined_df = channel_data.join( F.broadcast(channels_df), on=[self.config.container_id_col, self.config.channel_id_col], ) - container_count = channels_df.select(self.config.container_id_col).distinct().count() - return q, joined_df, container_count + return q, joined_df, container_count, plan + + def _plan_pushdowns(self, query, channels_q: DataFrame) -> PushdownPlan | None: + """Plan the pushdowns for *query*, or ``None`` when disabled or nothing applies.""" + if self.optimizer is None: + return None + plan = self.optimizer.optimize( + query.selections, PushdownContext.from_schema(channels_q.schema, self.config) + ) + if plan.is_empty: + return None + logger.info("Spark pushdown plan: %s", plan.describe()) + return plan + + def _apply_pushdowns( + self, + channels_q: DataFrame, + channels_df: DataFrame, + plan: PushdownPlan, + container_count: int, + ) -> DataFrame: + """Replace raw channel data with the derived rows of *plan* where possible. + + Each matched channel is tagged with the variants it ships + (:meth:`PushdownPlan.channel_variants`): its raw data (variant ``0``) + while any of its selectors still needs it, otherwise the derived + variants of its selectors, computed by their rules. The result carries + a :data:`VARIANT_COL` column and keeps the channel-data columns and types. + + The channel data is read once and partitioned by container exactly like + :meth:`_apply_grouped_map` does, so the rules' per-channel windows and + the grouped-map UDF share a single shuffle. + """ + key = [self.config.container_id_col, self.config.channel_id_col] + variants = channels_df.select( + *key, F.explode(plan.channel_variants(F.col("selector_ids"))).alias(VARIANT_COL) + ) + rows = channels_q.join(F.broadcast(variants), on=key).repartition( + container_count, self.config.container_id_col + ) + for rule, rule_variants in plan.variants_by_rule(): + rows = rule.to_spark(rows, rule_variants, self.config) + return rows @staticmethod def _query_contains_poi_selections(selections: Iterable[Any]) -> bool: @@ -1382,6 +1529,7 @@ def _solve_calculated_channels_udf( col_map: dict[str, str], container_tag_cols=None, container_metric_cols=None, + plan: PushdownPlan | None = None, ) -> pd.DataFrame: """ UDF to solve calculated channels for a single container. @@ -1409,6 +1557,8 @@ def _solve_calculated_channels_udf( container_metric_cols : list of str, optional Columns on *pdf* holding container-level metric values, exposed to expressions via ``cache.container_metrics``. + plan : PushdownPlan, optional + Pushdowns whose derived rows *pdf* carries. Returns ------- @@ -1421,6 +1571,7 @@ def _solve_calculated_channels_udf( col_map=col_map, container_tag_cols=container_tag_cols, container_metric_cols=container_metric_cols, + plan=plan, ) cid_col = col_map["cid"] ch_col = col_map["ch"] @@ -1472,7 +1623,7 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame Narrow DataFrame of calculated-channel samples. """ col_map = self.config.col_map - q, joined_df, container_count = self._prepare_channels_join(query, channels_df) + q, joined_df, container_count, plan = self._prepare_channels_join(query, channels_df) container_tag_cols, container_metric_cols = self._container_meta_cols(query) schema = self._build_calculated_channels_udf_schema(q) @@ -1483,6 +1634,7 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame col_map=col_map, container_tag_cols=container_tag_cols, container_metric_cols=container_metric_cols, + plan=plan, ), schema, F.PandasUDFType.GROUPED_MAP, diff --git a/src/impulse_query_engine/analyze/query/solvers/series_cache.py b/src/impulse_query_engine/analyze/query/solvers/series_cache.py index 33def4b8..fbf252a3 100644 --- a/src/impulse_query_engine/analyze/query/solvers/series_cache.py +++ b/src/impulse_query_engine/analyze/query/solvers/series_cache.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from typing import Any import pandas as pd @@ -70,6 +71,25 @@ def load_blob( """ pass + def materialized(self, expression) -> Any | None: + """Return the precomputed result of *expression*, if the cache holds one. + + Consulted by :meth:`TimeSeriesOp.build` before evaluating the node, so a + subtree computed ahead of time (e.g. pushed down to Spark) replaces its + in-memory evaluation. Default: no precomputed results. + + Parameters + ---------- + expression : TimeSeriesExpression + The node about to be built. + + Returns + ------- + Any or None + The precomputed result, or ``None`` to evaluate *expression* normally. + """ + return None + @property def container_tags(self) -> dict: """Container-level tag values available to UDFs. diff --git a/src/impulse_reporting/config/config_parser.py b/src/impulse_reporting/config/config_parser.py index 9bbde633..4744d963 100644 --- a/src/impulse_reporting/config/config_parser.py +++ b/src/impulse_reporting/config/config_parser.py @@ -388,6 +388,12 @@ class QueryEngine(BaseModel): When omitted, all default column names are used and no project/toolbox filtering is applied. + enable_pushdown : bool, default=False + Compute supported expression subtrees (currently + ``channel.resample(rate)`` on a physical channel) natively in Spark + before the pandas UDF, so fewer rows are shipped to it. A channel is + only pushed down when every use of it in a solve batch is. Meant for + downsampling. Notes ----- @@ -410,6 +416,7 @@ class QueryEngine(BaseModel): solver_config: SolverConfig | None = None max_channels_per_batch: int = 500 max_containers_per_batch: int | None = None + enable_pushdown: bool = False @model_validator(mode="before") @classmethod diff --git a/src/impulse_reporting/core/report.py b/src/impulse_reporting/core/report.py index 1a6b6f9f..111b9bf1 100644 --- a/src/impulse_reporting/core/report.py +++ b/src/impulse_reporting/core/report.py @@ -347,6 +347,7 @@ def create_solver(spark: SparkSession, config: ImpulseConfig) -> QuerySolver: is_raw_data=config.query_engine.data_type is DataType.RAW, drop_implausible_data=config.query_engine.drop_implausible_data, raw_encoder=config.query_engine.raw_encoder, + enable_pushdown=config.query_engine.enable_pushdown, ) case _: raise ValueError( diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/__init__.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py new file mode 100644 index 00000000..a50b2f89 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py @@ -0,0 +1,235 @@ +# pylint: disable=missing-function-docstring +"""Planning tests for the rule-based pushdown optimizer (no data involved).""" + +import math +import pickle + +import pyspark.sql.types as T +import pytest +from pyspark import cloudpickle + +from impulse_query_engine.analyze.metadata.tag_expression import TagSelector +from impulse_query_engine.analyze.metadata.time_series_expression import ( + SeriesType, + TimeSeriesExpression, + TimeSeriesSelector, +) +from impulse_query_engine.analyze.query.aggregations.histogram import HistogramDuration +from impulse_query_engine.analyze.query.aggregations.stats_aggregator import StatsAggregator +from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel +from impulse_query_engine.analyze.query.events.sequence_of_events_expression import ( + SequenceOfEventsExpression, +) +from impulse_query_engine.analyze.query.optimizer import ( + PushdownContext, + QueryOptimizer, + ResamplePushdownRule, +) +from impulse_query_engine.model.series.intervals import Intervals + +_LONG = PushdownContext(T.LongType(), T.LongType()) +_DOUBLE = PushdownContext(T.DoubleType(), T.DoubleType()) + + +def _channel(name: str, **kwargs) -> TimeSeriesSelector: + return TimeSeriesSelector(TagSelector("channel_name") == name, **kwargs) + + +def _plan(*selections, context=_LONG, rules=None): + return QueryOptimizer(rules).optimize(selections, context) + + +@pytest.fixture +def rpm(): + return _channel("rpm") + + +@pytest.fixture +def spd(): + return _channel("spd") + + +class TestPolicy: + def test_resample_only_selection_is_pushed(self, rpm): + plan = _plan(rpm.resample(1e6)) + assert not plan.is_empty + assert plan.pushed_selector_ids == {rpm.selector_id} + assert plan.raw_selector_ids == frozenset() + assert [p.describe() for p in plan.pushdowns_by_variant.values()] == [ + "resample(sample_rate=1000000.0)" + ] + + def test_raw_use_in_another_selection_blocks_pushdown(self, rpm): + assert _plan(rpm.resample(1e6), rpm.max()).is_empty + + def test_raw_use_in_the_same_expression_blocks_pushdown(self, rpm): + assert _plan(rpm.resample(1e6) + rpm).is_empty + + def test_other_channels_keep_raw_data(self, rpm, spd): + plan = _plan(rpm.resample(1e6), spd.max(), spd.resample(1e6)) + assert plan.pushed_selector_ids == {rpm.selector_id} + assert plan.raw_selector_ids == {spd.selector_id} + + def test_two_rates_give_two_variants(self, rpm): + plan = _plan(rpm.resample(1e6), rpm.resample(1e5)) + assert set(plan.variant_ids.values()) == {1, 2} + assert plan.variants_by_selector == {rpm.selector_id: (1, 2)} + + def test_equivalent_calls_share_one_variant(self, rpm): + plan = _plan(rpm.resample(1e6), rpm.resample(sample_rate=1_000_000)) + assert plan.variants_by_selector == {rpm.selector_id: (1,)} + + def test_equal_selectors_built_separately_are_one_selector(self): + plan = _plan(_channel("rpm").resample(1e6), _channel("rpm").max()) + assert plan.is_empty + + def test_node_shared_across_selections_and_branches(self, rpm): + resampled = rpm.resample(1e6) + plan = _plan(resampled + resampled, resampled.cumtrapz()) + assert plan.pushed_selector_ids == {rpm.selector_id} + + def test_non_expression_selection_gives_empty_plan(self, rpm): + assert _plan(rpm.resample(1e6), object()).is_empty + + +class TestDiscovery: + def test_inside_udf_argument(self, rpm): + # TimeSeriesUDF stores its first operand in ``optype`` too. + plan = _plan(rpm.resample(1e6).apply(lambda s: s)) + assert plan.pushed_selector_ids == {rpm.selector_id} + + def test_inside_calculated_channel(self, rpm): + plan = _plan(CalculatedChannel(rpm.resample(1e6) * 2, {"channel_name": "rpm_1hz"})) + assert plan.pushed_selector_ids == {rpm.selector_id} + + def test_inside_stats_aggregator_and_event(self, rpm, spd): + event = spd.resample(1e6) > 50 + stats = StatsAggregator( + input_expressions=[rpm.resample(1e6)], statistics=["max"], event_expression=event + ) + plan = _plan(stats, event) + assert plan.pushed_selector_ids == {rpm.selector_id, spd.selector_id} + + def test_inside_histogram_and_sequence_of_events(self, rpm, spd): + plan = _plan( + HistogramDuration(rpm.resample(1e6), bins=[0, 1000, 2000]), + SequenceOfEventsExpression([spd.resample(1e6) > 10, spd.resample(1e6) > 20]), + ) + assert plan.pushed_selector_ids == {rpm.selector_id, spd.selector_id} + + def test_aliased_selector(self): + alias = _channel("vehicle_speed", uses_alias=True) + assert _plan(alias.resample(1e6)).pushed_selector_ids == {alias.selector_id} + + def test_walk_disagreeing_with_get_selectors_keeps_raw_data(self, rpm): + # One selection over-counts a pushed subtree, another hides a raw use; the + # totals would balance, so only the per-selection guard keeps rpm raw. + resampled = rpm.resample(1e6) + over_counted = _Opaque({"a": resampled, "b": resampled}, reported=[rpm]) + hidden_raw_use = _Opaque({"a": rpm.resample(1e6)}, reported=[rpm, rpm]) + assert _plan(over_counted, hidden_raw_use).is_empty + + +class _Opaque(TimeSeriesExpression): + """Node whose attributes disagree with the selectors it reports.""" + + def __init__(self, attributes, reported): + super().__init__() + self.__dict__.update(attributes) + self._reported = lambda: list(reported) # not visible to the walk + + def build(self, cache): + return None + + def get_selectors(self): + return self._reported() + + def get_required_tag_exprs(self): + return set() + + def required_tags(self): + return set() + + def get_selector_expr(self): + return None + + def __str__(self): + return "_Opaque" + + +class TestResampleRule: + @pytest.mark.parametrize( + "make", + [ + lambda c: (c * 2).resample(1e6), + lambda c: c.resample(True), + lambda c: c.resample(0), + lambda c: c.resample(-1e6), + lambda c: c.resample(math.nan), + lambda c: c.resample(math.inf), + lambda c: c.resample(c), + lambda c: c.resample(1e6, 2), + lambda c: c.resample(1e6, sample_rate=1e6), + lambda c: c.resample(rate=1e6), + lambda c: c.rolling_average(1e6), + ], + ) + def test_unsupported_calls_are_not_matched(self, rpm, make): + assert ResamplePushdownRule().match(make(rpm)) is None + + def test_default_rate(self, rpm): + pushdown = ResamplePushdownRule().match(rpm.resample()) + assert pushdown.params == (("sample_rate", 1.0),) + assert pushdown.selector is rpm + + def test_points_in_time_selector_is_not_matched(self): + poi = _channel("dtc", series_type=SeriesType.POINTS_IN_TIME) + assert ResamplePushdownRule().match(poi.resample(1e6)) is None + + def test_fractional_rate_needs_double_time_columns(self, rpm): + assert _plan(rpm.resample(0.5), context=_LONG).is_empty + assert not _plan(rpm.resample(0.5), context=_DOUBLE).is_empty + + @pytest.mark.parametrize( + "time_type", [T.FloatType(), T.DecimalType(20, 0), T.TimestampType(), T.StringType()] + ) + def test_unsupported_time_column_types(self, rpm, time_type): + assert _plan(rpm.resample(1e6), context=PushdownContext(time_type, time_type)).is_empty + + +class TestResultTypeCheck: + def test_misdeclared_result_type_is_dropped(self, rpm): + class _IntervalsResample(ResamplePushdownRule): + def result_type(self, pushdown): + return Intervals + + assert _plan(rpm.resample(1e6), rules=[_IntervalsResample()]).is_empty + + def test_failing_type_inference_drops_candidate(self, rpm): + class _Exploding(ResamplePushdownRule): + def match(self, node): + pushdown = super().match(node) + if pushdown is not None: + node.evaluation_type = _raise + return pushdown + + def _raise(): + raise RuntimeError("boom") + + assert _plan(rpm.resample(1e6), rules=[_Exploding()]).is_empty + + +class TestPlanInExecutor: + @pytest.mark.parametrize("dumps", [pickle.dumps, cloudpickle.dumps]) + def test_plan_and_selections_survive_pickling(self, rpm, spd, dumps): + selections = [rpm.resample(1e6).cumtrapz(), spd.resample(1e6), spd.max()] + plan = _plan(*selections) + plan2, selections2 = pickle.loads(dumps((plan, selections))) + + pushed = selections2[0].args[0] + pushdown, variant = plan2.lookup(pushed) + assert variant == 1 + assert pushdown.selector.selector_id == rpm.selector_id + # spd is needed raw by spd.max(), so its resample is evaluated in pandas. + assert plan2.lookup(selections2[1]) is None + assert plan2.lookup(selections2[2]) is None diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py new file mode 100644 index 00000000..0363c8f2 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -0,0 +1,181 @@ +# pylint: disable=missing-function-docstring +"""Parity of the Spark ``resample`` pushdown with ``SampleSeries.resample``.""" + +import numpy as np +import numpy.testing as nptest +import pyspark.sql.types as T +import pytest + +from impulse_query_engine.analyze.metadata.tag_expression import TagSelector +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.query.optimizer import VARIANT_COL, ResamplePushdownRule +from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig +from impulse_query_engine.model.series.sample_series import SampleSeries +from tests.conftest import spark + +_RULE = ResamplePushdownRule() +_CONFIG = SolverConfig() + + +def _pushdown(rate): + selector = TimeSeriesSelector(TagSelector("channel_name") == "x") + return _RULE.match(selector.resample(rate)) + + +def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType()): + """Run the rule on ``{(cid, ch, variant): (tstarts, tends, values)}``. + + Returns ``{(cid, ch, variant): SampleSeries}`` built like the pandas cache does. + """ + schema = T.StructType( + [ + T.StructField("container_id", T.LongType()), + T.StructField("channel_id", T.LongType()), + T.StructField(VARIANT_COL, T.IntegerType()), + T.StructField("tstart", time_type), + T.StructField("tend", time_type), + T.StructField("value", T.DoubleType()), + ] + ) + to_time = int if isinstance(time_type, T.IntegralType) else float + rows = [ + (cid, ch, vid, to_time(ts), to_time(te), None if v is None else float(v)) + for (cid, ch, vid), (tstarts, tends, values) in series_by_key.items() + for ts, te, v in zip(tstarts, tends, values, strict=True) + ] + df = spark.createDataFrame(rows, schema=schema) + result = _RULE.to_spark(df, {vid: _pushdown(rate) for vid, rate in rates.items()}, _CONFIG) + assert result.schema == df.schema + + pdf = result.toPandas() + out = {} + for key, group in pdf.groupby(["container_id", "channel_id", VARIANT_COL]): + group = group.sort_values("tstart", kind="stable") + out[tuple(int(k) for k in key)] = SampleSeries( + group["tstart"], group["tend"], group["value"].astype(np.float64) + ) + return out + + +def _expected(series, rate): + tstarts, tends, values = series + return SampleSeries(tstarts, tends, [np.nan if v is None else v for v in values]).resample( + rate + ) + + +def _assert_same(actual, expected): + nptest.assert_array_equal(actual.tstarts, expected.tstarts) + nptest.assert_array_equal(actual.tends, expected.tends) + nptest.assert_array_equal(actual.values, expected.values) + + +def _where_lt(tstarts, tends, values, threshold): + s = SampleSeries(tstarts, tends, values) + w = s.where(s < threshold) + return list(w.tstarts), list(w.tends), list(w.values) + + +# Inputs of the SampleSeries.resample unit tests plus further edge cases, as +# (tstarts, tends, values, rate). Integral, so exact for LongType columns. +_CASES = { + "two_samples": ([0, 1], [1, 3], [1, 2], 1), + "long_first_sample": ([0, 4, 5], [4, 5, 6], [1, 3, 2], 1), + "uneven": ([0, 4, 6], [4, 6, 7], [1, 3, 2], 1), + "gap": ([0, 1, 6], [1, 3, 7], [1, 2, 4], 1), + "where_gap": (*_where_lt([0, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 10], [5, 5, 5, 4, 5, 1], 5), 1), + "where_gap_rate_2": ( + *_where_lt([0, 1, 2, 3, 4, 5, 10], [1, 2, 3, 4, 5, 10, 11], [5, 5, 5, 4, 5, 1, 2], 5), + 2, + ), + "trailing_point": ([0, 4, 5, 6], [4, 5, 6, 6], [1, 3, 2, 3], 1), + "where_trailing_point": ( + *_where_lt([0, 1, 2, 3, 4, 5, 10], [1, 2, 3, 4, 5, 10, 10], [5, 5, 5, 4, 5, 1, 2], 6), + 1, + ), + "interior_point": ([0, 2, 3], [1, 2, 4], [10, 20, 30], 1), + "interior_and_trailing_point": ([0, 2, 3], [1, 2, 3], [10, 20, 30], 1), + "isolated_trailing_point": ([0, 1, 5], [1, 2, 5], [1, 2, 3], 1), + "single_sample": ([5], [7], [1.0], 1), + "single_point": ([5], [5], [1.0], 1), + "nan_and_null_values": ([0, 1, 2, 3], [1, 2, 3, 4], [1.0, np.nan, None, 4.0], 1), + "grid_off_sample_boundaries": ([0, 3, 7, 12], [3, 7, 12, 20], [1, 2, 3, 4], 4), + "rate_beyond_block": ([0, 1, 2], [1, 2, 3], [1, 2, 3], 10), +} + + +@pytest.mark.parametrize("name", list(_CASES)) +def test_resample_matches_sample_series(spark, name): + tstarts, tends, values, rate = _CASES[name] + actual = _resample_in_spark(spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) + + +def test_all_cases_in_one_pass_and_other_variants_unchanged(spark): + series = {} + expected = {} + for i, (tstarts, tends, values, _rate) in enumerate(_CASES.values()): + # rate 1 -> variant 1, rate 2 -> variant 2; containers keep channels apart. + for vid, r in ((1, 1), (2, 2)): + series[(i, 7, vid)] = (tstarts, tends, values) + expected[(i, 7, vid)] = _expected((tstarts, tends, values), r) + # Raw rows (variant 0) and another rule's variant pass through untouched. + for vid in (0, 3): + series[(i, 7, vid)] = (tstarts, tends, values) + expected[(i, 7, vid)] = SampleSeries( + tstarts, tends, [np.nan if v is None else v for v in values] + ) + actual = _resample_in_spark(spark, series, {1: 1, 2: 2}) + assert actual.keys() == expected.keys() + for key, exp in expected.items(): + _assert_same(actual[key], exp) + + +def test_ns_epoch_longs_above_2_pow_53(spark): + rng = np.random.default_rng(7) + base = 1_700_000_000_123_456_789 + steps = rng.integers(9_000_000, 11_000_000, size=500) # ~10 ms, jittered + times = base + np.concatenate([[0], np.cumsum(steps)]) + tstarts, tends = times[:-1].tolist(), times[1:].tolist() + values = rng.normal(size=500).tolist() + rate = 100_000_000 # 100 ms + actual = _resample_in_spark(spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}) + expected = _expected((tstarts, tends, values), rate) + assert len(expected) < len(tstarts) / 5 + _assert_same(actual[(1, 1, 1)], expected) + + +def test_integer_time_columns(spark): + tstarts, tends, values, rate = _CASES["where_gap_rate_2"] + actual = _resample_in_spark( + spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}, time_type=T.IntegerType() + ) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) + + +def test_seconds_as_double_with_binary_exact_rate(spark): + tstarts = [10.0, 10.25, 11.0, 12.5, 14.0] + tends = [10.25, 11.0, 12.5, 14.0, 14.0] + values = [1.0, 2.0, 3.0, 4.0, 5.0] + actual = _resample_in_spark( + spark, {(1, 1, 1): (tstarts, tends, values)}, {1: 0.75}, time_type=T.DoubleType() + ) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), 0.75)) + + +def test_seconds_as_double_with_fractional_rate(spark): + # 0.1 is not binary-exact: grid points may differ from numpy by one ulp where + # numpy fuses the multiply-add. Sample boundaries stay clear of the grid, so the + # values must match exactly. + rng = np.random.default_rng(11) + times = 1234.5678 + np.concatenate([[0.0], np.cumsum(rng.uniform(0.031, 0.037, 400))]) + tstarts, tends = times[:-1].tolist(), times[1:].tolist() + values = rng.normal(size=400).tolist() + actual = _resample_in_spark( + spark, {(1, 1, 1): (tstarts, tends, values)}, {1: 0.1}, time_type=T.DoubleType() + )[(1, 1, 1)] + expected = _expected((tstarts, tends, values), 0.1) + assert len(actual) == len(expected) + nptest.assert_array_max_ulp(actual.tstarts, expected.tstarts, maxulp=1) + nptest.assert_array_max_ulp(actual.tends, expected.tends, maxulp=1) + nptest.assert_array_equal(actual.values, expected.values) diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py new file mode 100644 index 00000000..2ce8be9a --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py @@ -0,0 +1,308 @@ +# pylint: disable=missing-function-docstring +"""DefaultSolver with ``enable_pushdown``: identical results, fewer rows into the UDF.""" + +from collections import Counter + +import numpy as np +import numpy.testing as nptest +import pandas as pd +import pytest +from pyspark.sql import SparkSession + +from impulse_query_engine.analyze.metadata.tag_expression import TagSelector +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel +from impulse_query_engine.analyze.query.optimizer import ( + RAW_VARIANT, + VARIANT_COL, + PushdownPlan, + QueryOptimizer, + ResamplePushdownRule, +) +from impulse_query_engine.analyze.query.solvers.default_solver import ( + DefaultSolver, + TimeSeriesCache, +) +from impulse_query_engine.analyze.query.solvers.solver_config import ( + ChannelMappingConfig, + SolverConfig, + TableConfig, +) +from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig +from impulse_query_engine.model.series.intervals import Intervals +from impulse_query_engine.model.series.sample_series import SampleSeries +from tests.conftest import spark + +RATE = 1e7 # 10 s; the fixture channels are sampled roughly every 0.9 s + + +def _kvs_config() -> SolverConfig: + return SolverConfig( + project_id="SAMPLE_PROJECT", + container_tags=TableConfig(column_name_mapping={"element_id": "key"}), + container_metrics=TableConfig(column_name_mapping={"project": "project_id"}), + channel_mapping=ChannelMappingConfig(), + ) + + +def _unit_conversion_config() -> SolverConfig: + return SolverConfig( + project_id="SAMPLE_PROJECT", + container_metrics=TableConfig(column_name_mapping={"project": "project_id"}), + channel_mapping=ChannelMappingConfig(filters={"toolbox_id": "container_concept"}), + ) + + +@pytest.fixture +def udf_rows(monkeypatch): + """Rows reaching the grouped-map UDF, per solve: ``{(channel_id, variant): count}``.""" + recorded: list[Counter] = [] + original = DefaultSolver._apply_grouped_map + + def _recording(self, joined_df, *args, **kwargs): + has_variants = VARIANT_COL in joined_df.columns + cols = ["channel_id", VARIANT_COL] if has_variants else ["channel_id"] + counts = Counter() + for row in joined_df.groupBy(*cols).count().collect(): + variant = row[VARIANT_COL] if has_variants else None + counts[(row["channel_id"], variant)] += row["count"] + recorded.append(counts) + return original(self, joined_df, *args, **kwargs) + + monkeypatch.setattr(DefaultSolver, "_apply_grouped_map", _recording) + return recorded + + +def _solve_both(spark, db, make_selections, **solver_kwargs) -> tuple[pd.DataFrame, pd.DataFrame]: + """``toPandas`` results with pushdown off and on, sorted by container.""" + results = [] + for enable_pushdown in (False, True): + solver = DefaultSolver(spark, enable_pushdown=enable_pushdown, **solver_kwargs) + query = db.query + pdf = query.select(*make_selections(query)).toPandas(spark, solver=solver) + results.append(pdf.sort_values("container_id").reset_index(drop=True)) + return results[0], results[1] + + +def _assert_same(actual, expected): + assert list(actual.columns) == list(expected.columns) + assert len(actual) == len(expected) > 0 + for col in expected.columns: + for a, e in zip(actual[col], expected[col], strict=True): + if isinstance(e, SampleSeries): + assert isinstance(a, SampleSeries) + nptest.assert_array_equal(a.tstarts, e.tstarts) + nptest.assert_array_equal(a.tends, e.tends) + nptest.assert_array_equal(a.values, e.values) + else: + nptest.assert_equal(a, e) + + +def _derived_rows(pdf: pd.DataFrame, col: str) -> int: + return sum(len(s) for s in pdf[col]) + + +class TestIdenticalResults: + def test_resample_only(self, spark: SparkSession, key_value_store_db, udf_rows): + off, on = _solve_both( + spark, + key_value_store_db, + lambda q: [q.channel(channel_name="Engine RPM").resample(RATE).alias("rpm")], + config=_kvs_config(), + ) + _assert_same(on, off) + assert all(np.isfinite(s.values).all() and len(s) > 1 for s in on["rpm"]) + + rows_off, rows_on = udf_rows + assert set(rows_on) == {(5, 1)} + assert rows_on[(5, 1)] == _derived_rows(off, "rpm") < rows_off[(5, None)] / 5 + + def test_raw_use_of_the_same_channel_ships_raw_rows_only( + self, spark: SparkSession, key_value_store_db, udf_rows + ): + def selections(q): + rpm = q.channel(channel_name="Engine RPM") + return [rpm.resample(RATE).alias("rpm_10s"), rpm.max().alias("rpm_max")] + + off, on = _solve_both(spark, key_value_store_db, selections, config=_kvs_config()) + _assert_same(on, off) + rows_off, rows_on = udf_rows + # Nothing is pushed down, so the frame keeps today's shape. + assert rows_on == rows_off + + def test_mixed_channels(self, spark: SparkSession, key_value_store_db, udf_rows): + def selections(q): + rpm = q.channel(channel_name="Engine RPM") + spd = q.channel(channel_name="Vehicle Speed Sensor") + return [ + rpm.resample(RATE).alias("rpm_10s"), + spd.max().alias("spd_max"), + spd.resample(RATE).alias("spd_10s"), + ] + + off, on = _solve_both(spark, key_value_store_db, selections, config=_kvs_config()) + _assert_same(on, off) + rows_off, rows_on = udf_rows + assert rows_on == { + (7, RAW_VARIANT): rows_off[(7, None)], + (5, 1): _derived_rows(off, "rpm_10s"), + } + + def test_two_rates_and_nested_ops(self, spark: SparkSession, key_value_store_db, udf_rows): + def selections(q): + rpm = q.channel(channel_name="Engine RPM") + return [ + rpm.resample(RATE).cumtrapz().alias("rpm_integral"), + (rpm.resample(RATE / 2) > 1500).alias("rpm_high"), + rpm.resample(RATE / 2).mean().alias("rpm_mean"), + ] + + off, on = _solve_both(spark, key_value_store_db, selections, config=_kvs_config()) + _assert_same(on, off) + assert any(len(intervals) > 0 for intervals in on["rpm_high"]) + _, rows_on = udf_rows + assert set(rows_on) == {(5, 1), (5, 2)} + + def test_aliased_selector_with_unit_conversion( + self, spark: SparkSession, key_value_store_unit_conversion_db, udf_rows + ): + def selections(q): + speed = q.channel_with_alias(channel_alias="vehicle_speed") + return [speed.resample(RATE).alias("speed_10s")] + + off, on = _solve_both( + spark, key_value_store_unit_conversion_db, selections, config=_unit_conversion_config() + ) + _assert_same(on, off) + _, rows_on = udf_rows + assert {variant for _, variant in rows_on} == {1} + + def test_alongside_points_in_time_selection( + self, spark: SparkSession, basic_narrow_db, udf_rows + ): + def selections(q): + amb = q.channel(channel_name="Ambient Air Temperature") + dtc_count = q.poi_channel(channel_name="DTC_count") + rpm = q.channel(channel_name="Engine RPM") + return [ + amb.where(dtc_count.to_points_in_time()).alias("sampled"), + rpm.resample(RATE).alias("rpm_10s"), + ] + + off, on = _solve_both(spark, basic_narrow_db, selections) + _assert_same(on, off) + _, rows_on = udf_rows + assert rows_on[(91, RAW_VARIANT)] == 3 # POI rows keep the raw variant + assert {variant for (ch, variant) in rows_on if ch == 5} == {1} + + def test_raw_data_with_trailing_point( + self, spark: SparkSession, setup_raw_channels_db, mock_workspace_client + ): + tables = { + name: spark.read.table(f"spark_catalog.silver_raw.{name}") + for name in ("container_metrics", "channel_metrics", "channels") + } + db = MeasurementDB(MeasurementDBConfig.for_debug(tables), ws=mock_workspace_client) + off, on = _solve_both( + spark, + db, + lambda q: [q.channel(channel_name="Engine RPM").resample(2e6).alias("rpm")], + is_raw_data=True, + ) + _assert_same(on, off) + rpm = on["rpm"].iloc[0] + assert rpm.tstarts[-1] == rpm.tends[-1] # the series ends on a point + + +class TestCalculatedChannels: + def test_calculated_channel_with_resample( + self, spark: SparkSession, key_value_store_db, udf_rows + ): + results = [] + for enable_pushdown in (False, True): + q = key_value_store_db.query + rpm = q.channel(channel_name="Engine RPM") + channel = CalculatedChannel(rpm.resample(RATE) * 2, {"channel_name": "rpm_x2"}) + solver = DefaultSolver(spark, config=_kvs_config(), enable_pushdown=enable_pushdown) + df = q.select(channel).solve_calculated_channels(spark, solver=solver) + results.append( + df.drop("identity") + .toPandas() + .sort_values(["container_id", "tstart"]) + .reset_index(drop=True) + ) + off, on = results + assert len(off) > 0 + assert off["value"].gt(0).any() + pd.testing.assert_frame_equal(on, off) + assert set(udf_rows[1]) == {(5, 1)} + + +class TestTimeSeriesCacheVariants: + """Channel 5 carries raw rows (variant 0) and ``resample(10)`` rows (variant 1).""" + + @staticmethod + def _cache(rule, node) -> TimeSeriesCache: + pushdown = rule.match(node) + plan = PushdownPlan.build([rule], [pushdown], [pushdown.selector.selector_id]) + pdf = pd.DataFrame( + { + "container_id": [1] * 5, + "channel_id": [5] * 5, + VARIANT_COL: [1, 1, 0, 0, 0], + "tstart": [0, 10, 0, 5, 10], + "tend": [10, 20, 5, 10, 20], + "value": [1.0, 3.0, 1.0, 2.0, 3.0], + "conversion_factor": [2.0] * 5, + "selector_ids": [[pushdown.selector.selector_id], None, None, None, None], + } + ) + col_map = { + "cid": "container_id", + "ch": "channel_id", + "ts": "tstart", + "te": "tend", + "val": "value", + "conv": "conversion_factor", + } + return TimeSeriesCache(pdf, col_map=col_map, plan=plan) + + def test_raw_and_derived_rows_are_served_separately(self): + node = _rpm().resample(10) + cache = self._cache(ResamplePushdownRule(), node) + assert VARIANT_COL not in cache.mdf.columns + nptest.assert_array_equal(cache.load_blob(1, 5).tstarts, [0, 5, 10]) + derived = node.build(cache) + nptest.assert_array_equal(derived.tstarts, [0, 10]) + nptest.assert_array_equal(derived.values, [1.0, 3.0]) + + def test_conversion_factor_applies_to_aliased_derived_rows(self): + node = _rpm(uses_alias=True).resample(10) + cache = self._cache(ResamplePushdownRule(), node) + nptest.assert_array_equal(node.build(cache).values, [2.0, 6.0]) + + def test_falls_back_without_derived_rows(self): + rpm = _rpm() + cache = self._cache(ResamplePushdownRule(), rpm.resample(10)) + other_rate = rpm.resample(5) # not part of the plan + assert cache.materialized(other_rate) is None + nptest.assert_array_equal(other_rate.build(cache).tstarts, [0, 5, 10, 15]) + + def test_rule_decides_the_result_type(self): + class _IntervalsRule(ResamplePushdownRule): + def decode(self, rows, pushdown, cache): + return Intervals( + rows[cache.tstart_col].to_numpy(), rows[cache.tend_col].to_numpy() + ) + + node = _rpm().resample(10) + result = self._cache(_IntervalsRule(), node).materialized(node) + assert isinstance(result, Intervals) + nptest.assert_array_equal(result.tstarts, [0, 10]) + + def test_optimizer_default_rules(self): + assert [type(r) for r in QueryOptimizer().rules] == [ResamplePushdownRule] + + +def _rpm(uses_alias: bool = False) -> TimeSeriesSelector: + return TimeSeriesSelector(TagSelector("channel_name") == "rpm", uses_alias=uses_alias) diff --git a/tests/impulse_reporting/unit/config/config_parser_test.py b/tests/impulse_reporting/unit/config/config_parser_test.py index 7ed5b12d..2dc626aa 100644 --- a/tests/impulse_reporting/unit/config/config_parser_test.py +++ b/tests/impulse_reporting/unit/config/config_parser_test.py @@ -516,6 +516,18 @@ def test_impulse_config_from_dict_no_query_engine_provided(): assert config.query_engine.solver == Solvers.DEFAULT_SOLVER +@pytest.mark.parametrize( + ("query_engine", "expected"), [({}, False), ({"enable_pushdown": True}, True)] +) +def test_impulse_config_enable_pushdown(query_engine, expected): + """Spark pushdown is opt-in.""" + config_json = { + **impulse_config_JSON, + "query_engine": {"solver": "DefaultSolver", **query_engine}, + } + assert ImpulseConfig.model_validate(config_json).query_engine.enable_pushdown is expected + + def test_impulse_config_from_dict_no_measurement_dim_provided(): """Test ImpulseConfig with no measurement dimension info provided.""" config_json = impulse_config_JSON.copy() diff --git a/tests/impulse_reporting/unit/core/report_test.py b/tests/impulse_reporting/unit/core/report_test.py index a3e9fb5b..7bd5ef12 100644 --- a/tests/impulse_reporting/unit/core/report_test.py +++ b/tests/impulse_reporting/unit/core/report_test.py @@ -292,6 +292,25 @@ def test_create_solver_key_value_store_with_solver_config(): assert solver.config.project_id_col == "project_id" +@pytest.mark.parametrize("enable_pushdown", [False, True]) +def test_create_solver_forwards_enable_pushdown(enable_pushdown): + """``query_engine.enable_pushdown`` decides whether the solver gets an optimizer.""" + config = { + **DUMMY_KEY_VALUE_STORE_CONFIG, + "query_engine": { + **DUMMY_KEY_VALUE_STORE_CONFIG["query_engine"], + "enable_pushdown": enable_pushdown, + }, + } + report = Report( + name="test_report", + spark=None, + workspace_client=create_autospec(WorkspaceClient), + config=config, + ) + assert (report.get_solver().optimizer is not None) is enable_pushdown + + def test_create_solver_delta_default_config(): """DeltaSolver created without solver_config gets a default SolverConfig (no renames).""" from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver From 8201aebc13cf944d9e92a9b661719f746b3b01a9 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Wed, 14 Oct 2026 15:49:46 +0200 Subject: [PATCH 02/16] fix(query-engine): correct resample grid-point assignment to match pandas bucket semantics Update `ResamplePushdownRule` so each sample is assigned the bucket of the first grid point at or after its `tstart`, and only the last sample of a bucket emits the grid points it covers. This fixes cases where grid points land exactly on sample starts and avoids emitting duplicate or over-wide rows. Add unit tests covering grid-on-sample-start scenarios, including nanosecond-epoch timestamps that require the bucket-index correction step. --- .../analyze/query/optimizer/resample_rule.py | 179 +++++++++++------- .../query/optimizer/resample_rule_test.py | 13 ++ 2 files changed, 128 insertions(+), 64 deletions(-) diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py index 7c398b92..93011e83 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -4,7 +4,7 @@ import math import numbers -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any import pyspark.sql.functions as F @@ -37,6 +37,12 @@ class ResamplePushdownRule(PushdownRule): starts at or before it. Meant for downsampling: a rate finer than the native sample spacing yields more rows than the raw channel. + Each sample is assigned the bucket ``c`` of the first grid point at or after its + ``tstart``. The last sample of a bucket holds the value of grid points + ``c, c + 1, …`` up to the next non-empty bucket, so only those samples emit rows + and every row carries just the grid points it covers. Values are looked up + within a block; this matches the pandas evaluation unless intervals overlap. + Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. numpy builds that fuse this into a multiply-add can differ by one ulp for fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral @@ -78,7 +84,7 @@ def to_spark( self, rows: DataFrame, variants: Mapping[int, Pushdown], config: SolverConfig ) -> DataFrame: cid, ch = config.container_id_col, config.channel_id_col - ts, te, val = config.tstart_col, config.tend_col, config.value_col + ts, te = config.tstart_col, config.tend_col key = [cid, ch, VARIANT_COL] rate = F.try_element_at( F.create_map( @@ -91,13 +97,13 @@ def to_spark( F.col(VARIANT_COL), ) - # All per-channel state comes from windows sharing one partitioning and order - # (the pandas cache's tstart order), so the rows are sorted once and the - # incoming container partitioning is kept. + # Every window shares one partitioning and order (the pandas cache's tstart + # order), so the rows are sorted once and the container partitioning is kept. ordered = Window.partitionBy(*key).orderBy(ts, te) running = ordered.rowsBetween(Window.unboundedPreceding, Window.currentRow) + following = ordered.rowsBetween(Window.currentRow, Window.unboundedFollowing) whole = ordered.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) - ts_d, te_d = F.col("_rs_ts"), F.col("_rs_te") + ts_d, te_d, block = F.col("_rs_ts"), F.col("_rs_te"), F.col("_rs_block") prev_te = F.lag(te_d).over(ordered) df = ( rows.withColumn("_rs_ts", F.col(ts).cast("double")) @@ -107,12 +113,11 @@ def to_spark( .withColumn("_rs_last_ts", F.last(ts_d).over(whole)) .withColumn("_rs_last_te", F.last(te_d).over(whole)) .withColumn("_rs_block", F.sum("_rs_break").over(running)) - .withColumn( - "_rs_block_end", F.coalesce(F.lead("_rs_break").over(ordered), F.lit(1)) == 1 - ) ) - # Block bounds as seen on the block's last row: its first tstart, and the - # largest tend so far within the block (blocks are numbered in order). + # Block bounds on every row: the block's first tstart and its largest tend. + # Blocks are numbered in order, so max(struct(±block, tend)) over the rows + # before / after the current one stays within the current block. + te_no_nan = _ignore_nan(te_d).alias("te") df = ( df.withColumn( "_rs_t_min", @@ -120,91 +125,137 @@ def to_spark( ) .withColumn( "_rs_t_max", - F.max(F.struct(F.col("_rs_block"), _ignore_nan(te_d).alias("te"))) - .over(running) - .getField("te"), + F.greatest( + F.max(F.struct(block, te_no_nan)).over(running).getField("te"), + F.max(F.struct(-block, te_no_nan)).over(following).getField("te"), + ), ) - .withColumn("_rs_nblk", F.max("_rs_block").over(whole)) + .withColumn("_rs_nblk", F.max(block).over(whole)) ) # numpy.arange(t_min, t_max + rate, rate): length ceil(((t_max + rate) - t_min) / rate), # filled as t_min, t_min + rate, then t_min + k * ((t_min + rate) - t_min). t_min, t_max = F.col("_rs_t_min"), F.col("_rs_t_max") - length = F.ceil(((t_max + rate) - t_min) / rate) delta = (t_min + rate) - t_min + + def grid_point(k: Column) -> Column: + return ( + F.when(k == 0, t_min) + .when(k == 1, t_min + rate) + .otherwise(t_min + k.cast("double") * delta) + ) + + resampled = F.col(VARIANT_COL).isin(sorted(variants)) & (F.col("_rs_n") >= 2) # The last block keeps a grid point at t_max when the series ends with a # zero-duration sample. - inclusive = (F.col("_rs_block") == F.col("_rs_nblk")) & ( - F.col("_rs_last_ts") == F.col("_rs_last_te") + inclusive = (block == F.col("_rs_nblk")) & (F.col("_rs_last_ts") == F.col("_rs_last_te")) + length = F.greatest(F.ceil(((t_max + rate) - t_min) / rate), F.lit(0)) + + # Bucket of each sample, and the number of grid points of its block. + df = _with_first_grid_index( + df, "_rs_c", ts_d, F.lit(False), resampled, grid_point, t_min, delta ) - resampled = F.col(VARIANT_COL).isin(sorted(variants)) & (F.col("_rs_n") >= 2) - grid_points = F.filter( - F.transform( - F.sequence(F.lit(0).cast("long"), length - 1), - lambda k: F.when(k == 0, t_min) - .when(k == 1, t_min + rate) - .otherwise(t_min + k.cast("double") * delta), - ), - lambda g: (g < t_max) | (inclusive & (g == t_max)), + df = _with_first_grid_index( + df, "_rs_grid_end", t_max, inclusive, resampled, grid_point, t_min, delta ) + c, n_grid = F.col("_rs_c"), F.least(F.col("_rs_grid_end"), length) + + # The last sample of a bucket emits the grid points up to the next non-empty + # bucket of its block (or up to the block's end). + next_c, next_block = F.lead(c).over(ordered), F.lead(block).over(ordered) df = df.withColumn( - "_rs_grid", - F.when(resampled & F.col("_rs_block_end") & (length > 0), grid_points), + "_rs_c_stop", + F.when(next_block.isNull() | (next_block != block), n_grid).when( + next_c != c, F.least(next_c, n_grid) + ), ) - # One output element per row: the row itself for series this rule leaves - # alone (kind 2), otherwise the sample (kind 0) plus, on a block's last row, - # the block's grid points (kind 1). Fewer than two samples: returned as-is. def element(kind: int, t: Column, tend: Column) -> Column: return F.struct( - F.lit(kind).alias("kind"), t.alias("t"), tend.cast("double").alias("tend") + F.lit(kind).alias("kind"), + t.cast("double").alias("t"), + tend.cast("double").alias("tend"), ) + c_stop = F.col("_rs_c_stop") grid = F.transform( - F.col("_rs_grid"), - lambda g, i: element( - 1, g, F.coalesce(F.try_element_at(F.col("_rs_grid"), i + 2), t_max) + F.sequence(c, c_stop - 1), + lambda k: element( + 1, + grid_point(k), + F.when(k == n_grid - 1, t_max).otherwise(grid_point(k + 1)), ), ) - no_tend = F.lit(None).cast("double") - elements = ( - F.when(~resampled, F.array(element(2, ts_d, no_tend))) - .when( - F.col("_rs_grid").isNotNull(), F.concat(F.array(element(0, ts_d, no_tend)), grid) - ) - .otherwise(F.array(element(0, ts_d, no_tend))) - ) + no_time = F.lit(None) + # Series this rule leaves alone keep their row (kind 2); of the resampled + # series only the grid points survive (kind 1). explode drops the rest. + elements = F.when(~resampled, F.array(element(2, no_time, no_time))).when(c_stop > c, grid) df = df.withColumn("_rs_e", F.explode(elements)) - kind, t = F.col("_rs_e.kind"), F.col("_rs_e.t") - - # Previous-value lookup: each grid point takes the value of the last sample (in - # channel order) starting at or before it. Samples sort before grid points at - # equal times; the struct keeps null values distinguishable from "no sample". - as_of = ( - Window.partitionBy(*key) - .orderBy(t, kind, ts, te) - .rowsBetween(Window.unboundedPreceding, Window.currentRow) - ) - previous = F.last(F.when(kind == 0, F.struct(F.col(val).alias("value"))), ignorenulls=True) - df = df.withColumn("_rs_previous", previous.over(as_of)).where(kind != 0) - is_grid = kind == 1 + is_grid = F.col("_rs_e.kind") == 1 derived = { - ts: t.cast(rows.schema[ts].dataType), + ts: F.col("_rs_e.t").cast(rows.schema[ts].dataType), te: F.col("_rs_e.tend").cast(rows.schema[te].dataType), - # Beyond the last sample's end the pandas interpolation yields NaN. - val: F.when(t > F.col("_rs_last_te"), F.lit(None)) - .otherwise(F.col("_rs_previous.value")) - .cast(rows.schema[val].dataType), } return df.select( *[ - F.when(is_grid, derived[c]).otherwise(F.col(c)).alias(c) if c in derived else c - for c in rows.columns + ( + F.when(is_grid, derived[col]).otherwise(F.col(col)).alias(col) + if col in derived + else col + ) + for col in rows.columns ] ) +def _with_first_grid_index( + df: DataFrame, + name: str, + x: Column, + strict: Column, + active: Column, + grid_point: Callable[[Column], Column], + t_min: Column, + delta: Column, +) -> DataFrame: + """Add *name*: the smallest ``k >= 0`` with ``grid_point(k) >= x`` (``> x`` if *strict*). + + Estimated from ``(x - t_min) / delta`` and corrected by one step against the + exact grid points; a result that still is not exact raises instead of + misplacing a sample. Only evaluated where *active*. + """ + + def reached(k: Column) -> Column: + g = grid_point(k) + return F.when(strict, g > x).otherwise(g >= x) + + estimate, k = F.col(f"{name}_estimate"), F.col(name) + return ( + df.withColumn( + f"{name}_estimate", + F.when( + active, + F.when(delta > 0, F.greatest(F.ceil((x - t_min) / delta), F.lit(0))).otherwise( + F.raise_error( + F.lit("resample: sample_rate is below the resolution of the timestamps") + ) + ), + ), + ) + .withColumn( + name, + estimate, + ) + .withColumn( + name, + F.when(~active, F.lit(None)) + .when(reached(k) & ((k == 0) | ~reached(k - 1)), k) + .otherwise(F.raise_error(F.lit("resample: grid index could not be determined"))), + ) + ) + + def _sample_rate(args: tuple, kwargs: dict[str, Any]) -> float | None: """The resample rate of a ``resample(...)`` call, or ``None`` if unsupported.""" if len(args) > 1 or set(kwargs) - {_RATE} or (args and kwargs): diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py index 0363c8f2..03b58aa7 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -101,6 +101,7 @@ def _where_lt(tstarts, tends, values, threshold): "nan_and_null_values": ([0, 1, 2, 3], [1, 2, 3, 4], [1.0, np.nan, None, 4.0], 1), "grid_off_sample_boundaries": ([0, 3, 7, 12], [3, 7, 12, 20], [1, 2, 3, 4], 4), "rate_beyond_block": ([0, 1, 2], [1, 2, 3], [1, 2, 3], 10), + "grid_on_sample_starts": (list(range(10)), list(range(1, 11)), list(range(10)), 3), } @@ -145,6 +146,18 @@ def test_ns_epoch_longs_above_2_pow_53(spark): _assert_same(actual[(1, 1, 1)], expected) +def test_ns_epoch_grid_on_sample_starts(spark): + # Every fifth sample starts exactly on a grid point, the case where the bucket + # index estimate needs its correction step. + base = 1_700_000_000_123_456_789 + times = [base + i * 10_000_000 for i in range(301)] + tstarts, tends = times[:-1], times[1:] + values = [float(i) for i in range(300)] + rate = 50_000_000 + actual = _resample_in_spark(spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) + + def test_integer_time_columns(spark): tstarts, tends, values, rate = _CASES["where_gap_rate_2"] actual = _resample_in_spark( From e2014041a484c55c5c7cd895ba2781485f235fdd Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Wed, 14 Oct 2026 15:50:34 +0200 Subject: [PATCH 03/16] fix(query-engine): refine resample first-grid-index estimate correction Adjust `_with_first_grid_index` to nudge the estimated bucket index up or down only when the reachability check at the adjacent index differs, falling back to the raw estimate otherwise. This tightens the correction step introduced for nanosecond-epoch timestamps without over-adjusting cases where the initial estimate already lands on the correct grid point. --- .../analyze/query/optimizer/resample_rule.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py index 93011e83..21445eaa 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -245,7 +245,9 @@ def reached(k: Column) -> Column: ) .withColumn( name, - estimate, + F.when((estimate > 0) & reached(estimate - 1), estimate - 1) + .when(~reached(estimate), estimate + 1) + .otherwise(estimate), ) .withColumn( name, From f45101c969728c54789571ec0891081fd437f9f8 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Wed, 14 Oct 2026 15:56:14 +0200 Subject: [PATCH 04/16] docs(query-engine): clarify resample pushdown semantics and Spark implementation details Update API reference docs for `PushdownRule.to_spark` to describe the single-pass, lineage-preserving contract for replacing variant rows with derived rows. Expand `ResamplePushdownRule` and query solver docs to explain grid bucket assignment, last-sample row emission, and the bounded memory per row. --- .../analyze/query/optimizer/pushdown.md | 18 ++++++++++++------ .../analyze/query/optimizer/resample_rule.md | 6 ++++++ .../references/query_engine/query_solvers.md | 9 ++++++++- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md index 04fa9774..057ad252 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -130,19 +130,25 @@ def to_spark(rows: DataFrame, variants: Mapping[int, Pushdown], config: SolverConfig) -> DataFrame ``` -Compute all *variants* of this rule in one pass. +Replace the rows of this rule's *variants* with their derived rows. + +*rows* is read once and shared by all rules, so implementations must +stay a single lineage (no self-joins or unions of *rows*) and keep its +container partitioning (windows partitioned by container id, channel id +and :data:`VARIANT_COL` reuse it). **Arguments**: -- `rows` (`pyspark.sql.DataFrame`): Channel-data rows tagged with :data:`VARIANT_COL`, one copy per -variant a channel needs; only ids in *variants* occur. -- `variants` (`Mapping[int, Pushdown]`): Variant id to a representative pushdown carrying its parameters. +- `rows` (`pyspark.sql.DataFrame`): All channel-data rows of the solve, tagged with :data:`VARIANT_COL` +(one copy per variant a channel ships) and partitioned by container. +- `variants` (`Mapping[int, Pushdown]`): This rule's variant ids, each with a representative pushdown carrying +its parameters. - `config` (`SolverConfig`): Provides the internal column names. **Returns**: -`pyspark.sql.DataFrame`: Derived rows with the container id, channel id, :data:`VARIANT_COL`, -``tstart``, ``tend`` and ``value`` columns of *rows*, same types. +`pyspark.sql.DataFrame`: *rows* with the same columns and types, the rows of *variants* +replaced by derived rows and all other rows unchanged. #### decode diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md index 60f36085..18a9fd04 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md @@ -20,6 +20,12 @@ regular grid starting at the block's first ``tstart`` (built like starts at or before it. Meant for downsampling: a rate finer than the native sample spacing yields more rows than the raw channel. +Each sample is assigned the bucket ``c`` of the first grid point at or after its +``tstart``. The last sample of a bucket holds the value of grid points +``c, c + 1, …`` up to the next non-empty bucket, so only those samples emit rows +and every row carries just the grid points it covers. Values are looked up +within a block; this matches the pandas evaluation unless intervals overlap. + Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. numpy builds that fuse this into a multiply-add can differ by one ulp for fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index 30ecb684..378702c3 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -200,7 +200,14 @@ Results match the pandas evaluation exactly for integer timestamp columns with a `sample_rate` (e.g. µs or ns epochs), and for double timestamps whose grid is exact in binary floating point. For double timestamps with a fractional rate such as `0.1`, the grid points can differ in the last bit from numpy builds that fuse multiply-add operations. The order of samples -that share the same `tstart` is unspecified, as in the pandas evaluation. +that share the same `tstart` is unspecified, as in the pandas evaluation, and channels whose +`[tstart, tend)` intervals overlap are not supported (the RLE and interval encoders never +produce overlaps). + +Spark processes each channel in one pass over its samples, sorted by `tstart`: every sample is +assigned the grid bucket it falls into, and only the last sample of each bucket emits rows (the +grid points up to the next non-empty bucket). The memory per row is bounded by the grid points a +single sample covers, not by the length of the channel. ## API reference From 9221f4ff798aa9865816f27d913ece9ae745b3c5 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Wed, 14 Oct 2026 16:15:47 +0200 Subject: [PATCH 05/16] fix(query-engine): avoid quadratic window frames in resample pushdown Replace the `ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING` frame in `ResamplePushdownRule` with a two-pass approach using only linear frames. First pass computes per-block state with running/whole-partition windows; second pass derives block `t_max` and grid counts over the filtered last-sample rows. This prevents Spark from recomputing an unbounded-following frame per row. Add a unit test asserting the optimized plan contains no unbounded-following row frame. --- .../analyze/query/optimizer/resample_rule.py | 59 ++++++++++--------- .../query/optimizer/resample_rule_test.py | 11 ++++ 2 files changed, 41 insertions(+), 29 deletions(-) diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py index 21445eaa..529a2b06 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -97,11 +97,11 @@ def to_spark( F.col(VARIANT_COL), ) - # Every window shares one partitioning and order (the pandas cache's tstart - # order), so the rows are sorted once and the container partitioning is kept. + # Pass 1, over each series in the pandas cache's tstart order. Only running, + # whole-partition and neighbour windows (all linear in the series length); + # they share one sort and keep the container partitioning. ordered = Window.partitionBy(*key).orderBy(ts, te) running = ordered.rowsBetween(Window.unboundedPreceding, Window.currentRow) - following = ordered.rowsBetween(Window.currentRow, Window.unboundedFollowing) whole = ordered.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) ts_d, te_d, block = F.col("_rs_ts"), F.col("_rs_te"), F.col("_rs_block") prev_te = F.lag(te_d).over(ordered) @@ -113,29 +113,22 @@ def to_spark( .withColumn("_rs_last_ts", F.last(ts_d).over(whole)) .withColumn("_rs_last_te", F.last(te_d).over(whole)) .withColumn("_rs_block", F.sum("_rs_break").over(running)) - ) - # Block bounds on every row: the block's first tstart and its largest tend. - # Blocks are numbered in order, so max(struct(±block, tend)) over the rows - # before / after the current one stays within the current block. - te_no_nan = _ignore_nan(te_d).alias("te") - df = ( - df.withColumn( + .withColumn( "_rs_t_min", F.last(F.when(F.col("_rs_break") == 1, ts_d), ignorenulls=True).over(running), ) + # Largest tend so far within the block (blocks are numbered in order); on + # the block's last row this is the block's t_max. .withColumn( - "_rs_t_max", - F.greatest( - F.max(F.struct(block, te_no_nan)).over(running).getField("te"), - F.max(F.struct(-block, te_no_nan)).over(following).getField("te"), - ), + "_rs_t_max_so_far", + F.max(F.struct(block, _ignore_nan(te_d).alias("te"))).over(running).getField("te"), ) .withColumn("_rs_nblk", F.max(block).over(whole)) ) # numpy.arange(t_min, t_max + rate, rate): length ceil(((t_max + rate) - t_min) / rate), # filled as t_min, t_min + rate, then t_min + k * ((t_min + rate) - t_min). - t_min, t_max = F.col("_rs_t_min"), F.col("_rs_t_max") + t_min = F.col("_rs_t_min") delta = (t_min + rate) - t_min def grid_point(k: Column) -> Column: @@ -146,28 +139,36 @@ def grid_point(k: Column) -> Column: ) resampled = F.col(VARIANT_COL).isin(sorted(variants)) & (F.col("_rs_n") >= 2) + df = _with_first_grid_index( + df, "_rs_c", ts_d, F.lit(False), resampled, grid_point, t_min, delta + ) + c, next_c, next_block = F.col("_rs_c"), F.col("_rs_next_c"), F.lead(block).over(ordered) + df = ( + df.withColumn("_rs_next_c", F.lead(c).over(ordered)).withColumn( + "_rs_block_end", next_block.isNull() | (next_block != block) + ) + # Only the last sample of each bucket emits grid points. + .where(~resampled | F.col("_rs_block_end") | (next_c != c)) + ) + + # Pass 2, over the remaining rows (about one per grid point): the block's + # t_max, and from it the number of grid points of the block. + t_max = F.max("_rs_t_max_so_far").over(Window.partitionBy(*key, "_rs_block")) + df = df.withColumn("_rs_t_max", t_max) + t_max = F.col("_rs_t_max") # The last block keeps a grid point at t_max when the series ends with a # zero-duration sample. inclusive = (block == F.col("_rs_nblk")) & (F.col("_rs_last_ts") == F.col("_rs_last_te")) length = F.greatest(F.ceil(((t_max + rate) - t_min) / rate), F.lit(0)) - - # Bucket of each sample, and the number of grid points of its block. - df = _with_first_grid_index( - df, "_rs_c", ts_d, F.lit(False), resampled, grid_point, t_min, delta - ) df = _with_first_grid_index( df, "_rs_grid_end", t_max, inclusive, resampled, grid_point, t_min, delta ) - c, n_grid = F.col("_rs_c"), F.least(F.col("_rs_grid_end"), length) - - # The last sample of a bucket emits the grid points up to the next non-empty - # bucket of its block (or up to the block's end). - next_c, next_block = F.lead(c).over(ordered), F.lead(block).over(ordered) + n_grid = F.least(F.col("_rs_grid_end"), length) + # A bucket's last sample emits the grid points up to the next non-empty bucket + # of its block, or up to the block's end. df = df.withColumn( "_rs_c_stop", - F.when(next_block.isNull() | (next_block != block), n_grid).when( - next_c != c, F.least(next_c, n_grid) - ), + F.when(F.col("_rs_block_end"), n_grid).otherwise(F.least(next_c, n_grid)), ) def element(kind: int, t: Column, tend: Column) -> Column: diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py index 03b58aa7..85875ed0 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -146,6 +146,17 @@ def test_ns_epoch_longs_above_2_pow_53(spark): _assert_same(actual[(1, 1, 1)], expected) +def test_only_linear_window_frames(spark): + # Spark recomputes a ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING frame for every + # row (quadratic in the series length), so the rule must not use one. + schema = "container_id long, channel_id long, {} int, tstart long, tend long, value double" + df = spark.createDataFrame([(1, 1, 1, 0, 1, 1.0)], schema=schema.format(VARIANT_COL)) + result = _RULE.to_spark(df, {1: _pushdown(1)}, _CONFIG) + plan = result._jdf.queryExecution().optimizedPlan().toString() + assert "specifiedwindowframe(RowFrame, currentrow$(), unboundedfollowing$())" not in plan + assert "specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())" in plan + + def test_ns_epoch_grid_on_sample_starts(spark): # Every fifth sample starts exactly on a grid point, the case where the bucket # index estimate needs its correction step. From 989eeb90f9605cfc0c17e6554485c844ce457e10 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Wed, 14 Oct 2026 19:54:36 +0200 Subject: [PATCH 06/16] feat(query-engine): solve whole-pushdown calculated channels natively in Spark Add `PushdownPlan.native_targets` and `PushdownRule.emits_result_rows` so `DefaultSolver.solve_calculated_channels` can bypass the pandas UDF when every selection is a `CalculatedChannel` whose entire expression is a pushdown that emits its result rows directly. Native rows are produced from the derived variant rows and attached to the same identity map as the UDF path. Also make `DefaultSolver` resolve selectors to at most one channel per container, keeping the lowest `channel_id` (after alias priority), and give aliased and direct selections on the same tags distinct `selector_id`s. Update docs and add unit/integration tests covering native calculated channels, channel resolution deduplication, and pushdown parity. --- .../analyze/query/optimizer/plan.md | 22 ++ .../analyze/query/optimizer/pushdown.md | 11 + .../analyze/query/solvers/default_solver.md | 4 + .../references/query_engine/query_solvers.md | 11 + .../query_engine/tsal/defining_expressions.md | 11 + .../metadata/time_series_expression.py | 9 +- .../analyze/query/optimizer/plan.py | 52 +++- .../analyze/query/optimizer/pushdown.py | 12 + .../analyze/query/solvers/default_solver.py | 198 +++++++++++++- .../unit/analyze/query/optimizer/plan_test.py | 108 ++++++++ .../default_solver_channel_resolution_test.py | 183 +++++++++++++ ..._solver_native_calculated_channels_test.py | 247 ++++++++++++++++++ .../integration/calculated_channel_test.py | 49 ++++ 13 files changed, 900 insertions(+), 17 deletions(-) create mode 100644 tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md index 9bf7bb41..c549f9f6 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md @@ -76,6 +76,28 @@ def lookup(node: TimeSeriesExpression) -> tuple[Pushdown, int] | None Return ``(pushdown, variant)`` when *node* is a pushed-down subtree. +#### native\_targets + +```python +def native_targets( + selections: Iterable[Any] +) -> dict[int, dict[int, list[tuple[int, bool]]]] | None +``` + +Sources of calculated channels that are pushed down as a whole. + +**Arguments**: + +- `selections` (`Iterable`): The selections of a ``solve_calculated_channels`` call. + +**Returns**: + +`dict or None`: ``{variant: {selector_id: [(channel_id, uses_alias), ...]}}``: the +calculated channels (by output ``channel_id``) whose rows are the +given variant's rows of the channel resolved for ``selector_id``. +``None`` unless every selection is a ``CalculatedChannel`` whose +whole expression is a pushdown emitting its result rows. + #### channel\_variants ```python diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md index 057ad252..74dbe7c3 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -160,3 +160,14 @@ def decode(rows: pd.DataFrame, pushdown: Pushdown, Build the subtree's result from one container/channel/variant slice. +#### emits\_result\_rows + +```python +def emits_result_rows(pushdown: Pushdown) -> bool +``` + +Whether *pushdown*'s derived rows are its result's samples as they are. + +Holds when the result is a :class:`SampleSeries` built by the default + + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md index 463e0f3c..3eba53ee 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md @@ -172,6 +172,10 @@ frame; the selectors are applied directly to ``channel_metrics`` columns, the result is restricted to the candidate containers, and ``selector_ids`` is computed from the matching selectors. +In both modes a selector resolves to at most one channel per container: +when several channels of a container match, the lowest ``channel_id`` +is kept. + **Arguments**: - `spark` (`SparkSession`): Spark session used for query execution. diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index 378702c3..16614469 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -145,6 +145,8 @@ Every selection must be a [`CalculatedChannel`](../report/channel.md) (identity need not match across selections), and each wrapped expression must evaluate to a `SampleSeries`. The stage reuses the same metadata filter pipeline as `solve()` to resolve the input channels, then runs a grouped-map UDF per container. Only `DefaultSolver` implements it — the base `QuerySolver` raises `NotImplementedError`. +The `channel_id` column of `channels` must be an integer or long column to hold the calculated +channels' ids. ```python from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel @@ -209,6 +211,15 @@ assigned the grid bucket it falls into, and only the last sample of each bucket grid points up to the next non-empty bucket). The memory per row is bounded by the grid points a single sample covers, not by the length of the channel. +### Calculated channels without the pandas UDF + +When every calculated channel of a `solve_calculated_channels` call (in a report: of a solve +batch) is a pushed-down expression as a whole, e.g. +`CalculatedChannel(veh_spd.resample(1e6), {...})`, its rows are produced entirely in Spark and no +pandas UDF runs. If any calculated channel of the call needs the pandas evaluation, e.g. +`veh_spd.resample(1e6) * 3.6`, all of them are evaluated in the UDF as usual; their pushed-down +parts still run in Spark. Either way the rows are identical. + ## API reference Auto-generated symbol-level docs: diff --git a/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md b/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md index a8bd2647..aa4650c1 100644 --- a/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md +++ b/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md @@ -31,6 +31,14 @@ Each tag passed to `channel(...)` must be resolvable by the solver — either as table is configured — see [Query Solvers](../query_solvers.md#how-defaultsolver-adapts). ::: +:::note When a selection matches several channels +A selection resolves to at most one channel per container. If several channels of a container +match its tags (e.g. the same `channel_name` under two `data_key`s), the one with the lowest +`channel_id` is used; add tags such as `data_key` to select a specific channel. Earlier versions +picked an arbitrary match, so recompute incremental reports in full if their selections were +ambiguous. +::: + ### Points-in-Time (POI) channels Not every channel can be interpreted as a continuous signal within its `[tstart, tend)` intervals. @@ -95,6 +103,9 @@ channel-identifying columns (by default `channel_name` and `data_key`). See [Query Solvers](../query_solvers.md) for details. ::: +If several mapping entries of an alias resolve in a container, the one with the lowest `priority` +wins; without priorities, or on ties, the lowest `channel_id` is used. + When the `channel_mapping` table carries `source_unit` and `target_unit` columns and the report config sets `source.unit_conversion_table`, values returned from `channel_with_alias()` are automatically converted from source to target unit before any expression is evaluated. Constants and parameters in expressions over an aliased selector diff --git a/src/impulse_query_engine/analyze/metadata/time_series_expression.py b/src/impulse_query_engine/analyze/metadata/time_series_expression.py index 8b335011..f6c87c78 100644 --- a/src/impulse_query_engine/analyze/metadata/time_series_expression.py +++ b/src/impulse_query_engine/analyze/metadata/time_series_expression.py @@ -796,8 +796,12 @@ def value_type(self) -> SeriesValueType: @property def selector_id(self) -> int: # series_type / value_type keep distinct-typed selections of the same tag - # apart; CONTINUOUS uses the bare ``str(expr)`` hash for back-compat. + # apart, and an aliased selection resolves differently from a direct one on + # the same tags; a direct CONTINUOUS selection uses the bare ``str(expr)`` + # hash for back-compat. if self._series_type is SeriesType.CONTINUOUS: + if self._uses_alias: + return zlib.crc32(f"alias|{self._expr}".encode()) return zlib.crc32(str(self._expr).encode()) return zlib.crc32(f"{self._series_type}|{self._expr}|{self._value_type}".encode()) @@ -868,7 +872,8 @@ def build(self, cache: SeriesCache) -> SampleSeries | PointsInTimeSeries: if self._series_type is SeriesType.POINTS_IN_TIME: return PointsInTimeSeries.empty(value_type=self._value_type) return SampleSeries.empty() - # TODO: select candidate + # DefaultSolver resolves a selector to at most one channel per container; + # other solvers may pass several candidates, of which the first is used. mid = candidates.container_id.iloc[0] cid = candidates.channel_id.iloc[0] # The selector is the source of truth for the series type: pass it to the diff --git a/src/impulse_query_engine/analyze/query/optimizer/plan.py b/src/impulse_query_engine/analyze/query/optimizer/plan.py index 5e30299e..90df8048 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/plan.py +++ b/src/impulse_query_engine/analyze/query/optimizer/plan.py @@ -4,6 +4,7 @@ from collections.abc import Iterable, Mapping from dataclasses import dataclass, field +from typing import Any import pyspark.sql.functions as F from pyspark.sql import Column @@ -106,13 +107,55 @@ def lookup(self, node: TimeSeriesExpression) -> tuple[Pushdown, int] | None: return None return None + def native_targets( + self, selections: Iterable[Any] + ) -> dict[int, dict[int, list[tuple[int, bool]]]] | None: + """Sources of calculated channels that are pushed down as a whole. + + Parameters + ---------- + selections : Iterable + The selections of a ``solve_calculated_channels`` call. + + Returns + ------- + dict or None + ``{variant: {selector_id: [(channel_id, uses_alias), ...]}}``: the + calculated channels (by output ``channel_id``) whose rows are the + given variant's rows of the channel resolved for ``selector_id``. + ``None`` unless every selection is a ``CalculatedChannel`` whose + whole expression is a pushdown emitting its result rows. + """ + from impulse_query_engine.analyze.query.channels.calculated_channel import ( + CalculatedChannel, + ) + + selections = list(selections) + if not selections: + return None + targets: dict[int, dict[int, list[tuple[int, bool]]]] = {} + for selection in selections: + if not isinstance(selection, CalculatedChannel): + return None + hit = self.lookup(selection.expr) + if hit is None or not hit[0].rule.emits_result_rows(hit[0]): + return None + pushdown, variant = hit + by_selector = targets.setdefault(variant, {}) + by_selector.setdefault(pushdown.selector.selector_id, []).append( + (selection.channel_id, pushdown.selector.uses_alias) + ) + return targets + def channel_variants(self, selector_ids: Column) -> Column: """Variant ids a matched channel ships, given its ``selector_ids`` array. A channel ships only its raw data (``[0]``) unless every one of its selectors is pushed down; then it ships the derived variants instead. """ - ids = F.transform(selector_ids, lambda sid: sid.cast("long")) + ids = F.filter( + F.transform(selector_ids, lambda sid: sid.cast("long")), lambda sid: sid.isNotNull() + ) pushed = sorted(self.variants_by_selector) needs_raw = F.exists(ids, lambda sid: ~sid.isin(pushed)) | (F.size(ids) == 0) mapping = F.create_map( @@ -125,8 +168,13 @@ def channel_variants(self, selector_ids: Column) -> Column: ) ] ) + no_variants = F.array().cast("array") derived = F.array_distinct( - F.flatten(F.transform(ids, lambda sid: F.try_element_at(mapping, sid))) + F.flatten( + F.transform( + ids, lambda sid: F.coalesce(F.try_element_at(mapping, sid), no_variants) + ) + ) ) return F.when(F.coalesce(needs_raw, F.lit(True)), F.array(F.lit(RAW_VARIANT))).otherwise( derived diff --git a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py index 7ac8bb16..5443d6c5 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -146,3 +146,15 @@ def to_spark( def decode(self, rows: pd.DataFrame, pushdown: Pushdown, cache: TimeSeriesCache) -> Any: """Build the subtree's result from one container/channel/variant slice.""" return cache.to_sample_series(rows, uses_alias=pushdown.selector.uses_alias) + + def emits_result_rows(self, pushdown: Pushdown) -> bool: + """Whether *pushdown*'s derived rows are its result's samples as they are. + + Holds when the result is a :class:`SampleSeries` built by the default + :meth:`decode`: the variant's rows, with values multiplied by the + conversion factor for aliased selectors, are then exactly the result, so + it can be emitted without the pandas evaluation. + """ + return ( + type(self).decode is PushdownRule.decode and self.result_type(pushdown) is SampleSeries + ) diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index bc5aa214..f66443c7 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -8,7 +8,7 @@ import pandas as pd import pyspark.sql.functions as F import pyspark.sql.types as T -from pyspark.sql import DataFrame, Window +from pyspark.sql import Column, DataFrame, Window from impulse_query_engine.analyze.metadata.metric_expression import MetricExpression from impulse_query_engine.analyze.metadata.tag_expression import TagExpression @@ -618,6 +618,10 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: columns, the result is restricted to the candidate containers, and ``selector_ids`` is computed from the matching selectors. + In both modes a selector resolves to at most one channel per container: + when several channels of a container match, the lowest ``channel_id`` + is kept. + Parameters ---------- spark : SparkSession @@ -654,10 +658,7 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: on=[container_id_col, channel_id_col], how="inner", ) - metrics = metrics.withColumn("selector_ids", F.array(F.col("selector_id"))).drop( - "selector_id" - ) - return metrics.select(container_id_col, channel_id_col, "selector_ids") + return self._one_channel_per_selector(metrics) # Wide mode: selectors apply directly to channel_metrics columns. channel_metrics_df = channel_metrics_df.where(self._build_expr(selectors)) @@ -666,10 +667,28 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: on=[container_id_col], how="inner", ) - result = result.withColumn( - "selector_ids", F.array(self._build_selector_id_expr(selectors)) + result = result.withColumn("selector_id", self._build_selector_id_expr(selectors)) + return self._one_channel_per_selector(result) + + def _one_channel_per_selector(self, matches: DataFrame) -> DataFrame: + """Keep the lowest matching ``channel_id`` per container and direct selector. + + *matches* carries one ``selector_id`` per ``(container_id, channel_id)`` + row (the first selector matching the channel). Returns + ``(container_id, channel_id, selector_ids)`` with one row per container + and selector. + """ + container_id_col = self.config.container_id_col + channel_id_col = self.config.channel_id_col + return ( + matches.groupBy(container_id_col, "selector_id") + .agg(F.min(channel_id_col).alias(channel_id_col)) + .select( + container_id_col, + channel_id_col, + F.array(F.col("selector_id")).alias("selector_ids"), + ) ) - return result.select(container_id_col, channel_id_col, "selector_ids") def filter_aliased_channel_metrics( self, spark, db: MeasurementDB, container_df, selectors @@ -801,8 +820,9 @@ def filter_aliased_channel_metrics( resolved = resolved.withColumn(source_unit_col, F.col("_map_source_unit")) resolved = resolved.withColumn(target_unit_col, F.col("_map_target_unit")) + # Ties in priority (e.g. no priorities configured) go to the lowest channel_id. dedup_window = Window.partitionBy(container_id_col, channel_alias_col).orderBy( - F.col(alias_priority_col).asc_nulls_last() + F.col(alias_priority_col).asc_nulls_last(), F.col(channel_id_col).asc() ) resolved = resolved.withColumn("_rank", F.row_number().over(dedup_window)) resolved = resolved.where(F.col("_rank") == 1).drop("_rank") @@ -885,7 +905,20 @@ def resolve_channel_selections( ] if has_unit_cols: aliased_solve_cols.extend([source_unit_col, target_unit_col]) - aliased_for_union = aliased_channel_metrics_df.select(*aliased_solve_cols) + # An aliased selector matching several aliases of a container resolves to + # one channel, like a direct selector: best alias priority, then lowest + # channel_id. + per_selector = Window.partitionBy( + self.config.container_id_col, F.col("selector_ids")[0] + ).orderBy( + F.col(self.config.alias_priority_col).asc_nulls_last(), + F.col(self.config.channel_id_col).asc(), + ) + aliased_for_union = ( + aliased_channel_metrics_df.withColumn("_rank", F.row_number().over(per_selector)) + .where(F.col("_rank") == 1) + .select(*aliased_solve_cols) + ) merged = channel_metrics_df.unionByName( aliased_for_union, allowMissingColumns=has_unit_cols @@ -1621,9 +1654,23 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame ------- pyspark.sql.DataFrame Narrow DataFrame of calculated-channel samples. + + Notes + ----- + When every selection is a pushdown as a whole (see + :meth:`PushdownPlan.native_targets`), the rows are produced in Spark from + the derived rows and no pandas UDF runs. """ col_map = self.config.col_map q, joined_df, container_count, plan = self._prepare_channels_join(query, channels_df) + self._validate_calculated_channel_id_type(q) + + targets = plan.native_targets(selections) if plan is not None else None + if targets is not None: + logger.info("Calculated channels solved natively in Spark, without the pandas UDF") + native = self._native_calculated_channel_rows(joined_df, q, targets) + return self._attach_identity(native, selections) + container_tag_cols, container_metric_cols = self._container_meta_cols(query) schema = self._build_calculated_channels_udf_schema(q) @@ -1647,8 +1694,14 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame container_meta_cols=[*container_tag_cols, *container_metric_cols], ) - # Identity is constant per channel_id, so attach it here via a - # channel_id-keyed CASE instead of the UDF shipping a copy per row. + return self._attach_identity(res, selections) + + def _attach_identity(self, rows: DataFrame, selections) -> DataFrame: + """Add the ``identity`` map of each calculated channel to its rows. + + Identity is constant per channel_id, so it is attached via a + channel_id-keyed CASE instead of being shipped with every row. + """ channel_id_col = self.config.channel_id_col identity_expr = None for cc in selections: @@ -1659,4 +1712,123 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame if identity_expr is None else identity_expr.when(branch, id_map) ) - return res.withColumn("identity", identity_expr) + return rows.withColumn("identity", identity_expr) + + def _validate_calculated_channel_id_type(self, channels_q: DataFrame) -> None: + """Raise unless the channel-data ``channel_id`` can hold calculated-channel ids. + + :attr:`CalculatedChannel.channel_id` is a 31-bit integer, which only fits + ``IntegerType`` and ``LongType`` columns. + """ + dtype = channels_q.schema[self.config.channel_id_col].dataType + if not isinstance(dtype, T.IntegerType | T.LongType): + raise ValueError( + "Calculated channels need an integer or long channel_id column, got " + f"{dtype.simpleString()}." + ) + + def _native_calculated_channel_rows( + self, + joined_df: DataFrame, + channels_q: DataFrame, + targets: dict[int, dict[int, list[tuple[int, bool]]]], + ) -> DataFrame: + """Calculated-channel rows taken directly from the pushed-down rows. + + *targets* (:meth:`PushdownPlan.native_targets`) maps each variant and + selector to the calculated channels built from them. Every derived row + is emitted once per such channel, so the result equals the pandas + evaluation: values get the conversion factor for aliased selectors, and + the round-trip through float64 and Arrow is mirrored (integral times via + double, NaN as null). + + Returns + ------- + pyspark.sql.DataFrame + ``[container_id, channel_id, tstart, tend, value]`` as typed by + :meth:`_build_calculated_channels_udf_schema`. + """ + cfg = self.config + cid, ch = cfg.container_id_col, cfg.channel_id_col + ts, te, val, conv = cfg.tstart_col, cfg.tend_col, cfg.value_col, cfg.conversion_factor_col + target_type = T.ArrayType( + T.StructType( + [ + T.StructField("_cc_channel_id", T.LongType()), + T.StructField("_cc_uses_alias", T.BooleanType()), + ] + ) + ) + by_variant = F.create_map( + *[ + item + for variant, by_selector in sorted(targets.items()) + for item in ( + F.lit(variant), + F.create_map( + *[ + entry + for selector_id, channels in sorted(by_selector.items()) + for entry in ( + F.lit(selector_id).cast("long"), + F.array( + *[ + F.struct( + F.lit(channel_id).cast("long").alias("_cc_channel_id"), + F.lit(uses_alias).alias("_cc_uses_alias"), + ) + for channel_id, uses_alias in channels + ] + ), + ) + ] + ), + ) + ] + ) + by_selector = F.try_element_at(by_variant, F.col(VARIANT_COL)) + selector_ids = F.array_distinct( + F.filter( + F.transform("selector_ids", lambda sid: sid.cast("long")), + lambda sid: sid.isNotNull(), + ) + ) + no_targets = F.array().cast(target_type) + channel_targets = F.flatten( + F.transform( + selector_ids, + lambda sid: F.coalesce(F.try_element_at(by_selector, sid), no_targets), + ) + ) + has_conversion = conv in joined_df.columns + rows = joined_df.where(F.col(VARIANT_COL) != RAW_VARIANT).select( + cid, + ts, + te, + val, + *([conv] if has_conversion else []), + F.inline(channel_targets), + ) + + value = F.col(val).cast("double") + if has_conversion: + factor = F.col(conv) + value = F.when(F.col("_cc_uses_alias") & factor.isNotNull(), value * factor).otherwise( + value + ) + + def as_udf_output(col: Column, dtype: T.DataType) -> Column: + # The UDF emits float64 arrays; Arrow casts them back and masks NaN as null. + col = col.cast("double") + if isinstance(dtype, T.IntegralType): + return col.cast(dtype) + return F.when(F.isnan(col), F.lit(None)).otherwise(col).cast(dtype) + + schema = channels_q.schema + return rows.select( + cid, + F.col("_cc_channel_id").cast(schema[ch].dataType).alias(ch), + as_udf_output(F.col(ts), schema[ts].dataType).alias(ts), + as_udf_output(F.col(te), schema[te].dataType).alias(te), + as_udf_output(value, T.DoubleType()).alias(val), + ) diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py new file mode 100644 index 00000000..2466d93e --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py @@ -0,0 +1,108 @@ +# pylint: disable=missing-function-docstring +"""PushdownPlan: channel variant tagging and native calculated-channel targets.""" + +import pyspark.sql.functions as F +import pyspark.sql.types as T +import pytest + +from impulse_query_engine.analyze.metadata.tag_expression import TagSelector +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel +from impulse_query_engine.analyze.query.optimizer import ( + PushdownContext, + QueryOptimizer, + ResamplePushdownRule, +) +from tests.conftest import spark + +_LONG = PushdownContext(T.LongType(), T.LongType()) + + +def _channel(name, **kwargs): + return TimeSeriesSelector(TagSelector("channel_name") == name, **kwargs) + + +def _cc(expr, name): + return CalculatedChannel(expr, {"channel_name": name}) + + +@pytest.fixture +def selectors(): + return {"pushed": _channel("rpm"), "pushed_2": _channel("spd"), "raw": _channel("temp")} + + +def test_channel_variants(spark, selectors): + pushed, pushed_2, raw = selectors["pushed"], selectors["pushed_2"], selectors["raw"] + plan = QueryOptimizer().optimize( + [pushed.resample(1e6), pushed_2.resample(5e5), raw.max()], _LONG + ) + pushed_variant = plan.variants_by_selector[pushed.selector_id] + pushed_2_variant = plan.variants_by_selector[pushed_2.selector_id] + cases = { + "none": ([], [0]), + "null": ([None], [0]), + "raw": ([raw.selector_id], [0]), + "pushed": ([pushed.selector_id], list(pushed_variant)), + "raw_and_pushed": ([raw.selector_id, pushed.selector_id], [0]), + "pushed_and_null": ([pushed.selector_id, None], list(pushed_variant)), + "two_pushed": ( + [pushed.selector_id, pushed_2.selector_id], + list(pushed_variant + pushed_2_variant), + ), + } + df = spark.createDataFrame( + [(name, ids) for name, (ids, _) in cases.items()], + schema="name string, selector_ids array", + ) + actual = { + r["name"]: r["variants"] + for r in df.select( + "name", plan.channel_variants(F.col("selector_ids")).alias("variants") + ).collect() + } + assert actual == {name: expected for name, (_, expected) in cases.items()} + + +class TestNativeTargets: + def test_whole_expression_pushdowns(self, selectors): + rpm, spd = selectors["pushed"], selectors["pushed_2"] + aliased = _channel("vehicle_speed", uses_alias=True) + channels = [ + _cc(rpm.resample(1e6), "rpm_1s"), + _cc(rpm.resample(1e6), "rpm_1s_copy"), + _cc(spd.resample(5e5), "spd"), + _cc(aliased.resample(1e6), "vehicle_speed"), + ] + plan = QueryOptimizer().optimize(channels, _LONG) + v_1s = plan.variant_ids[("resample", (("sample_rate", 1e6),))] + v_half = plan.variant_ids[("resample", (("sample_rate", 5e5),))] + assert plan.native_targets(channels) == { + v_1s: { + rpm.selector_id: [ + (channels[0].channel_id, False), + (channels[1].channel_id, False), + ], + aliased.selector_id: [(channels[3].channel_id, True)], + }, + v_half: {spd.selector_id: [(channels[2].channel_id, False)]}, + } + + def test_needs_every_selection_to_be_a_whole_pushdown(self, selectors): + rpm, spd = selectors["pushed"], selectors["pushed_2"] + native = _cc(rpm.resample(1e6), "rpm_1s") + for other in (_cc(spd.resample(1e6) * 2, "spd_x2"), spd.resample(1e6)): + selections = [native, other] + plan = QueryOptimizer().optimize(selections, _LONG) + assert not plan.is_empty + assert plan.native_targets(selections) is None + assert QueryOptimizer().optimize([native], _LONG).native_targets([]) is None + + def test_rule_with_custom_decode_is_not_native(self, selectors): + class _CustomDecode(ResamplePushdownRule): + def decode(self, rows, pushdown, cache): + return super().decode(rows, pushdown, cache) + + channels = [_cc(selectors["pushed"].resample(1e6), "rpm_1s")] + plan = QueryOptimizer([_CustomDecode()]).optimize(channels, _LONG) + assert not plan.is_empty + assert plan.native_targets(channels) is None diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py new file mode 100644 index 00000000..d2be8d8a --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py @@ -0,0 +1,183 @@ +# pylint: disable=missing-function-docstring +"""DefaultSolver resolves every selector to at most one channel per container. + +When several channels of a container match a selection, the lowest ``channel_id`` +wins (after the alias priority for aliased selections), independent of row order. +""" + +import os +import zlib + +import numpy as np +import pandas as pd +import pyspark.sql.functions as F +import pytest + +from impulse_query_engine.analyze.metadata.tag_expression import TagSelector +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver +from impulse_query_engine.analyze.query.solvers.solver_config import ( + ChannelMappingConfig, + SolverConfig, + TableConfig, +) +from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig +from tests.conftest import spark + +_ALIAS_SCHEMA = "spark_catalog.silver_key_value_store_alias" + + +def _alias_config() -> SolverConfig: + return SolverConfig( + project_id="SAMPLE_PROJECT", + container_metrics=TableConfig(column_name_mapping={"project": "project_id"}), + channel_mapping=ChannelMappingConfig(filters={"toolbox_id": "container_concept"}), + ) + + +def _resolved(spark, solver, query, selector) -> dict: + """``{container_id: [channel_id, ...]}`` the filter pipeline resolves *selector* to.""" + matches = query.select(selector)._run_filter_pipeline(spark, solver, None) + rows = matches.where( + F.exists("selector_ids", lambda sid: sid.cast("long") == selector.selector_id) + ).collect() + resolved: dict = {} + for row in rows: + resolved.setdefault(row["container_id"], []).append(row["channel_id"]) + return {cid: sorted(channels) for cid, channels in resolved.items()} + + +def _alias_tables(spark) -> dict: + names = ( + "container_tags", + "container_metrics", + "channel_metrics", + "channels", + "channel_mapping", + "channel_tags", + ) + return {name: spark.read.table(f"{_ALIAS_SCHEMA}.{name}") for name in names} + + +def _alias_db(tables, ws, with_channel_tags=False) -> MeasurementDB: + tables = dict(tables) + if not with_channel_tags: + tables.pop("channel_tags") + cfg = MeasurementDBConfig.for_debug(tables) + cfg.channel_mapping_table = "channel_mapping" + return MeasurementDB(cfg, ws=ws) + + +def _with_channel(channel_metrics, container_id, channel_id, **columns): + """*channel_metrics* plus a copy of container 1's Engine RPM row as a new channel.""" + template = channel_metrics.where( + (F.col("container_id") == 1) & (F.col("channel_name") == "Engine RPM") + ) + row = template.withColumn("container_id", F.lit(container_id).cast("long")).withColumn( + "channel_id", F.lit(channel_id).cast(channel_metrics.schema["channel_id"].dataType) + ) + for name, value in columns.items(): + row = row.withColumn(name, F.lit(value)) + return channel_metrics.unionByName(row.select(*channel_metrics.columns)) + + +class TestDirectSelectors: + def test_wide_mode_keeps_lowest_channel_id(self, spark, basic_narrow_db): + # unit="C" matches the ambient and intake temperatures: channels 6 and 9 in + # containers 1 and 2, channels 2 and 8 in container 3. + q = basic_narrow_db.query + celsius = q.channel(unit="C") + assert _resolved(spark, DefaultSolver(spark), q, celsius) == {1: [6], 2: [6], 3: [2]} + + def test_wide_mode_solves_on_the_resolved_channel(self, spark, basic_narrow_db): + q = basic_narrow_db.query + result = q.select(q.channel(unit="C").max().alias("t_max")).solve( + spark, solver=DefaultSolver(spark) + ) + actual = {r["container_id"]: r["t_max"] for r in result.collect()} + + base_path = os.path.dirname(os.path.abspath(__file__)) + base_path = base_path[: base_path.find("tests")] + data = pd.read_csv(f"{base_path}/tests/unit/data/basic_narrow_csv/channel_data.csv") + expected = { + cid: np.nanmax(data[(data.container_id == cid) & (data.channel_id == ch)].value) + for cid, ch in {1: 6, 2: 6, 3: 2}.items() + } + assert actual == pytest.approx(expected) + assert all(np.isfinite(v) for v in actual.values()) + + def test_eav_mode_keeps_lowest_channel_id( + self, spark, setup_key_value_store_alias_db, mock_workspace_client + ): + # Re-tag channel 9 of container 1 and channels 2 and 8 of container 3 as + # "Engine RPM": container 1 keeps 5 (the original), container 3 keeps 2. + tables = _alias_tables(spark) + tags = tables["channel_tags"] + retagged = [(1, 9), (3, 2), (3, 8)] + keep = ~F.lit(False) + for cid, ch in retagged: + keep = keep & ~((F.col("container_id") == cid) & (F.col("channel_id") == ch)) + extra = spark.createDataFrame( + [(cid, ch, "channel_name", "Engine RPM") for cid, ch in retagged], + schema=tags.schema, + ) + tables["channel_tags"] = tags.where(keep).unionByName(extra) + db = _alias_db(tables, mock_workspace_client, with_channel_tags=True) + + q = db.query + rpm = q.channel(channel_name="Engine RPM") + assert _resolved(spark, DefaultSolver(spark, config=_alias_config()), q, rpm) == { + 1: [5], + 2: [5], + 3: [2], + } + + +class TestAliasedSelectors: + @pytest.mark.parametrize(("extra_channel", "expected"), [(2, 2), (8, 5)]) + def test_priority_tie_keeps_lowest_channel_id( + self, spark, setup_key_value_store_alias_db, mock_workspace_client, extra_channel, expected + ): + # Container 1 gets a second engine_speed candidate (EngSpd); both candidates + # have no priority, so the lower channel_id wins. + tables = _alias_tables(spark) + tables["channel_metrics"] = _with_channel( + tables["channel_metrics"], + 1, + extra_channel, + channel_name="EngSpd", + data_key="ProjSpecREC_10Hz", + ) + db = _alias_db(tables, mock_workspace_client) + q = db.query + engine_speed = q.channel_with_alias(channel_alias="engine_speed") + solver = DefaultSolver(spark, config=_alias_config()) + assert _resolved(spark, solver, q, engine_speed)[1] == [expected] + + def test_selector_matching_several_aliases_resolves_one_channel( + self, spark, key_value_store_alias_db + ): + q = key_value_store_alias_db.query + solver = DefaultSolver(spark, config=_alias_config()) + per_alias = [ + _resolved(spark, solver, q, q.channel_with_alias(channel_alias=alias)) + for alias in ("engine_speed", "vehicle_speed") + ] + both = TimeSeriesSelector( + (TagSelector("channel_alias") == "engine_speed") + | (TagSelector("channel_alias") == "vehicle_speed"), + uses_alias=True, + ) + resolved = _resolved(spark, solver, q, both) + assert resolved + for cid, channels in resolved.items(): + candidates = [ch for one in per_alias for ch in one.get(cid, [])] + assert channels == [min(candidates)] + + +def test_aliased_and_direct_selectors_on_the_same_tags_have_distinct_ids(): + expr = TagSelector("channel_name") == "Engine RPM" + direct = TimeSeriesSelector(expr) + aliased = TimeSeriesSelector(expr, uses_alias=True) + assert direct.selector_id == zlib.crc32(str(expr).encode()) + assert aliased.selector_id != direct.selector_id diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py new file mode 100644 index 00000000..a57c172b --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py @@ -0,0 +1,247 @@ +# pylint: disable=missing-function-docstring +"""Calculated channels that are pushdowns as a whole are solved without the pandas UDF. + +Each case solves the same calculated channels with pushdown off (pandas UDF) and on +(native when every channel qualifies) and requires identical rows, types, nulls and +NaNs. +""" + +import pandas as pd +import pyspark.sql.functions as F +import pytest +from pyspark.sql import DataFrame + +from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver +from impulse_query_engine.analyze.query.solvers.solver_config import ( + ChannelMappingConfig, + SolverConfig, + TableConfig, +) +from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig +from tests.conftest import spark + +RATE = 1e7 # 10 s in µs; the fixture channels are sampled roughly every 0.9 s +_KVS_SCHEMA = "spark_catalog.silver_key_value_store" + + +def _kvs_config() -> SolverConfig: + return SolverConfig( + project_id="SAMPLE_PROJECT", + container_tags=TableConfig(column_name_mapping={"element_id": "key"}), + container_metrics=TableConfig(column_name_mapping={"project": "project_id"}), + channel_mapping=ChannelMappingConfig(), + ) + + +def _unit_conversion_config() -> SolverConfig: + return SolverConfig( + project_id="SAMPLE_PROJECT", + container_metrics=TableConfig(column_name_mapping={"project": "project_id"}), + channel_mapping=ChannelMappingConfig(filters={"toolbox_id": "container_concept"}), + ) + + +@pytest.fixture +def udf_calls(monkeypatch) -> list: + """Number of grouped-map UDF runs, per solve.""" + calls: list[int] = [] + original = DefaultSolver._apply_grouped_map + + def _counting(self, *args, **kwargs): + calls[-1] += 1 + return original(self, *args, **kwargs) + + monkeypatch.setattr(DefaultSolver, "_apply_grouped_map", _counting) + return calls + + +def _solve(spark, db, make_channels, enable_pushdown, udf_calls=None, **solver_kwargs): + if udf_calls is not None: + udf_calls.append(0) + q = db.query + solver = DefaultSolver(spark, enable_pushdown=enable_pushdown, **solver_kwargs) + return q.select(*make_channels(q)).solve_calculated_channels(spark, solver=solver) + + +def _nan_and_null_counts(df: DataFrame) -> tuple: + return tuple( + df.select( + F.sum(F.isnan("value").cast("int")), F.sum(F.col("value").isNull().cast("int")) + ).first() + ) + + +def _rows(df: DataFrame) -> pd.DataFrame: + pdf = df.toPandas() + pdf["identity"] = pdf["identity"].map(lambda m: tuple(sorted(m.items()))) + return pdf.sort_values(["container_id", "channel_id", "tstart", "tend"]).reset_index(drop=True) + + +def _assert_same(native: DataFrame, udf: DataFrame) -> pd.DataFrame: + assert [(f.name, f.dataType) for f in native.schema] == [ + (f.name, f.dataType) for f in udf.schema + ] + assert _nan_and_null_counts(native) == _nan_and_null_counts(udf) + expected = _rows(udf) + assert len(expected) > 0 + pd.testing.assert_frame_equal(_rows(native), expected, check_exact=True) + return expected + + +def _solve_both(spark, db, make_channels, udf_calls, **solver_kwargs) -> pd.DataFrame: + udf = _solve(spark, db, make_channels, False, udf_calls, **solver_kwargs) + native = _solve(spark, db, make_channels, True, udf_calls, **solver_kwargs) + return _assert_same(native, udf), native + + +class TestNativeCalculatedChannels: + def test_resampled_channels_skip_the_udf(self, spark, key_value_store_db, udf_calls): + def channels(q): + return [ + CalculatedChannel( + q.channel(channel_name="Engine RPM").resample(RATE), {"channel_name": "rpm"} + ), + CalculatedChannel( + q.channel(channel_name="Vehicle Speed Sensor").resample(RATE / 2), + {"channel_name": "speed", "data_key": "CALC"}, + ), + ] + + rows, native = _solve_both( + spark, key_value_store_db, channels, udf_calls, config=_kvs_config() + ) + assert udf_calls == [1, 0] + assert rows["value"].gt(0).any() + assert rows["channel_id"].nunique() == 2 + plan = native._jdf.queryExecution().executedPlan().toString() + assert "FlatMapGroupsInPandas" not in plan + + def test_same_expression_under_two_identities_and_two_rates( + self, spark, key_value_store_db, udf_calls + ): + def channels(q): + rpm = q.channel(channel_name="Engine RPM") + return [ + CalculatedChannel(rpm.resample(RATE), {"channel_name": "rpm_a"}), + CalculatedChannel(rpm.resample(RATE), {"channel_name": "rpm_b"}), + CalculatedChannel(rpm.resample(RATE / 4), {"channel_name": "rpm_c"}), + ] + + rows, _ = _solve_both(spark, key_value_store_db, channels, udf_calls, config=_kvs_config()) + assert udf_calls == [1, 0] + assert rows["channel_id"].nunique() == 3 + + def test_mixed_selection_uses_the_udf(self, spark, key_value_store_db, udf_calls): + def channels(q): + rpm = q.channel(channel_name="Engine RPM") + return [ + CalculatedChannel(rpm.resample(RATE), {"channel_name": "rpm"}), + CalculatedChannel(rpm.resample(RATE) * 2, {"channel_name": "rpm_x2"}), + ] + + _solve_both(spark, key_value_store_db, channels, udf_calls, config=_kvs_config()) + assert udf_calls == [1, 1] + + def test_aliased_and_direct_selection_of_one_channel_with_unit_conversion( + self, spark, key_value_store_unit_conversion_db, udf_calls + ): + # vehicle_speed resolves to the "Vehicle Speed Sensor" channel; only the + # aliased selection converts its values. + def channels(q): + return [ + CalculatedChannel( + q.channel_with_alias(channel_alias="vehicle_speed").resample(RATE), + {"channel_name": "speed_converted"}, + ), + CalculatedChannel( + q.channel(channel_name="Vehicle Speed Sensor").resample(RATE), + {"channel_name": "speed_raw"}, + ), + ] + + rows, _ = _solve_both( + spark, + key_value_store_unit_conversion_db, + channels, + udf_calls, + config=_unit_conversion_config(), + ) + assert udf_calls == [1, 0] + by_channel = rows.groupby("channel_id")["value"].sum() + assert len(by_channel) == 2 and by_channel.iloc[0] != by_channel.iloc[1] + + def test_raw_data(self, spark, setup_raw_channels_db, mock_workspace_client, udf_calls): + tables = { + name: spark.read.table(f"spark_catalog.silver_raw.{name}") + for name in ("container_metrics", "channel_metrics", "channels") + } + db = MeasurementDB(MeasurementDBConfig.for_debug(tables), ws=mock_workspace_client) + + def channels(q): + rpm = q.channel(channel_name="Engine RPM") + return [CalculatedChannel(rpm.resample(2e6), {"channel_name": "rpm_2s"})] + + _solve_both(spark, db, channels, udf_calls, is_raw_data=True) + assert udf_calls == [1, 0] + + def test_ns_timestamps_nan_values_and_a_single_sample( + self, spark, setup_key_value_store_db, mock_workspace_client, udf_calls + ): + # Timestamps beyond 2^53 (not exact in float64), NaN values, and container 3's + # Engine RPM reduced to a single sample (passed through by resample). + tables = { + name: spark.read.table(f"{_KVS_SCHEMA}.{name}") + for name in ("container_tags", "container_metrics", "channel_metrics", "channels") + } + data = tables["channels"] + rpm_3 = (F.col("container_id") == 3) & (F.col("channel_id") == 5) + first_3 = data.where(rpm_3).agg(F.min("tstart")).first()[0] + rpm_1 = (F.col("container_id") == 1) & (F.col("channel_id") == 5) + first_1 = data.where(rpm_1).agg(F.min("tstart")).first()[0] + tables["channels"] = ( + data.where(~rpm_3 | (F.col("tstart") == first_3)) + .withColumn("tstart", F.col("tstart") * 1000 + 1) + .withColumn("tend", F.col("tend") * 1000 + 1) + .withColumn( + "value", + # tstart is already scaled here; every sample but the first becomes NaN. + F.when( + rpm_1 & (F.col("tstart") > first_1 * 1000 + 1), F.lit(float("nan")) + ).otherwise(F.col("value").cast("double")), + ) + ) + db = MeasurementDB(MeasurementDBConfig.for_debug(tables), ws=mock_workspace_client) + + def channels(q): + rpm = q.channel(channel_name="Engine RPM") + return [CalculatedChannel(rpm.resample(RATE * 1000), {"channel_name": "rpm_ns"})] + + rows, native = _solve_both(spark, db, channels, udf_calls, config=_kvs_config()) + assert udf_calls == [1, 0] + assert rows["tstart"].max() > 2**53 + assert len(rows[rows.container_id == 3]) == 1 + assert _nan_and_null_counts(native)[1] > 0 + + +@pytest.mark.parametrize("channel_id_type", ["short", "string"]) +@pytest.mark.parametrize("enable_pushdown", [False, True]) +def test_channel_id_must_be_integer_or_long( + spark, setup_key_value_store_db, mock_workspace_client, channel_id_type, enable_pushdown +): + tables = { + name: spark.read.table(f"{_KVS_SCHEMA}.{name}") + for name in ("container_tags", "container_metrics", "channel_metrics", "channels") + } + for name in ("channel_metrics", "channels"): + tables[name] = tables[name].withColumn( + "channel_id", F.col("channel_id").cast(channel_id_type) + ) + db = MeasurementDB(MeasurementDBConfig.for_debug(tables), ws=mock_workspace_client) + + def channels(q): + rpm = q.channel(channel_name="Engine RPM") + return [CalculatedChannel(rpm.resample(RATE), {"channel_name": "rpm"})] + + with pytest.raises(ValueError, match="integer or long channel_id"): + _solve(spark, db, channels, enable_pushdown, config=_kvs_config()) diff --git a/tests/impulse_reporting/integration/calculated_channel_test.py b/tests/impulse_reporting/integration/calculated_channel_test.py index 481f7dea..090d821b 100644 --- a/tests/impulse_reporting/integration/calculated_channel_test.py +++ b/tests/impulse_reporting/integration/calculated_channel_test.py @@ -9,6 +9,7 @@ from unittest.mock import create_autospec +import pandas as pd import pyspark.sql.functions as F import pyspark.sql.types as T import pytest @@ -629,3 +630,51 @@ def test_channel_metrics_custom_kpis(spark): assert "max" in metrics.columns assert "mean" not in metrics.columns assert "duration" not in metrics.columns + + +def test_resampled_channel_facts_are_identical_with_pushdown(spark, monkeypatch): + """A resample-only calculated channel is solved in Spark without the pandas UDF; + its gold facts equal the ones written without pushdown.""" + from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver + + udf_runs = [] + original = DefaultSolver._apply_grouped_map + + def _counting(self, *args, **kwargs): + udf_runs[-1] += 1 + return original(self, *args, **kwargs) + + monkeypatch.setattr(DefaultSolver, "_apply_grouped_map", _counting) + facts = [] + for enable_pushdown in (False, True): + udf_runs.append(0) + config = _config() + config.query_engine = QueryEngine(enable_pushdown=enable_pushdown) + report = Report( + name="calc_channel_report", + spark=spark, + workspace_client=create_autospec(WorkspaceClient), + config=dict(config), + ) + q = report.get_db().query + report.add_calculated_channel( + CalculatedChannel( + name="speed_10s", + expr=q.channel(channel_name="Vehicle Speed Sensor").resample(1e7), + identity={"channel_name": "speed_10s", "data_key": "CALC"}, + ) + ) + report.determine_report() + report.persist_results() + facts.append( + spark.read.table(_FACT) + .drop("_created_at") + .toPandas() + .sort_values(["container_id", "channel_id", "tstart"]) + .reset_index(drop=True) + ) + off, on = facts + assert udf_runs == [1, 0] + assert len(off) > 0 + assert off["value"].gt(0).any() + pd.testing.assert_frame_equal(on, off, check_exact=True) From 155d7405498a9cb7924a70e83c955cd27f0fdc72 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Thu, 15 Oct 2026 07:51:34 +0200 Subject: [PATCH 07/16] feat(query-engine): chunk resample pushdown to bound per-task memory Update `ResamplePushdownRule` to process each series in equal-time chunks of about `chunk_rows` samples (default 100,000) instead of whole series. Per-chunk windows keep worker memory bounded; chunk-boundary state (gaps, block start) is resolved from a small per-chunk summary broadcast-joined back. The summary shares one extra narrow read of the channel key/time columns, so Spark reads channel data once in full plus one narrow read. - Add `chunk_rows` parameter with validation. - Update `PushdownRule.to_spark` contract to allow broadcast-joining small summaries derived from `rows`. - Stop filtering raw variant rows before the pushdown input so the narrow summary read can be shared across rules. - Add unit tests for chunking, chunk-size validation, random series across chunk sizes, and parity with existing cases. - Update API docs and query solver docs to describe chunked resampling and the two-channel-data-read plan. --- .../analyze/query/optimizer/pushdown.md | 11 +- .../analyze/query/optimizer/resample_rule.md | 9 + .../references/query_engine/query_solvers.md | 16 +- .../analyze/query/optimizer/pushdown.py | 11 +- .../analyze/query/optimizer/resample_rule.py | 181 ++++++++++++++---- .../analyze/query/solvers/default_solver.py | 5 +- .../query/optimizer/resample_rule_test.py | 129 +++++++++++-- ..._solver_native_calculated_channels_test.py | 7 + .../solvers/default_solver_pushdown_test.py | 39 ++++ 9 files changed, 343 insertions(+), 65 deletions(-) diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md index 74dbe7c3..fec552b9 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -132,10 +132,13 @@ def to_spark(rows: DataFrame, variants: Mapping[int, Pushdown], Replace the rows of this rule's *variants* with their derived rows. -*rows* is read once and shared by all rules, so implementations must -stay a single lineage (no self-joins or unions of *rows*) and keep its -container partitioning (windows partitioned by container id, channel id -and :data:`VARIANT_COL` reuse it). +*rows* is shared by all rules and partitioned by container; the result +must keep that partitioning (windows partitioned by container id, +channel id and :data:`VARIANT_COL` reuse it). Implementations must not +union *rows* or join it to a non-broadcast frame. Small summaries +derived from *rows* (per series or per chunk) may be broadcast-joined +back; Spark reads the channel data again for them, so summaries should +read only the key and time columns to share one narrow read. **Arguments**: diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md index 18a9fd04..e5720897 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md @@ -26,9 +26,18 @@ Each sample is assigned the bucket ``c`` of the first grid point at or after its and every row carries just the grid points it covers. Values are looked up within a block; this matches the pandas evaluation unless intervals overlap. +Every series is cut into equal time slices (chunks) holding about *chunk_rows* +samples on average, and the windows over samples run per chunk, so a worker +holds one chunk of a series at a time instead of the whole series. Gaps at +chunk boundaries, the block a chunk starts in and that block's start are +resolved from one summary row per chunk. + Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. numpy builds that fuse this into a multiply-add can differ by one ulp for fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral grids are exact. +**Arguments**: + +- `chunk_rows` (`int`): Average number of samples per chunk (default 100 000). diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index 16614469..44f818fd 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -206,10 +206,18 @@ that share the same `tstart` is unspecified, as in the pandas evaluation, and ch `[tstart, tend)` intervals overlap are not supported (the RLE and interval encoders never produce overlaps). -Spark processes each channel in one pass over its samples, sorted by `tstart`: every sample is -assigned the grid bucket it falls into, and only the last sample of each bucket emits rows (the -grid points up to the next non-empty bucket). The memory per row is bounded by the grid points a -single sample covers, not by the length of the channel. +Spark cuts each channel into chunks: equal time slices that hold about 100 000 samples on +average. It processes one chunk at a time, sorted by `tstart`: every sample is assigned the grid +bucket it falls into, and only the last sample of each bucket emits rows (the grid points up to +the next non-empty bucket). What crosses a chunk boundary comes from one summary row per chunk: +whether a gap starts there, and which block (and block start) the chunk continues. A task +therefore holds one chunk of a channel, plus the resampled rows of one continuous block, rather +than the whole channel. If the sampling density varies strongly within a channel, a chunk can hold +more samples. + +The summaries need the channel's time columns a second time, so with pushdown Spark reads the +channel data twice: once in full, and once with only `container_id`, `channel_id`, `tstart` and +`tend`. The full rows are still shuffled only once. ### Calculated channels without the pandas UDF diff --git a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py index 5443d6c5..2af656cc 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -120,10 +120,13 @@ def to_spark( ) -> DataFrame: """Replace the rows of this rule's *variants* with their derived rows. - *rows* is read once and shared by all rules, so implementations must - stay a single lineage (no self-joins or unions of *rows*) and keep its - container partitioning (windows partitioned by container id, channel id - and :data:`VARIANT_COL` reuse it). + *rows* is shared by all rules and partitioned by container; the result + must keep that partitioning (windows partitioned by container id, + channel id and :data:`VARIANT_COL` reuse it). Implementations must not + union *rows* or join it to a non-broadcast frame. Small summaries + derived from *rows* (per series or per chunk) may be broadcast-joined + back; Spark reads the channel data again for them, so summaries should + read only the key and time columns to share one narrow read. Parameters ---------- diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py index 529a2b06..630bbd9d 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -43,14 +43,32 @@ class ResamplePushdownRule(PushdownRule): and every row carries just the grid points it covers. Values are looked up within a block; this matches the pandas evaluation unless intervals overlap. + Every series is cut into equal time slices (chunks) holding about *chunk_rows* + samples on average, and the windows over samples run per chunk, so a worker + holds one chunk of a series at a time instead of the whole series. Gaps at + chunk boundaries, the block a chunk starts in and that block's start are + resolved from one summary row per chunk. + Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. numpy builds that fuse this into a multiply-add can differ by one ulp for fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral grids are exact. + + Parameters + ---------- + chunk_rows : int, optional + Average number of samples per chunk (default 100 000). """ name = "resample" + def __init__(self, chunk_rows: int = 100_000): + if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, numbers.Integral): + raise TypeError(f"chunk_rows must be an integer, got {chunk_rows!r}") + if chunk_rows < 1: + raise ValueError(f"chunk_rows must be positive, got {chunk_rows}") + self.chunk_rows = int(chunk_rows) + def match(self, node: TimeSeriesExpression) -> Pushdown | None: if ( type(node) is not TimeSeriesOp @@ -96,34 +114,104 @@ def to_spark( ), F.col(VARIANT_COL), ) + ts_d, te_d = F.col(ts).cast("double"), F.col(te).cast("double") + # The pandas cache's sample order. + order = F.struct(ts, te) + + # Series boundaries. This summary and the chunk summary below read only the + # key and time columns, so Spark serves both from one extra narrow read. + series = rows.groupBy(*key).agg( + F.min(ts_d).alias("_rs_first"), + F.max(_ignore_nan(te_d)).alias("_rs_end"), + F.count(F.lit(1)).alias("_rs_n"), + ) + df = rows.join(F.broadcast(series), on=key) + + # Equal time slices of about chunk_rows samples; monotone in tstart, so the + # chunks of a series follow its sample order. + first, end = F.col("_rs_first"), F.col("_rs_end") + n_chunks = F.ceil(F.col("_rs_n") / F.lit(self.chunk_rows)) + df = df.withColumn( + "_rs_chunk", + F.when( + end > first, + F.least(F.floor((ts_d - first) / (end - first) * n_chunks), n_chunks - 1), + ).otherwise(F.lit(0).cast("long")), + ) - # Pass 1, over each series in the pandas cache's tstart order. Only running, - # whole-partition and neighbour windows (all linear in the series length); - # they share one sort and keep the container partitioning. - ordered = Window.partitionBy(*key).orderBy(ts, te) - running = ordered.rowsBetween(Window.unboundedPreceding, Window.currentRow) - whole = ordered.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) - ts_d, te_d, block = F.col("_rs_ts"), F.col("_rs_te"), F.col("_rs_block") - prev_te = F.lag(te_d).over(ordered) + # Pass 1, per chunk: only running and neighbour windows (linear in the chunk + # size), sharing one sort and the container partitioning. A chunk splits into + # segments at its gaps; segment 0 continues whatever block the previous chunk + # ended in, which the chunk summary resolves. + in_chunk = Window.partitionBy(*key, "_rs_chunk").orderBy(ts, te) + running = in_chunk.rowsBetween(Window.unboundedPreceding, Window.currentRow) + prev_te = F.lag(te_d).over(in_chunk) + is_break = F.col("_rs_break") == 1 df = ( - rows.withColumn("_rs_ts", F.col(ts).cast("double")) - .withColumn("_rs_te", F.col(te).cast("double")) - .withColumn("_rs_break", (prev_te.isNull() | ((ts_d - prev_te) != 0)).cast("long")) - .withColumn("_rs_n", F.count(F.lit(1)).over(whole)) - .withColumn("_rs_last_ts", F.last(ts_d).over(whole)) - .withColumn("_rs_last_te", F.last(te_d).over(whole)) - .withColumn("_rs_block", F.sum("_rs_break").over(running)) + df.withColumn( + "_rs_break", + F.when(F.row_number().over(in_chunk) == 1, F.lit(0).cast("long")).otherwise( + (prev_te.isNull() | ((ts_d - prev_te) != 0)).cast("long") + ), + ) + .withColumn("_rs_seg", F.sum("_rs_break").over(running)) .withColumn( - "_rs_t_min", - F.last(F.when(F.col("_rs_break") == 1, ts_d), ignorenulls=True).over(running), + "_rs_seg_t_min", F.last(F.when(is_break, ts_d), ignorenulls=True).over(running) ) - # Largest tend so far within the block (blocks are numbered in order); on - # the block's last row this is the block's t_max. + # Largest tend so far within the segment (segments are numbered in order). .withColumn( - "_rs_t_max_so_far", - F.max(F.struct(block, _ignore_nan(te_d).alias("te"))).over(running).getField("te"), + "_rs_seg_t_max", + F.max(F.struct(F.col("_rs_seg"), _ignore_nan(te_d).alias("te"))) + .over(running) + .getField("te"), ) - .withColumn("_rs_nblk", F.max(block).over(whole)) + .withColumn("_rs_next_ts", F.lead(ts_d).over(in_chunk)) + .withColumn("_rs_next_break", F.lead("_rs_break").over(in_chunk)) + ) + + # One summary row per chunk, scanned in chunk order: whether the chunk starts + # a block, the block of its segment 0 and that block's start, plus the number + # of blocks and the last sample of the series. + chunks = df.groupBy(*key, "_rs_chunk").agg( + F.min_by(ts_d, order).alias("_rs_head_ts"), + F.max_by(F.struct(ts_d.alias("ts"), te_d.alias("te")), order).alias("_rs_tail"), + F.sum("_rs_break").alias("_rs_inner_breaks"), + F.max(F.when(is_break, ts_d)).alias("_rs_last_start"), + ) + by_chunk = Window.partitionBy(*key).orderBy("_rs_chunk") + all_chunks = by_chunk.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) + head_ts, head_break = F.col("_rs_head_ts"), F.col("_rs_head_break") == 1 + prev_tail_te = F.lag(F.col("_rs_tail.te")).over(by_chunk) + last_start = F.coalesce(F.col("_rs_last_start"), F.when(head_break, head_ts)) + chunks = ( + chunks.withColumn( + "_rs_head_break", + (prev_tail_te.isNull() | ((head_ts - prev_tail_te) != 0)).cast("long"), + ) + .withColumn("_rs_breaks", F.col("_rs_head_break") + F.col("_rs_inner_breaks")) + .withColumn( + "_rs_block_base", + F.sum("_rs_breaks").over( + by_chunk.rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + - F.col("_rs_inner_breaks"), + ) + .withColumn( + "_rs_head_t_min", + F.when(head_break, head_ts).otherwise( + F.last(last_start, ignorenulls=True).over( + by_chunk.rowsBetween(Window.unboundedPreceding, -1) + ) + ), + ) + .withColumn("_rs_nblk", F.sum("_rs_breaks").over(all_chunks)) + .withColumn("_rs_last", F.last("_rs_tail").over(all_chunks)) + .select(*key, "_rs_chunk", "_rs_block_base", "_rs_head_t_min", "_rs_nblk", "_rs_last") + ) + df = ( + df.join(F.broadcast(chunks), on=[*key, "_rs_chunk"]) + .withColumn("_rs_block", F.col("_rs_block_base") + F.col("_rs_seg")) + .withColumn("_rs_t_min", F.coalesce(F.col("_rs_seg_t_min"), F.col("_rs_head_t_min"))) ) # numpy.arange(t_min, t_max + rate, rate): length ceil(((t_max + rate) - t_min) / rate), @@ -139,26 +227,44 @@ def grid_point(k: Column) -> Column: ) resampled = F.col(VARIANT_COL).isin(sorted(variants)) & (F.col("_rs_n") >= 2) + next_in_block = F.coalesce(F.col("_rs_next_break") == 0, F.lit(False)) df = _with_first_grid_index( df, "_rs_c", ts_d, F.lit(False), resampled, grid_point, t_min, delta ) - c, next_c, next_block = F.col("_rs_c"), F.col("_rs_next_c"), F.lead(block).over(ordered) - df = ( - df.withColumn("_rs_next_c", F.lead(c).over(ordered)).withColumn( - "_rs_block_end", next_block.isNull() | (next_block != block) - ) - # Only the last sample of each bucket emits grid points. - .where(~resampled | F.col("_rs_block_end") | (next_c != c)) + df = _with_first_grid_index( + df, + "_rs_next_sample_c", + F.col("_rs_next_ts"), + F.lit(False), + resampled & next_in_block, + grid_point, + t_min, + delta, ) - - # Pass 2, over the remaining rows (about one per grid point): the block's - # t_max, and from it the number of grid points of the block. - t_max = F.max("_rs_t_max_so_far").over(Window.partitionBy(*key, "_rs_block")) - df = df.withColumn("_rs_t_max", t_max) - t_max = F.col("_rs_t_max") + c = F.col("_rs_c") + # Only the last sample of each bucket emits grid points. A chunk's last sample + # is kept too: its bucket may continue in the next chunk. + df = df.where(~resampled | ~next_in_block | (F.col("_rs_next_sample_c") != c)) + + # Pass 2, over the remaining rows of each block (about one per grid point): the + # next non-empty bucket, and the block's t_max (each segment's last sample + # carries its maximum). A chunk's last sample whose bucket continues in the + # next chunk is followed by its own bucket and emits nothing. Rows this rule + # leaves alone stay grouped by chunk. + in_block = Window.partitionBy( + *key, "_rs_block", F.when(~resampled, F.col("_rs_chunk")) + ).orderBy("_rs_c", ts, te) + whole_block = in_block.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) + df = df.withColumn("_rs_next_c", F.lead("_rs_c").over(in_block)).withColumn( + "_rs_t_max", F.max("_rs_seg_t_max").over(whole_block) + ) + next_c, t_max = F.col("_rs_next_c"), F.col("_rs_t_max") + block_end = next_c.isNull() # The last block keeps a grid point at t_max when the series ends with a # zero-duration sample. - inclusive = (block == F.col("_rs_nblk")) & (F.col("_rs_last_ts") == F.col("_rs_last_te")) + inclusive = (F.col("_rs_block") == F.col("_rs_nblk")) & ( + F.col("_rs_last.ts") == F.col("_rs_last.te") + ) length = F.greatest(F.ceil(((t_max + rate) - t_min) / rate), F.lit(0)) df = _with_first_grid_index( df, "_rs_grid_end", t_max, inclusive, resampled, grid_point, t_min, delta @@ -167,8 +273,7 @@ def grid_point(k: Column) -> Column: # A bucket's last sample emits the grid points up to the next non-empty bucket # of its block, or up to the block's end. df = df.withColumn( - "_rs_c_stop", - F.when(F.col("_rs_block_end"), n_grid).otherwise(F.least(next_c, n_grid)), + "_rs_c_stop", F.when(block_end, n_grid).otherwise(F.least(next_c, n_grid)) ) def element(kind: int, t: Column, tend: Column) -> Column: diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index f66443c7..1b4c599d 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -1801,7 +1801,10 @@ def _native_calculated_channel_rows( ) ) has_conversion = conv in joined_df.columns - rows = joined_df.where(F.col(VARIANT_COL) != RAW_VARIANT).select( + # Raw rows (variant 0) have no targets, so inline drops them. Not filtering + # them explicitly keeps the filter out of the pushdown rules' summaries, which + # would otherwise stop sharing one read of the channel data. + rows = joined_df.select( cid, ts, te, diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py index 85875ed0..cbac91d7 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -1,6 +1,8 @@ # pylint: disable=missing-function-docstring """Parity of the Spark ``resample`` pushdown with ``SampleSeries.resample``.""" +import re + import numpy as np import numpy.testing as nptest import pyspark.sql.types as T @@ -22,7 +24,7 @@ def _pushdown(rate): return _RULE.match(selector.resample(rate)) -def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType()): +def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType(), chunk_rows=None): """Run the rule on ``{(cid, ch, variant): (tstarts, tends, values)}``. Returns ``{(cid, ch, variant): SampleSeries}`` built like the pandas cache does. @@ -44,7 +46,8 @@ def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType()): for ts, te, v in zip(tstarts, tends, values, strict=True) ] df = spark.createDataFrame(rows, schema=schema) - result = _RULE.to_spark(df, {vid: _pushdown(rate) for vid, rate in rates.items()}, _CONFIG) + rule = _RULE if chunk_rows is None else ResamplePushdownRule(chunk_rows=chunk_rows) + result = rule.to_spark(df, {vid: _pushdown(rate) for vid, rate in rates.items()}, _CONFIG) assert result.schema == df.schema pdf = result.toPandas() @@ -112,7 +115,10 @@ def test_resample_matches_sample_series(spark, name): _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) -def test_all_cases_in_one_pass_and_other_variants_unchanged(spark): +# Few samples per chunk put chunk boundaries at gaps, inside buckets and next to +# zero-duration samples, with empty chunks in between. +@pytest.mark.parametrize("chunk_rows", [None, 1, 2, 3]) +def test_all_cases_in_one_pass_and_other_variants_unchanged(spark, chunk_rows): series = {} expected = {} for i, (tstarts, tends, values, _rate) in enumerate(_CASES.values()): @@ -126,13 +132,14 @@ def test_all_cases_in_one_pass_and_other_variants_unchanged(spark): expected[(i, 7, vid)] = SampleSeries( tstarts, tends, [np.nan if v is None else v for v in values] ) - actual = _resample_in_spark(spark, series, {1: 1, 2: 2}) + actual = _resample_in_spark(spark, series, {1: 1, 2: 2}, chunk_rows=chunk_rows) assert actual.keys() == expected.keys() for key, exp in expected.items(): _assert_same(actual[key], exp) -def test_ns_epoch_longs_above_2_pow_53(spark): +@pytest.mark.parametrize("chunk_rows", [None, 7]) +def test_ns_epoch_longs_above_2_pow_53(spark, chunk_rows): rng = np.random.default_rng(7) base = 1_700_000_000_123_456_789 steps = rng.integers(9_000_000, 11_000_000, size=500) # ~10 ms, jittered @@ -140,7 +147,9 @@ def test_ns_epoch_longs_above_2_pow_53(spark): tstarts, tends = times[:-1].tolist(), times[1:].tolist() values = rng.normal(size=500).tolist() rate = 100_000_000 # 100 ms - actual = _resample_in_spark(spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}) + actual = _resample_in_spark( + spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}, chunk_rows=chunk_rows + ) expected = _expected((tstarts, tends, values), rate) assert len(expected) < len(tstarts) / 5 _assert_same(actual[(1, 1, 1)], expected) @@ -157,7 +166,82 @@ def test_only_linear_window_frames(spark): assert "specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())" in plan -def test_ns_epoch_grid_on_sample_starts(spark): +def _random_series(rng, n): + """Integral samples with gaps, zero-duration samples and NaN values, in tstart order.""" + tstarts, tends, values = [], [], [] + t, after_point = 0, False + for _ in range(n): + if rng.random() < 0.15: + t += int(rng.integers(1, 6)) + duration = 0 if not after_point and rng.random() < 0.1 else int(rng.integers(1, 5)) + tstarts.append(t) + tends.append(t + duration) + values.append(np.nan if rng.random() < 0.1 else float(rng.integers(0, 100))) + t, after_point = t + duration, duration == 0 + return tstarts, tends, values + + +def test_random_series_across_chunk_sizes(spark): + rng = np.random.default_rng(3) + rates = {1: 1, 2: 3, 3: 7} + series, expected = {}, {} + for cid, n in enumerate([1, 2, 40, 60, 60, 80]): + tstarts, tends, values = _random_series(rng, n) + for vid, rate in rates.items(): + series[(cid, 1, vid)] = (tstarts, tends, values) + expected[(cid, 1, vid)] = _expected((tstarts, tends, values), rate) + series[(cid, 1, 0)] = (tstarts, tends, values) + expected[(cid, 1, 0)] = SampleSeries(tstarts, tends, values) + assert sum(len(e) for e in expected.values()) > 0 + for chunk_rows in (1, 2, 3, 5, 8, 13): + actual = _resample_in_spark(spark, series, rates, chunk_rows=chunk_rows) + assert actual.keys() == expected.keys() + for key, exp in expected.items(): + _assert_same(actual[key], exp) + + +def _top_level_items(text): + items, depth, start = [], 0, 0 + for i, char in enumerate(text): + depth += {"(": 1, ")": -1}.get(char, 0) + if char == "," and depth == 0: + items.append(text[start:i].strip()) + start = i + 1 + return [*items, text[start:].strip()] + + +def test_windows_over_samples_are_chunked(spark): + # Spark buffers a whole window partition, so no window over the samples may be + # partitioned by the series alone; only the chunk summaries are. + schema = "container_id long, channel_id long, {} int, tstart long, tend long, value double" + df = spark.createDataFrame([(1, 1, 1, 0, 1, 1.0)], schema=schema.format(VARIANT_COL)) + result = _RULE.to_spark(df, {1: _pushdown(1)}, _CONFIG) + plan = result._jdf.queryExecution().optimizedPlan().toString() + specs = re.findall(r"windowspecdefinition\((.*?), specifiedwindowframe", plan) + series_key = ["container_id", "channel_id", VARIANT_COL] + seen = set() + for spec in specs: + items = _top_level_items(re.sub(r"#\d+L?", "", spec)) + partition = [item for item in items if not item.endswith(("NULLS FIRST", "NULLS LAST"))] + order = [item.split(" ")[0] for item in items[len(partition) :]] + seen.add((tuple(partition), tuple(order))) + if partition == series_key: + assert order == ["_rs_chunk"] + else: + assert partition[:3] == series_key + assert {"_rs_chunk", "_rs_block"} & set(partition) + assert any(order[-2:] == ("tstart", "tend") for _, order in seen) + assert any(partition == tuple(series_key) for partition, _ in seen) + + +@pytest.mark.parametrize("chunk_rows", [0, -1, 1.5, True]) +def test_chunk_rows_must_be_a_positive_integer(chunk_rows): + with pytest.raises((TypeError, ValueError), match="chunk_rows"): + ResamplePushdownRule(chunk_rows=chunk_rows) + + +@pytest.mark.parametrize("chunk_rows", [None, 7]) +def test_ns_epoch_grid_on_sample_starts(spark, chunk_rows): # Every fifth sample starts exactly on a grid point, the case where the bucket # index estimate needs its correction step. base = 1_700_000_000_123_456_789 @@ -165,29 +249,42 @@ def test_ns_epoch_grid_on_sample_starts(spark): tstarts, tends = times[:-1], times[1:] values = [float(i) for i in range(300)] rate = 50_000_000 - actual = _resample_in_spark(spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}) + actual = _resample_in_spark( + spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}, chunk_rows=chunk_rows + ) _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) -def test_integer_time_columns(spark): +@pytest.mark.parametrize("chunk_rows", [None, 2]) +def test_integer_time_columns(spark, chunk_rows): tstarts, tends, values, rate = _CASES["where_gap_rate_2"] actual = _resample_in_spark( - spark, {(1, 1, 1): (tstarts, tends, values)}, {1: rate}, time_type=T.IntegerType() + spark, + {(1, 1, 1): (tstarts, tends, values)}, + {1: rate}, + time_type=T.IntegerType(), + chunk_rows=chunk_rows, ) _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) -def test_seconds_as_double_with_binary_exact_rate(spark): +@pytest.mark.parametrize("chunk_rows", [None, 2]) +def test_seconds_as_double_with_binary_exact_rate(spark, chunk_rows): tstarts = [10.0, 10.25, 11.0, 12.5, 14.0] tends = [10.25, 11.0, 12.5, 14.0, 14.0] values = [1.0, 2.0, 3.0, 4.0, 5.0] actual = _resample_in_spark( - spark, {(1, 1, 1): (tstarts, tends, values)}, {1: 0.75}, time_type=T.DoubleType() + spark, + {(1, 1, 1): (tstarts, tends, values)}, + {1: 0.75}, + time_type=T.DoubleType(), + chunk_rows=chunk_rows, ) _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), 0.75)) -def test_seconds_as_double_with_fractional_rate(spark): +@pytest.mark.parametrize("chunk_rows", [None, 7]) +def test_seconds_as_double_with_fractional_rate(spark, chunk_rows): # 0.1 is not binary-exact: grid points may differ from numpy by one ulp where # numpy fuses the multiply-add. Sample boundaries stay clear of the grid, so the # values must match exactly. @@ -196,7 +293,11 @@ def test_seconds_as_double_with_fractional_rate(spark): tstarts, tends = times[:-1].tolist(), times[1:].tolist() values = rng.normal(size=400).tolist() actual = _resample_in_spark( - spark, {(1, 1, 1): (tstarts, tends, values)}, {1: 0.1}, time_type=T.DoubleType() + spark, + {(1, 1, 1): (tstarts, tends, values)}, + {1: 0.1}, + time_type=T.DoubleType(), + chunk_rows=chunk_rows, )[(1, 1, 1)] expected = _expected((tstarts, tends, values), 0.1) assert len(actual) == len(expected) diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py index a57c172b..2a4fb035 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py @@ -20,6 +20,11 @@ ) from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig from tests.conftest import spark +from tests.impulse_query_engine.unit.analyze.query.solvers.default_solver_pushdown_test import ( + CHANNEL_DATA_SCANS, + channel_data_scans, + small_chunks, +) RATE = 1e7 # 10 s in µs; the fixture channels are sampled roughly every 0.9 s _KVS_SCHEMA = "spark_catalog.silver_key_value_store" @@ -95,6 +100,7 @@ def _solve_both(spark, db, make_channels, udf_calls, **solver_kwargs) -> pd.Data return _assert_same(native, udf), native +@pytest.mark.usefixtures("small_chunks") class TestNativeCalculatedChannels: def test_resampled_channels_skip_the_udf(self, spark, key_value_store_db, udf_calls): def channels(q): @@ -116,6 +122,7 @@ def channels(q): assert rows["channel_id"].nunique() == 2 plan = native._jdf.queryExecution().executedPlan().toString() assert "FlatMapGroupsInPandas" not in plan + assert channel_data_scans(native) == CHANNEL_DATA_SCANS def test_same_expression_under_two_identities_and_two_rates( self, spark, key_value_store_db, udf_calls diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py index 2ce8be9a..bc5c334e 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py @@ -1,6 +1,7 @@ # pylint: disable=missing-function-docstring """DefaultSolver with ``enable_pushdown``: identical results, fewer rows into the UDF.""" +import re from collections import Counter import numpy as np @@ -18,6 +19,7 @@ PushdownPlan, QueryOptimizer, ResamplePushdownRule, + query_optimizer, ) from impulse_query_engine.analyze.query.solvers.default_solver import ( DefaultSolver, @@ -53,6 +55,30 @@ def _unit_conversion_config() -> SolverConfig: ) +@pytest.fixture +def small_chunks(monkeypatch): + """Resample in chunks of a few samples, so every channel spans many chunks.""" + monkeypatch.setattr( + query_optimizer, "default_rules", lambda: (ResamplePushdownRule(chunk_rows=7),) + ) + + +def channel_data_scans(df) -> list[str]: + """Columns of each channels-table scan in the final plan of the executed *df*.""" + final = df._jdf.queryExecution().executedPlan().toString().split("== Initial Plan ==")[0] + return sorted( + re.sub(r"#\d+L?", "", columns) + for columns in re.findall(r"FileScan parquet \S+\.channels\[([^\]]*)\]", final) + ) + + +#: One full read for the rows, one narrow read shared by the resample summaries. +CHANNEL_DATA_SCANS = [ + "container_id,channel_id,tstart,tend", + "container_id,channel_id,tstart,tend,value", +] + + @pytest.fixture def udf_rows(monkeypatch): """Rows reaching the grouped-map UDF, per solve: ``{(channel_id, variant): count}``.""" @@ -102,6 +128,7 @@ def _derived_rows(pdf: pd.DataFrame, col: str) -> int: return sum(len(s) for s in pdf[col]) +@pytest.mark.usefixtures("small_chunks") class TestIdenticalResults: def test_resample_only(self, spark: SparkSession, key_value_store_db, udf_rows): off, on = _solve_both( @@ -117,6 +144,18 @@ def test_resample_only(self, spark: SparkSession, key_value_store_db, udf_rows): assert set(rows_on) == {(5, 1)} assert rows_on[(5, 1)] == _derived_rows(off, "rpm") < rows_off[(5, None)] / 5 + def test_reads_the_channel_data_once_plus_one_narrow_read( + self, spark: SparkSession, key_value_store_db + ): + q = key_value_store_db.query + result = q.select( + q.channel(channel_name="Engine RPM").resample(RATE).alias("rpm"), + q.channel(channel_name="Vehicle Speed Sensor").max().alias("speed_max"), + ).solve(spark, solver=DefaultSolver(spark, config=_kvs_config(), enable_pushdown=True)) + rows = result.collect() # executes result's own plan, which the check reads + assert rows and all(row["speed_max"] > 0 for row in rows) + assert channel_data_scans(result) == CHANNEL_DATA_SCANS + def test_raw_use_of_the_same_channel_ships_raw_rows_only( self, spark: SparkSession, key_value_store_db, udf_rows ): From db0b212259726704e7772a9d100c54740877b5cb Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Thu, 15 Oct 2026 09:53:06 +0200 Subject: [PATCH 08/16] feat(query-engine): shared ordered pipeline for pushdown rules Replace per-rule `PushdownRule.to_spark` with a shared `SeriesPipeline` that evaluates all ordered rules in one pass over the channel data. Rules now declare hooks (`row_columns`, `keep`, `kept_columns`, `emit`) operating on a `SeriesContext` of Spark columns, so reads and shuffles are shared across rules. Add `context.py` with `SeriesContext`, `Sample`, `Run`, and `RunBreak`; add `pipeline.py` with chunking, per-chunk windows, chunk summaries, and bounded-memory series processing. Convert `ResamplePushdownRule` to the new hook interface and remove its `chunk_rows` parameter. Update `DefaultSolver` to use `apply_pushdowns`, regenerate API docs, and expand query solver docs to describe the shared pushdown pass. --- .../analyze/query/optimizer/context.md | 265 +++++++++ .../analyze/query/optimizer/pipeline.md | 58 ++ .../analyze/query/optimizer/pushdown.md | 73 ++- .../analyze/query/optimizer/resample_rule.md | 36 +- docs/impulse/docs/references/api/sidebar.json | 2 + .../references/query_engine/query_solvers.md | 35 +- docs/impulse/pydoc-markdown.yml | 2 + .../analyze/query/optimizer/__init__.py | 9 + .../analyze/query/optimizer/context.py | 280 +++++++++ .../analyze/query/optimizer/pipeline.py | 547 ++++++++++++++++++ .../analyze/query/optimizer/pushdown.py | 82 ++- .../analyze/query/optimizer/resample_rule.py | 356 +++--------- .../analyze/query/solvers/default_solver.py | 63 +- .../query/optimizer/pipeline_helpers.py | 183 ++++++ .../analyze/query/optimizer/pipeline_test.py | 235 ++++++++ .../query/optimizer/resample_rule_test.py | 124 ++-- ..._solver_native_calculated_channels_test.py | 37 +- .../solvers/default_solver_pushdown_test.py | 69 ++- 18 files changed, 1993 insertions(+), 463 deletions(-) create mode 100644 docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/context.md create mode 100644 docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pipeline.md create mode 100644 src/impulse_query_engine/analyze/query/optimizer/context.py create mode 100644 src/impulse_query_engine/analyze/query/optimizer/pipeline.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/context.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/context.md new file mode 100644 index 00000000..b127041b --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/context.md @@ -0,0 +1,265 @@ +--- +sidebar_label: context +title: impulse_query_engine.analyze.query.optimizer.context +--- + +What a pushdown rule sees of the samples it computes on. + +A rule never handles DataFrames. The pipeline (:mod:`.pipeline`) evaluates the +rule's hooks with a :class:`SeriesContext`, whose properties are Spark columns +describing the current sample of one of the rule's series: its times and value, +its neighbours, the run (e.g. continuous block) it belongs to, and facts about +the whole series. Columns returned by a hook become part of the shared pass +over the channel data. + + +#### ELEMENT\_TYPE + +Element type of the arrays returned by :meth:`PushdownRule.emit`. + + +#### BLOCK + +Name of the built-in run of continuous samples. + + +#### run\_col + +```python +def run_col(run: str, name: str) -> str +``` + +Name of the pipeline column *name* of the run *run*. + + +## Sample + +```python +class Sample() +``` + +One sample as columns: ``tstart`` and ``tend`` as double, and its value as double. + + +#### RunBreak + +A run definition: whether the sample *cur* starts a new run after *prev*. It + +is called with ``prev`` columns that are null before the series' first +sample; a null result counts as a break. + + +## Run + +```python +class Run() +``` + +Where the current sample sits in its run, a maximal stretch without a break. + +**Arguments**: + +- `starts` (`Column`): The sample starts its run. +- `id` (`Column`): Number of the run within the series, from 1. +- `start` (`Column`): ``tstart`` of the run's first sample. +- `next_starts` (`Column`): The next sample starts a new run; null for the series' last sample. +- `is_last` (`Column`): The run is the last of its series. + +#### block\_break + +```python +def block_break(prev: Sample, cur: Sample) -> Column +``` + +Continuous blocks: a sample continues its block when it starts where the previous ended. + + +## SeriesContext + +```python +class SeriesContext() +``` + +Columns a rule's hooks build on, for the current sample of one of its series. + +Every series is one channel's samples of one variant, ordered by +``(tstart, tend, value)``. Rules running in series order (``ordered``) see +neighbours, runs and series facts; per-sample rules see only the sample. + +**Arguments**: + +- `rule` (`PushdownRule`): The rule the hooks belong to. +- `variants` (`Mapping[int, Pushdown]`): The rule's variants in this solve. +- `key` (`Sequence[str]`): Columns identifying a series. +- `order` (`Sequence[str]`): The channel-data ``tstart``, ``tend`` and value columns: the series order. +- `ordered` (`bool`): Whether the pipeline runs the rule in series order. +- `runs` (`Mapping[str, str]`): The rule's run names to pipeline run names. +- `kept` (`bool`): Whether the hook runs on the kept samples (``kept_columns``, ``emit``). + +#### kept\_stage + +```python +def kept_stage() -> SeriesContext +``` + +This context, for the hooks that run on the kept samples. + + +#### active + +```python +def active() -> Column +``` + +The sample belongs to one of the rule's variants. + + +#### name + +```python +def name(column: str) -> str +``` + +Pipeline name of the rule's column *column*. + + +#### col + +```python +def col(column: str) -> Column +``` + +The rule's column *column*, defined by an earlier hook entry. + + +#### param + +```python +def param(name: str) -> Column +``` + +The pushdown parameter *name* of the sample's variant. + + +#### tstart + +```python +def tstart() -> Column +``` + +The sample's ``tstart`` as double. + + +#### tend + +```python +def tend() -> Column +``` + +The sample's ``tend`` as double. + + +#### value + +```python +def value() -> Column +``` + +The sample's value as double. + + +#### sample + +```python +def sample() -> Sample +``` + +The current sample. + + +#### element + +```python +def element(tstart: Column | None = None, + tend: Column | None = None, + value: Column | None = None) -> Column +``` + +One emitted row; ``tstart`` / ``tend`` left out keep the sample's own times. + + +#### prev + +```python +def prev() -> Sample +``` + +The previous sample of the series; null columns for its first sample. + + +#### next + +```python +def next() -> Sample +``` + +The next sample of the series; null columns for its last sample. + + +#### n + +```python +def n() -> Column +``` + +Number of samples in the series. + + +#### last + +```python +def last() -> Sample +``` + +The series' last sample. + + +#### in\_chunk + +```python +def in_chunk() -> WindowSpec +``` + +Window over the sample's chunk in series order, for local neighbour access. + +Chunks are consecutive slices of a series; use :attr:`prev`, :attr:`next` +and runs for anything that must see across chunk boundaries. + + +#### block + +```python +def block() -> Run +``` + +The continuous block the sample belongs to (gaps end blocks). + + +#### run + +```python +def run(name: str) -> Run +``` + +The sample's run of the rule's run definition *name* (see :meth:`PushdownRule.runs`). + + +#### block\_end + +```python +def block_end() -> Column +``` + +``nanmax`` of the ``tend`` values of the sample's block (kept samples only). + + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pipeline.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pipeline.md new file mode 100644 index 00000000..b19e517d --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pipeline.md @@ -0,0 +1,58 @@ +--- +sidebar_label: pipeline +title: impulse_query_engine.analyze.query.optimizer.pipeline +--- + +Spark execution of a pushdown plan: one shared pass for all rules. + + +#### CHUNK\_ROWS + +Average number of samples per chunk of a series. + + +#### apply\_pushdowns + +```python +def apply_pushdowns(rows: DataFrame, + plan: PushdownPlan, + config, + include_raw: bool = True) -> DataFrame +``` + +The rows a solve ships: raw rows (when *include_raw*) and all derived rows. + +**Arguments**: + +- `rows` (`pyspark.sql.DataFrame`): Channel data ``(container_id, channel_id, tstart, tend, value)`` tagged +with :data:`VARIANT_COL`, one copy per variant a channel ships. +- `plan` (`PushdownPlan`): The solve's pushdown plan. +- `config` (`SolverConfig`): Provides the internal column names. +- `include_raw` (`bool`): Whether to keep the raw rows (variant 0). + +**Returns**: + +`pyspark.sql.DataFrame`: Rows with the columns and types of *rows*. + +## SeriesPipeline + +```python +class SeriesPipeline() +``` + +The ordered pass shared by all rules that need their series in order. + +**Arguments**: + +- `rules` (`Sequence of (PushdownRule, Mapping[int, Pushdown])`): Ordered rules with their variants in this solve. +- `config` (`SolverConfig`): Provides the internal column names. + +#### run + +```python +def run(rows: DataFrame) -> DataFrame +``` + +Replace the samples of the rules' variants with their derived rows. + + diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md index fec552b9..3cef9462 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -81,11 +81,7 @@ class PushdownRule(abc.ABC) Recognizes a subtree and computes it in Spark ahead of the pandas UDF. -A rule's derived rows use the channel-data shape -``(container_id, channel_id, variant, tstart, tend, value)``; :meth:`decode` -turns one container/channel/variant slice of them back into the object the -original subtree evaluates to. The result must equal evaluating the subtree -in pandas on the same channel data. +Planning and decoding run in Python: :meth:`match` recognizes the subtree, #### name @@ -93,6 +89,13 @@ in pandas on the same channel data. Unique rule name, part of every :attr:`Pushdown.variant_key`. +#### ordered + +Whether the rule needs its series in order (neighbours, runs, series + +facts). Per-sample rules set ``False`` and skip the series scan. + + #### match ```python @@ -123,35 +126,57 @@ def result_type(pushdown: Pushdown) -> type Core-model class :meth:`decode` returns for *pushdown*. -#### to\_spark +#### runs ```python -def to_spark(rows: DataFrame, variants: Mapping[int, Pushdown], - config: SolverConfig) -> DataFrame +def runs() -> Mapping[str, RunBreak] ``` -Replace the rows of this rule's *variants* with their derived rows. +Run definitions beyond the built-in continuous blocks, by name. -*rows* is shared by all rules and partitioned by container; the result -must keep that partitioning (windows partitioned by container id, -channel id and :data:`VARIANT_COL` reuse it). Implementations must not -union *rows* or join it to a non-broadcast frame. Small summaries -derived from *rows* (per series or per chunk) may be broadcast-joined -back; Spark reads the channel data again for them, so summaries should -read only the key and time columns to share one narrow read. +Each maps ``(prev, cur)`` samples to whether ``cur`` starts a new run; -**Arguments**: -- `rows` (`pyspark.sql.DataFrame`): All channel-data rows of the solve, tagged with :data:`VARIANT_COL` -(one copy per variant a channel ships) and partitioned by container. -- `variants` (`Mapping[int, Pushdown]`): This rule's variant ids, each with a representative pushdown carrying -its parameters. -- `config` (`SolverConfig`): Provides the internal column names. +#### row\_columns + +```python +def row_columns(ctx: SeriesContext) -> Mapping[str, Column] +``` + +Columns computed for every sample, in order; refer to them via ``ctx.col``. + + +#### keep + +```python +def keep(ctx: SeriesContext) -> Column +``` + +Whether a sample can emit rows; the others are dropped early. + + +#### kept\_columns + +```python +def kept_columns(ctx: SeriesContext) -> Mapping[str, Column] +``` + +Columns computed for the kept samples, in order (``ctx.block_end`` available). + + +#### emit + +```python +def emit(ctx: SeriesContext) -> Column +``` + +The derived rows of a kept sample. **Returns**: -`pyspark.sql.DataFrame`: *rows* with the same columns and types, the rows of *variants* -replaced by derived rows and all other rows unchanged. +`Column`: An array of :data:`~.context.ELEMENT_TYPE` structs built with +:meth:`SeriesContext.element`; empty or null emits nothing. Times +left out keep the sample's own ``tstart`` / ``tend``. #### decode diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md index e5720897..af0cdf38 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md @@ -22,22 +22,40 @@ sample spacing yields more rows than the raw channel. Each sample is assigned the bucket ``c`` of the first grid point at or after its ``tstart``. The last sample of a bucket holds the value of grid points -``c, c + 1, …`` up to the next non-empty bucket, so only those samples emit rows +``c, c + 1, …`` up to the next sample's bucket, so only those samples are kept and every row carries just the grid points it covers. Values are looked up within a block; this matches the pandas evaluation unless intervals overlap. -Every series is cut into equal time slices (chunks) holding about *chunk_rows* -samples on average, and the windows over samples run per chunk, so a worker -holds one chunk of a series at a time instead of the whole series. Gaps at -chunk boundaries, the block a chunk starts in and that block's start are -resolved from one summary row per chunk. - Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. numpy builds that fuse this into a multiply-add can differ by one ulp for fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral grids are exact. -**Arguments**: -- `chunk_rows` (`int`): Average number of samples per chunk (default 100 000). +## \_Grid + +```python +class _Grid() +``` + +The numpy.arange grid of the sample's block. + +``numpy.arange(t_min, t_max + rate, rate)`` has length +``ceil(((t_max + rate) - t_min) / rate)`` and is filled as ``t_min``, +``t_min + rate``, then ``t_min + k * ((t_min + rate) - t_min)``. + + +#### first\_index + +```python +def first_index(name: str, x: Column, strict: Column, + active: Column) -> dict[str, Column] +``` + +Columns ending in *name*: the smallest ``k >= 0`` with ``point(k) >= x`` (``> x`` if *strict*). + +Estimated from ``(x - t_min) / delta`` and corrected by one step against the +exact grid points; a result that still is not exact raises instead of +misplacing a sample. Only evaluated where *active*. + diff --git a/docs/impulse/docs/references/api/sidebar.json b/docs/impulse/docs/references/api/sidebar.json index 90c19dfd..92c8d3f9 100644 --- a/docs/impulse/docs/references/api/sidebar.json +++ b/docs/impulse/docs/references/api/sidebar.json @@ -32,6 +32,8 @@ }, { "items": [ + "references/api/impulse_query_engine/analyze/query/optimizer/context", + "references/api/impulse_query_engine/analyze/query/optimizer/pipeline", "references/api/impulse_query_engine/analyze/query/optimizer/plan", "references/api/impulse_query_engine/analyze/query/optimizer/pushdown", "references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer", diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index 44f818fd..099b3ee3 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -206,18 +206,24 @@ that share the same `tstart` is unspecified, as in the pandas evaluation, and ch `[tstart, tend)` intervals overlap are not supported (the RLE and interval encoders never produce overlaps). -Spark cuts each channel into chunks: equal time slices that hold about 100 000 samples on -average. It processes one chunk at a time, sorted by `tstart`: every sample is assigned the grid -bucket it falls into, and only the last sample of each bucket emits rows (the grid points up to -the next non-empty bucket). What crosses a chunk boundary comes from one summary row per chunk: -whether a gap starts there, and which block (and block start) the chunk continues. A task -therefore holds one chunk of a channel, plus the resampled rows of one continuous block, rather -than the whole channel. If the sampling density varies strongly within a channel, a chunk can hold -more samples. - -The summaries need the channel's time columns a second time, so with pushdown Spark reads the -channel data twice: once in full, and once with only `container_id`, `channel_id`, `tstart` and -`tend`. The full rows are still shuffled only once. +#### How Spark runs pushed-down subtrees + +All pushed-down subtrees of a solve share one pass over the channel data: + +- Spark cuts each pushed channel into chunks, equal time slices of about 100 000 samples on + average, and processes the chunks in parallel. A task holds one chunk at a time, not a whole + channel or container; with strongly varying sampling density a chunk can hold more samples. +- Within a chunk the samples are ordered by `tstart`, then `tend`, then value; samples with equal + times are ordered by value. For `resample`, each sample is assigned the grid bucket it falls + into, and only the last sample of each bucket emits rows: the grid points up to the next + sample's bucket. +- What crosses a chunk boundary (neighbouring samples, gaps, the start and end of a continuous + block) comes from one summary row per chunk. + +The channel data is read **twice**, however many expressions are pushed down: once in full, and +once with only `container_id`, `channel_id`, `tstart` and `tend` (the channel boundaries that +define the chunks). Channels that are also used outside pushed-down subtrees, e.g. `rpm.max()` +next to `spd.resample(1e6)`, are read once more for their raw rows. ### Calculated channels without the pandas UDF @@ -239,4 +245,7 @@ Auto-generated symbol-level docs: — per-table column mappings, filters, and project scoping. - [`QueryOptimizer`](../api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md) and [`PushdownRule`](../api/impulse_query_engine/analyze/query/optimizer/pushdown.md) - — the Spark pushdown optimizer and the interface for its rules. + — the Spark pushdown optimizer and the interface for its rules; + [`pipeline`](../api/impulse_query_engine/analyze/query/optimizer/pipeline.md) and + [`SeriesContext`](../api/impulse_query_engine/analyze/query/optimizer/context.md) — how rules + run in Spark and what they can use. diff --git a/docs/impulse/pydoc-markdown.yml b/docs/impulse/pydoc-markdown.yml index 515a8dfe..7fa9bb8d 100644 --- a/docs/impulse/pydoc-markdown.yml +++ b/docs/impulse/pydoc-markdown.yml @@ -25,6 +25,8 @@ loaders: - impulse_query_engine.analyze.query.solvers.solver_config - impulse_query_engine.analyze.query.solvers.query_solver - impulse_query_engine.analyze.query.optimizer.pushdown + - impulse_query_engine.analyze.query.optimizer.context + - impulse_query_engine.analyze.query.optimizer.pipeline - impulse_query_engine.analyze.query.optimizer.plan - impulse_query_engine.analyze.query.optimizer.query_optimizer - impulse_query_engine.analyze.query.optimizer.resample_rule diff --git a/src/impulse_query_engine/analyze/query/optimizer/__init__.py b/src/impulse_query_engine/analyze/query/optimizer/__init__.py index 4dad882a..a3a218a3 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/__init__.py +++ b/src/impulse_query_engine/analyze/query/optimizer/__init__.py @@ -1,9 +1,12 @@ +from .context import ELEMENT_TYPE, Run, RunBreak, Sample, SeriesContext +from .pipeline import SeriesPipeline, apply_pushdowns from .plan import PushdownPlan from .pushdown import RAW_VARIANT, VARIANT_COL, Pushdown, PushdownContext, PushdownRule from .query_optimizer import QueryOptimizer, default_rules from .resample_rule import ResamplePushdownRule __all__ = [ + "ELEMENT_TYPE", "RAW_VARIANT", "VARIANT_COL", "Pushdown", @@ -12,5 +15,11 @@ "PushdownRule", "QueryOptimizer", "ResamplePushdownRule", + "Run", + "RunBreak", + "Sample", + "SeriesContext", + "SeriesPipeline", + "apply_pushdowns", "default_rules", ] diff --git a/src/impulse_query_engine/analyze/query/optimizer/context.py b/src/impulse_query_engine/analyze/query/optimizer/context.py new file mode 100644 index 00000000..f270a3ab --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/context.py @@ -0,0 +1,280 @@ +"""What a pushdown rule sees of the samples it computes on. + +A rule never handles DataFrames. The pipeline (:mod:`.pipeline`) evaluates the +rule's hooks with a :class:`SeriesContext`, whose properties are Spark columns +describing the current sample of one of the rule's series: its times and value, +its neighbours, the run (e.g. continuous block) it belongs to, and facts about +the whole series. Columns returned by a hook become part of the shared pass +over the channel data. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import pyspark.sql.functions as F +import pyspark.sql.types as T +from pyspark.sql import Column, Window, WindowSpec + +from .pushdown import VARIANT_COL + +if TYPE_CHECKING: + from .pushdown import Pushdown, PushdownRule + +#: Element type of the arrays returned by :meth:`PushdownRule.emit`. +ELEMENT_TYPE = T.StructType( + [ + T.StructField("tstart", T.DoubleType()), + T.StructField("tend", T.DoubleType()), + T.StructField("value", T.DoubleType()), + ] +) + +#: Name of the built-in run of continuous samples. +BLOCK = "block" + +# Columns the pipeline provides, shared by all rules. +TS, TE = "_pd_ts", "_pd_te" +CHUNK = "_pd_chunk" +PREV_TS, PREV_TE, PREV_VALUE = "_pd_prev_ts", "_pd_prev_te", "_pd_prev_value" +NEXT_TS, NEXT_TE, NEXT_VALUE = "_pd_next_ts", "_pd_next_te", "_pd_next_value" +COUNT = "_pd_n" +LAST = "_pd_last" +BLOCK_END = "_pd_block_end" + + +def run_col(run: str, name: str) -> str: + """Name of the pipeline column *name* of the run *run*.""" + return f"_pd_run_{run}_{name}" + + +@dataclass(frozen=True) +class Sample: + """One sample as columns: ``tstart`` and ``tend`` as double, and its value as double.""" + + tstart: Column + tend: Column + value: Column + + +#: A run definition: whether the sample *cur* starts a new run after *prev*. It +#: is called with ``prev`` columns that are null before the series' first +#: sample; a null result counts as a break. +RunBreak = Callable[[Sample, Sample], Column] + + +@dataclass(frozen=True) +class Run: + """Where the current sample sits in its run, a maximal stretch without a break. + + Attributes + ---------- + starts : Column + The sample starts its run. + id : Column + Number of the run within the series, from 1. + start : Column + ``tstart`` of the run's first sample. + next_starts : Column + The next sample starts a new run; null for the series' last sample. + is_last : Column + The run is the last of its series. + """ + + starts: Column + id: Column + start: Column + next_starts: Column + is_last: Column + + +def block_break(prev: Sample, cur: Sample) -> Column: + """Continuous blocks: a sample continues its block when it starts where the previous ended.""" + return prev.tend.isNull() | ((cur.tstart - prev.tend) != 0) + + +class SeriesContext: + """Columns a rule's hooks build on, for the current sample of one of its series. + + Every series is one channel's samples of one variant, ordered by + ``(tstart, tend, value)``. Rules running in series order (``ordered``) see + neighbours, runs and series facts; per-sample rules see only the sample. + + Parameters + ---------- + rule : PushdownRule + The rule the hooks belong to. + variants : Mapping[int, Pushdown] + The rule's variants in this solve. + key : Sequence[str] + Columns identifying a series. + order : Sequence[str] + The channel-data ``tstart``, ``tend`` and value columns: the series order. + ordered : bool + Whether the pipeline runs the rule in series order. + runs : Mapping[str, str] + The rule's run names to pipeline run names. + kept : bool + Whether the hook runs on the kept samples (``kept_columns``, ``emit``). + """ + + def __init__( + self, + rule: PushdownRule, + variants: Mapping[int, Pushdown], + key: Sequence[str], + order: Sequence[str], + ordered: bool, + runs: Mapping[str, str], + kept: bool = False, + ): + self._rule = rule + self._variants = variants + self._key = tuple(key) + self._order = tuple(order) + self._ordered = ordered + self._runs = dict(runs) + self._kept = kept + + def kept_stage(self) -> SeriesContext: + """This context, for the hooks that run on the kept samples.""" + return SeriesContext( + self._rule, + self._variants, + self._key, + self._order, + self._ordered, + self._runs, + kept=True, + ) + + # --- always available ------------------------------------------------------- + + @property + def active(self) -> Column: + """The sample belongs to one of the rule's variants.""" + return F.col(VARIANT_COL).isin(sorted(self._variants)) + + def name(self, column: str) -> str: + """Pipeline name of the rule's column *column*.""" + return f"_pd_{self._rule.name}_{column}" + + def col(self, column: str) -> Column: + """The rule's column *column*, defined by an earlier hook entry.""" + return F.col(self.name(column)) + + def param(self, name: str) -> Column: + """The pushdown parameter *name* of the sample's variant.""" + return F.try_element_at( + F.create_map( + *[ + item + for vid, pushdown in sorted(self._variants.items()) + for item in (F.lit(vid), F.lit(dict(pushdown.params)[name])) + ] + ), + F.col(VARIANT_COL), + ) + + @property + def tstart(self) -> Column: + """The sample's ``tstart`` as double.""" + return F.col(TS) + + @property + def tend(self) -> Column: + """The sample's ``tend`` as double.""" + return F.col(TE) + + @property + def value(self) -> Column: + """The sample's value as double.""" + return F.col(self._order[2]).cast("double") + + @property + def sample(self) -> Sample: + """The current sample.""" + return Sample(self.tstart, self.tend, self.value) + + @staticmethod + def element( + tstart: Column | None = None, tend: Column | None = None, value: Column | None = None + ) -> Column: + """One emitted row; ``tstart`` / ``tend`` left out keep the sample's own times.""" + + def field(col: Column | None, name: str) -> Column: + return (F.lit(None) if col is None else col).cast("double").alias(name) + + return F.struct(field(tstart, "tstart"), field(tend, "tend"), field(value, "value")) + + # --- rules running in series order -------------------------------------------- + + def _require_ordered(self, what: str) -> None: + if not self._ordered: + raise TypeError(f"{self._rule.name}: {what} needs an ordered rule") + + @property + def prev(self) -> Sample: + """The previous sample of the series; null columns for its first sample.""" + self._require_ordered("prev") + return Sample(F.col(PREV_TS), F.col(PREV_TE), F.col(PREV_VALUE)) + + @property + def next(self) -> Sample: + """The next sample of the series; null columns for its last sample.""" + self._require_ordered("next") + return Sample(F.col(NEXT_TS), F.col(NEXT_TE), F.col(NEXT_VALUE)) + + @property + def n(self) -> Column: + """Number of samples in the series.""" + self._require_ordered("n") + return F.col(COUNT) + + @property + def last(self) -> Sample: + """The series' last sample.""" + self._require_ordered("last") + last = F.col(LAST) + return Sample(last.getField("tstart"), last.getField("tend"), last.getField("value")) + + @property + def in_chunk(self) -> WindowSpec: + """Window over the sample's chunk in series order, for local neighbour access. + + Chunks are consecutive slices of a series; use :attr:`prev`, :attr:`next` + and runs for anything that must see across chunk boundaries. + """ + self._require_ordered("in_chunk") + return Window.partitionBy(*self._key, CHUNK).orderBy(*self._order) + + @property + def block(self) -> Run: + """The continuous block the sample belongs to (gaps end blocks).""" + self._require_ordered("block") + return _run(BLOCK) + + def run(self, name: str) -> Run: + """The sample's run of the rule's run definition *name* (see :meth:`PushdownRule.runs`).""" + self._require_ordered("run") + return _run(self._runs[name]) + + @property + def block_end(self) -> Column: + """``nanmax`` of the ``tend`` values of the sample's block (kept samples only).""" + self._require_ordered("block_end") + if not self._kept: + raise TypeError(f"{self._rule.name}: block_end is available to kept samples only") + return F.col(BLOCK_END) + + +def _run(run: str) -> Run: + return Run( + starts=F.col(run_col(run, "starts")), + id=F.col(run_col(run, "id")), + start=F.col(run_col(run, "start")), + next_starts=F.col(run_col(run, "next_starts")), + is_last=F.col(run_col(run, "is_last")), + ) diff --git a/src/impulse_query_engine/analyze/query/optimizer/pipeline.py b/src/impulse_query_engine/analyze/query/optimizer/pipeline.py new file mode 100644 index 00000000..ab94b687 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/pipeline.py @@ -0,0 +1,547 @@ +"""Spark execution of a pushdown plan: one shared pass for all rules. + +:func:`apply_pushdowns` turns the tagged channel data of a solve into the rows +the solver ships: raw rows where still needed, plus every rule's derived rows. +Rules contribute columns only (see :class:`~.pushdown.PushdownRule`), so the +number of reads of the channel data does not grow with the number of rules: + +* **Ordered rules** (:class:`SeriesPipeline`) read the channel data once in full + and once narrowly (series boundaries, aggregated before their shuffle). The + samples are shuffled once, by series and chunk, and processed per chunk in + parallel. +* **Raw rows and per-sample rules** share one more read, only when present. + +The ordered pass, per series (one channel's samples of one variant, ordered by +``(tstart, tend, value)``): + +1. Series boundaries give each sample its chunk: equal time slices holding about + :data:`CHUNK_ROWS` samples on average. +2. A window per chunk derives each sample's neighbours and run breaks. +3. One summary row per chunk, scanned in chunk order, carries what crosses chunk + boundaries: neighbours at chunk edges, run numbers and starts, the series' + sample count and last sample, and block ends. +4. Rules add per-sample columns (``row_columns``) and select samples (``keep``). +5. Kept samples get their block end; rules add ``kept_columns``. +6. Rules ``emit`` the derived rows. + +Windows never span more than a chunk of samples (or the chunk summaries of a +series), so the memory per task is bounded by the chunk size. + +Keeping one full read for the chunk pass and its summaries relies on both +seeing the same projected rows: a fixed projection before the shuffle, a chunk +id that cannot be null, the value as the last sort key, and summaries that +depend on the chunk window. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import reduce + +import pyspark.sql.functions as F +from pyspark.sql import Column, DataFrame, Window + +from .context import ( + BLOCK, + BLOCK_END, + CHUNK, + COUNT, + LAST, + NEXT_TE, + NEXT_TS, + NEXT_VALUE, + PREV_TE, + PREV_TS, + PREV_VALUE, + TE, + TS, + RunBreak, + Sample, + SeriesContext, + block_break, + run_col, +) +from .plan import PushdownPlan +from .pushdown import RAW_VARIANT, VARIANT_COL, Pushdown, PushdownRule + +#: Average number of samples per chunk of a series. +CHUNK_ROWS = 100_000 + +_ELEMENT = "_pd_element" + + +def apply_pushdowns( + rows: DataFrame, plan: PushdownPlan, config, include_raw: bool = True +) -> DataFrame: + """The rows a solve ships: raw rows (when *include_raw*) and all derived rows. + + Parameters + ---------- + rows : pyspark.sql.DataFrame + Channel data ``(container_id, channel_id, tstart, tend, value)`` tagged + with :data:`VARIANT_COL`, one copy per variant a channel ships. + plan : PushdownPlan + The solve's pushdown plan. + config : SolverConfig + Provides the internal column names. + include_raw : bool + Whether to keep the raw rows (variant 0). + + Returns + ------- + pyspark.sql.DataFrame + Rows with the columns and types of *rows*. + """ + groups = plan.variants_by_rule() + ordered = [(rule, variants) for rule, variants in groups if rule.ordered] + per_sample = [(rule, variants) for rule, variants in groups if not rule.ordered] + variant = F.col(VARIANT_COL) + parts = [] + row_variants = ([RAW_VARIANT] if include_raw else []) + _variant_ids(per_sample) + if row_variants: + parts.append(_row_branch(rows.where(variant.isin(row_variants)), per_sample, config)) + if ordered: + pipeline = SeriesPipeline(ordered, config) + parts.append(pipeline.run(rows.where(variant.isin(_variant_ids(ordered))))) + if not parts: + return rows.limit(0) + return reduce(lambda a, b: a.unionByName(b), parts) + + +def _variant_ids(groups: Sequence[tuple[PushdownRule, Mapping[int, Pushdown]]]) -> list[int]: + return sorted(vid for _, variants in groups for vid in variants) + + +@dataclass(frozen=True) +class _Run: + """A run definition of the pipeline, active for *variants* (all when ``None``).""" + + name: str + breaks: RunBreak + variants: tuple[int, ...] | None + + +class SeriesPipeline: + """The ordered pass shared by all rules that need their series in order. + + Parameters + ---------- + rules : Sequence of (PushdownRule, Mapping[int, Pushdown]) + Ordered rules with their variants in this solve. + config : SolverConfig + Provides the internal column names. + """ + + def __init__(self, rules: Sequence[tuple[PushdownRule, Mapping[int, Pushdown]]], config): + self.rules = list(rules) + self.config = config + self.key = [config.container_id_col, config.channel_id_col, VARIANT_COL] + self.order = [config.tstart_col, config.tend_col, config.value_col] + self.runs = [_Run(BLOCK, block_break, None)] + self.rule_runs: dict[str, dict[str, str]] = {} + for rule, variants in self.rules: + names = {} + for name, breaks in rule.runs().items(): + names[name] = f"{rule.name}_{name}" + self.runs.append(_Run(names[name], breaks, tuple(sorted(variants)))) + self.rule_runs[rule.name] = names + + def _context(self, rule: PushdownRule, variants: Mapping[int, Pushdown]) -> SeriesContext: + return SeriesContext(rule, variants, self.key, self.order, True, self.rule_runs[rule.name]) + + def run(self, rows: DataFrame) -> DataFrame: + """Replace the samples of the rules' variants with their derived rows.""" + df = self._chunked(rows) + df = self._chunk_scan(df) + df = df.join(F.broadcast(self._chunk_summaries(df)), on=[*self.key, CHUNK]) + df = self._resolve_neighbours_and_runs(df) + contexts = [(rule, self._context(rule, variants)) for rule, variants in self.rules] + for rule, ctx in contexts: + df = _with_rule_columns(df, ctx, rule.row_columns(ctx)) + df = _keep(df, contexts, extra=_flag(run_col(BLOCK, "ends"))) + df = self._with_block_end(df) + kept = [(rule, ctx.kept_stage()) for rule, ctx in contexts] + for rule, ctx in kept: + df = _with_rule_columns(df, ctx, rule.kept_columns(ctx)) + return _emit(df, kept, rows, self.order, raw=None) + + # --- stages ---------------------------------------------------------------------- + + def _chunked(self, rows: DataFrame) -> DataFrame: + """Series boundaries and each sample's chunk, projected for the shuffle.""" + ts, te, value = self.order + ts_d, te_d = F.col(ts).cast("double"), F.col(te).cast("double") + series = rows.groupBy(*self.key).agg( + F.min(ts_d).alias("_pd_first"), + F.max(_nan_free(te_d)).alias("_pd_end"), + F.count(F.lit(1)).alias("_pd_count"), + ) + first, end = F.col("_pd_first"), F.col("_pd_end") + n_chunks = F.ceil(F.col("_pd_count") / F.lit(CHUNK_ROWS)) + chunk = F.coalesce( + F.when( + end > first, + F.least(F.floor((ts_d - first) / (end - first) * n_chunks), n_chunks - 1), + ), + F.lit(0).cast("long"), + ) + return rows.join(F.broadcast(series), on=self.key).select( + *self.key, ts, te, value, ts_d.alias(TS), te_d.alias(TE), chunk.alias(CHUNK) + ) + + def _chunk_scan(self, df: DataFrame) -> DataFrame: + """Neighbours within the chunk, and each run's breaks and segments.""" + value = self.order[2] + in_chunk = Window.partitionBy(*self.key, CHUNK).orderBy(*self.order) + running = in_chunk.rowsBetween(Window.unboundedPreceding, Window.currentRow) + df = ( + df.withColumn("_pd_first_in_chunk", F.row_number().over(in_chunk) == 1) + .withColumn("_pd_last_in_chunk", F.lead(F.lit(True)).over(in_chunk).isNull()) + .withColumn("_pd_lag_ts", F.lag(TS).over(in_chunk)) + .withColumn("_pd_lag_te", F.lag(TE).over(in_chunk)) + .withColumn("_pd_lag_value", F.lag(value).over(in_chunk)) + .withColumn("_pd_lead_ts", F.lead(TS).over(in_chunk)) + .withColumn("_pd_lead_te", F.lead(TE).over(in_chunk)) + .withColumn("_pd_lead_value", F.lead(value).over(in_chunk)) + ) + prev = Sample( + F.col("_pd_lag_ts"), F.col("_pd_lag_te"), F.col("_pd_lag_value").cast("double") + ) + cur = Sample(F.col(TS), F.col(TE), F.col(value).cast("double")) + for run in self.runs: + brk = F.col(run_col(run.name, "brk")) + df = ( + df.withColumn( + run_col(run.name, "brk"), + _masked( + run, + F.when(F.col("_pd_first_in_chunk"), F.lit(0).cast("long")).otherwise( + _breaks(run, prev, cur) + ), + ), + ) + .withColumn(run_col(run.name, "seg"), F.sum(brk).over(running)) + .withColumn( + run_col(run.name, "seg_start"), + F.last(F.when(brk == 1, F.col(TS)), ignorenulls=True).over(running), + ) + .withColumn(run_col(run.name, "lead_brk"), F.lead(brk).over(in_chunk)) + ) + return df + + def _chunk_summaries(self, df: DataFrame) -> DataFrame: + """One row per chunk with everything that crosses chunk boundaries.""" + order = F.struct(*self.order) + sample = F.struct( + F.col(TS).alias("tstart"), + F.col(TE).alias("tend"), + F.col(self.order[2]).cast("double").alias("value"), + ) + block_brk, block_seg = F.col(run_col(BLOCK, "brk")), F.col(run_col(BLOCK, "seg")) + # Head and tail are taken via the chunk window's flags, so the summaries always + # run on the chunk-partitioned samples and share their read. + head = F.when(F.col("_pd_first_in_chunk"), F.lit(1)) + tail = F.when(F.col("_pd_last_in_chunk"), F.lit(1)) + aggregates = [ + F.max_by(sample, head).alias("_pd_head"), + F.max_by(sample, tail).alias("_pd_tail"), + F.max_by(block_brk, tail).alias("_pd_tail_brk"), + F.count(F.lit(1)).alias("_pd_chunk_count"), + # The block of the chunk's first sample: its last sample within the chunk. + F.max_by(sample, F.when(block_seg == 0, order)).alias("_pd_seg0_last"), + F.sum(F.when(block_seg == 0, 1).otherwise(0)).alias("_pd_seg0_count"), + ] + for run in self.runs: + brk = F.col(run_col(run.name, "brk")) + aggregates += [ + F.sum(brk).alias(run_col(run.name, "inner")), + F.max(F.when(brk == 1, F.col(TS))).alias(run_col(run.name, "last_start")), + ] + chunks = df.groupBy(*self.key, CHUNK).agg(*aggregates) + + by_chunk = Window.partitionBy(*self.key).orderBy(CHUNK) + before = by_chunk.rowsBetween(Window.unboundedPreceding, -1) + so_far = by_chunk.rowsBetween(Window.unboundedPreceding, Window.currentRow) + everything = by_chunk.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) + later = ( + Window.partitionBy(*self.key) + .orderBy(F.col(CHUNK).desc()) + .rowsBetween(Window.unboundedPreceding, -1) + ) + head, tail = F.col("_pd_head"), F.col("_pd_tail") + for field in ("tstart", "tend", "value"): + chunks = chunks.withColumn( + f"_pd_prev_tail_{field}", F.lag(tail.getField(field)).over(by_chunk) + ).withColumn(f"_pd_next_head_{field}", F.lead(head.getField(field)).over(by_chunk)) + prev_tail = Sample(*(F.col(f"_pd_prev_tail_{f}") for f in ("tstart", "tend", "value"))) + head_sample = Sample(*(head.getField(f) for f in ("tstart", "tend", "value"))) + keep = [*self.key, CHUNK, "_pd_head", "_pd_tail"] + keep += [f"_pd_prev_tail_{f}" for f in ("tstart", "tend", "value")] + keep += [f"_pd_next_head_{f}" for f in ("tstart", "tend", "value")] + for run in self.runs: + head_brk = F.col(run_col(run.name, "head_brk")) + inner = F.col(run_col(run.name, "inner")) + starts_here = F.coalesce( + F.col(run_col(run.name, "last_start")), + F.when(head_brk == 1, head.getField("tstart")), + ) + chunks = ( + chunks.withColumn( + run_col(run.name, "head_brk"), + _masked(run, _breaks(run, prev_tail, head_sample)), + ) + .withColumn(run_col(run.name, "breaks"), head_brk + inner) + .withColumn( + run_col(run.name, "base"), + F.sum(run_col(run.name, "breaks")).over(so_far) - inner, + ) + .withColumn( + run_col(run.name, "head_start"), + F.when(head_brk == 1, head.getField("tstart")).otherwise( + F.last(starts_here, ignorenulls=True).over(before) + ), + ) + .withColumn(run_col(run.name, "next_head_brk"), F.lead(head_brk).over(by_chunk)) + .withColumn( + run_col(run.name, "total"), F.sum(run_col(run.name, "breaks")).over(everything) + ) + ) + keep += [ + run_col(run.name, c) + for c in ("head_brk", "base", "head_start", "next_head_brk", "total") + ] + + # Block ends, from the block's last sample (see _block_end_at), wrapped in a + # struct: a block's end may be null, and must not be mistaken for none. + head_brk = F.col(run_col(BLOCK, "head_brk")) == 1 + count = F.col("_pd_chunk_count") + tail_single = F.when(count == 1, head_brk).otherwise(F.col("_pd_tail_brk") == 1) + seg0_single = (F.col("_pd_seg0_count") == 1) & head_brk + tail_ends = F.coalesce(F.col(run_col(BLOCK, "next_head_brk")) == 1, F.lit(True)) + tail_end = _found(_block_end_at(tail, tail_single)) + chunks = ( + chunks.withColumn(COUNT, F.sum(count).over(everything)) + .withColumn(LAST, F.last(tail).over(everything)) + # The end of the block the chunk starts in, if that block ends here. + .withColumn( + "_pd_first_end", + F.when( + F.col(run_col(BLOCK, "inner")) > 0, + _found(_block_end_at(F.col("_pd_seg0_last"), seg0_single)), + ).when(tail_ends, tail_end), + ) + .withColumn( + "_pd_tail_block_end", + F.when(tail_ends, tail_end).otherwise( + F.last("_pd_first_end", ignorenulls=True).over(later) + ), + ) + ) + return chunks.select(*keep, COUNT, LAST, "_pd_tail_block_end") + + def _resolve_neighbours_and_runs(self, df: DataFrame) -> DataFrame: + """Neighbours and run facts of every sample, across chunk boundaries.""" + first, last = F.col("_pd_first_in_chunk"), F.col("_pd_last_in_chunk") + df = ( + df.withColumn( + PREV_TS, + F.when(first, F.col("_pd_prev_tail_tstart")).otherwise(F.col("_pd_lag_ts")), + ) + .withColumn( + PREV_TE, F.when(first, F.col("_pd_prev_tail_tend")).otherwise(F.col("_pd_lag_te")) + ) + .withColumn( + PREV_VALUE, + F.when(first, F.col("_pd_prev_tail_value")).otherwise( + F.col("_pd_lag_value").cast("double") + ), + ) + .withColumn( + NEXT_TS, + F.when(last, F.col("_pd_next_head_tstart")).otherwise(F.col("_pd_lead_ts")), + ) + .withColumn( + NEXT_TE, F.when(last, F.col("_pd_next_head_tend")).otherwise(F.col("_pd_lead_te")) + ) + .withColumn( + NEXT_VALUE, + F.when(last, F.col("_pd_next_head_value")).otherwise( + F.col("_pd_lead_value").cast("double") + ), + ) + ) + for run in self.runs: + + def c(name: str, run: _Run = run) -> Column: + return F.col(run_col(run.name, name)) + + df = ( + df.withColumn( + run_col(run.name, "starts"), + F.when(first, c("head_brk") == 1).otherwise(c("brk") == 1), + ) + .withColumn(run_col(run.name, "id"), c("base") + c("seg")) + .withColumn( + run_col(run.name, "start"), F.coalesce(c("seg_start"), c("head_start")) + ) + .withColumn( + run_col(run.name, "next_starts"), + F.when(last, c("next_head_brk") == 1).otherwise(c("lead_brk") == 1), + ) + .withColumn(run_col(run.name, "is_last"), c("id") == c("total")) + ) + # The last sample of a block (also the series' last sample). + return df.withColumn( + run_col(BLOCK, "ends"), + F.coalesce(F.col(run_col(BLOCK, "next_starts")), F.lit(True)), + ) + + def _with_block_end(self, df: DataFrame) -> DataFrame: + """Each kept sample's block end: the nearest block end at or after it. + + Every block's last sample is kept, so within a chunk a reverse scan over + the kept samples finds it; a block that continues past the chunk ends + where the chunk summary says. + """ + reverse = ( + Window.partitionBy(*self.key, CHUNK) + .orderBy(*[F.col(c).desc() for c in self.order]) + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + here = F.when( + F.col(run_col(BLOCK, "ends")), + _found( + _block_end_at( + F.struct(F.col(TS).alias("tstart"), F.col(TE).alias("tend")), + F.col(run_col(BLOCK, "starts")), + ) + ), + ) + found = F.coalesce( + F.last(here, ignorenulls=True).over(reverse), F.col("_pd_tail_block_end") + ) + return df.withColumn(BLOCK_END, found.getField("end")) + + +def _row_branch( + rows: DataFrame, rules: Sequence[tuple[PushdownRule, Mapping[int, Pushdown]]], config +) -> DataFrame: + """Raw rows unchanged, and the derived rows of per-sample rules.""" + if not rules: + return rows + key = [config.container_id_col, config.channel_id_col, VARIANT_COL] + order = [config.tstart_col, config.tend_col, config.value_col] + df = rows.withColumn(TS, F.col(order[0]).cast("double")).withColumn( + TE, F.col(order[1]).cast("double") + ) + contexts = [ + (rule, SeriesContext(rule, variants, key, order, False, {})) for rule, variants in rules + ] + for rule, ctx in contexts: + df = _with_rule_columns(df, ctx, rule.row_columns(ctx)) + df = _keep(df, contexts, extra=F.col(VARIANT_COL) == RAW_VARIANT) + kept = [(rule, ctx.kept_stage()) for rule, ctx in contexts] + for rule, ctx in kept: + df = _with_rule_columns(df, ctx, rule.kept_columns(ctx)) + return _emit(df, kept, rows, order, raw=F.col(VARIANT_COL) == RAW_VARIANT) + + +# --- helpers shared by both branches ------------------------------------------------ + + +def _with_rule_columns(df: DataFrame, ctx: SeriesContext, columns: Mapping[str, Column]): + """Add a rule's columns, defined for the rule's samples only.""" + for name, column in columns.items(): + df = df.withColumn(ctx.name(name), F.when(ctx.active, column)) + return df + + +def _keep_col(rule: PushdownRule) -> str: + return f"_pd_keep_{rule.name}" + + +def _keep(df: DataFrame, contexts, extra: Column) -> DataFrame: + """Drop the samples no rule keeps (*extra* marks samples kept regardless).""" + keep = extra + for rule, ctx in contexts: + df = df.withColumn( + _keep_col(rule), F.coalesce(F.when(ctx.active, rule.keep(ctx)), F.lit(False)) + ) + keep = keep | F.col(_keep_col(rule)) + return df.where(keep) + + +def _emit( + df: DataFrame, contexts, rows: DataFrame, order: Sequence[str], raw: Column | None +) -> DataFrame: + """Explode the rules' emitted rows into rows with the columns and types of *rows*. + + *raw* marks rows passed through unchanged. + """ + elements = None + for rule, ctx in contexts: + emitted = F.when(F.col(_keep_col(rule)), rule.emit(ctx)) + elements = ( + F.when(ctx.active, emitted) if elements is None else elements.when(ctx.active, emitted) + ) + original = F.lit(False) if raw is None else raw + if raw is not None: + # One placeholder element keeps each raw row. + elements = F.when(raw, F.array(SeriesContext.element())).otherwise(elements) + element = F.col(_ELEMENT) + fields = dict(zip(order, ("tstart", "tend", "value"), strict=True)) + + def out(col: str) -> Column: + if col not in df.columns: + # Columns beyond the channel-data shape (e.g. POI string values) are not + # carried through the series scan; derived rows have none. + return F.lit(None).cast(rows.schema[col].dataType).alias(col) + if col not in fields: + return F.col(col) + derived = element.getField(fields[col]).cast(rows.schema[col].dataType) + # Times left out of an element keep the sample's own. + keep = ( + original + if fields[col] == "value" + else original | element.getField(fields[col]).isNull() + ) + return F.when(keep, F.col(col)).otherwise(derived).alias(col) + + return df.withColumn(_ELEMENT, F.explode(elements)).select(*[out(c) for c in rows.columns]) + + +def _breaks(run: _Run, prev: Sample, cur: Sample) -> Column: + """1 where *cur* starts a new run after *prev* (a null break counts as one).""" + return F.coalesce(run.breaks(prev, cur), F.lit(True)).cast("long") + + +def _masked(run: _Run, column: Column) -> Column: + if run.variants is None: + return column + return F.when(F.col(VARIANT_COL).isin(list(run.variants)), column) + + +def _flag(name: str) -> Column: + return F.coalesce(F.col(name), F.lit(False)) + + +def _nan_free(col: Column) -> Column: + """Map NaN to null so ``min`` / ``max`` behave like ``numpy.nanmin`` / ``nanmax``.""" + return F.when(~F.isnan(col), col) + + +def _found(end: Column) -> Column: + """*end* in a struct that is not null even when *end* is.""" + return F.struct(end.alias("end")) + + +def _block_end_at(last_sample: Column, single: Column) -> Column: + """``nanmax`` of a block's ``tend`` values, from its last sample. + + Within a block every ``tend`` equals the next sample's ``tstart``, so the + maximum is the last ``tend`` or, for blocks of several samples, the last + ``tstart`` if larger (NaN ``tend`` ignored). + """ + tend = _nan_free(last_sample.getField("tend")) + return F.when(single, tend).otherwise(F.greatest(last_sample.getField("tstart"), tend)) diff --git a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py index 2af656cc..4da2f34e 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -7,8 +7,9 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar +import pyspark.sql.functions as F import pyspark.sql.types as T -from pyspark.sql import DataFrame +from pyspark.sql import Column from impulse_query_engine.analyze.metadata.time_series_expression import ( TimeSeriesExpression, @@ -22,6 +23,8 @@ from impulse_query_engine.analyze.query.solvers.default_solver import TimeSeriesCache from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig + from .context import RunBreak, SeriesContext + #: Column tagging each channel-data row with the series variant it belongs to. VARIANT_COL = "__impulse_variant" #: Variant id of the unmodified (raw) channel data. @@ -88,15 +91,30 @@ def describe(self) -> str: class PushdownRule(abc.ABC): """Recognizes a subtree and computes it in Spark ahead of the pandas UDF. - A rule's derived rows use the channel-data shape - ``(container_id, channel_id, variant, tstart, tend, value)``; :meth:`decode` - turns one container/channel/variant slice of them back into the object the - original subtree evaluates to. The result must equal evaluating the subtree - in pandas on the same channel data. + Planning and decoding run in Python: :meth:`match` recognizes the subtree, + :meth:`decode` turns one container/channel/variant slice of the derived rows + back into the object the subtree evaluates to. The result must equal + evaluating the subtree in pandas on the same channel data. + + The Spark side is a set of hooks returning columns, evaluated by the shared + pushdown pipeline (:mod:`.pipeline`) for the samples of the rule's variants, + one series (channel and variant) at a time, with a :class:`SeriesContext`: + + 1. :meth:`row_columns` per sample; + 2. :meth:`keep` selects the samples that go on; + 3. :meth:`kept_columns` per kept sample; + 4. :meth:`emit` the derived rows of a kept sample. + + Rules never handle DataFrames, so every rule shares the pipeline's reads and + shuffles. Derived rows have the channel-data shape + ``(container_id, channel_id, variant, tstart, tend, value)``. """ #: Unique rule name, part of every :attr:`Pushdown.variant_key`. name: ClassVar[str] + #: Whether the rule needs its series in order (neighbours, runs, series + #: facts). Per-sample rules set ``False`` and skip the series scan. + ordered: ClassVar[bool] = True @abc.abstractmethod def match(self, node: TimeSeriesExpression) -> Pushdown | None: @@ -114,36 +132,36 @@ def result_type(self, pushdown: Pushdown) -> type: """Core-model class :meth:`decode` returns for *pushdown*.""" return SampleSeries + def runs(self) -> Mapping[str, RunBreak]: + """Run definitions beyond the built-in continuous blocks, by name. + + Each maps ``(prev, cur)`` samples to whether ``cur`` starts a new run; + :meth:`SeriesContext.run` then describes the sample's run. + """ + return {} + + def row_columns(self, ctx: SeriesContext) -> Mapping[str, Column]: + """Columns computed for every sample, in order; refer to them via ``ctx.col``.""" + return {} + + def keep(self, ctx: SeriesContext) -> Column: + """Whether a sample can emit rows; the others are dropped early.""" + return F.lit(True) + + def kept_columns(self, ctx: SeriesContext) -> Mapping[str, Column]: + """Columns computed for the kept samples, in order (``ctx.block_end`` available).""" + return {} + @abc.abstractmethod - def to_spark( - self, rows: DataFrame, variants: Mapping[int, Pushdown], config: SolverConfig - ) -> DataFrame: - """Replace the rows of this rule's *variants* with their derived rows. - - *rows* is shared by all rules and partitioned by container; the result - must keep that partitioning (windows partitioned by container id, - channel id and :data:`VARIANT_COL` reuse it). Implementations must not - union *rows* or join it to a non-broadcast frame. Small summaries - derived from *rows* (per series or per chunk) may be broadcast-joined - back; Spark reads the channel data again for them, so summaries should - read only the key and time columns to share one narrow read. - - Parameters - ---------- - rows : pyspark.sql.DataFrame - All channel-data rows of the solve, tagged with :data:`VARIANT_COL` - (one copy per variant a channel ships) and partitioned by container. - variants : Mapping[int, Pushdown] - This rule's variant ids, each with a representative pushdown carrying - its parameters. - config : SolverConfig - Provides the internal column names. + def emit(self, ctx: SeriesContext) -> Column: + """The derived rows of a kept sample. Returns ------- - pyspark.sql.DataFrame - *rows* with the same columns and types, the rows of *variants* - replaced by derived rows and all other rows unchanged. + Column + An array of :data:`~.context.ELEMENT_TYPE` structs built with + :meth:`SeriesContext.element`; empty or null emits nothing. Times + left out keep the sample's own ``tstart`` / ``tend``. """ def decode(self, rows: pd.DataFrame, pushdown: Pushdown, cache: TimeSeriesCache) -> Any: diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py index 630bbd9d..0ad2cf6f 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -4,12 +4,12 @@ import math import numbers -from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping +from typing import Any import pyspark.sql.functions as F import pyspark.sql.types as T -from pyspark.sql import Column, DataFrame, Window +from pyspark.sql import Column from impulse_query_engine.analyze.metadata.time_series_expression import ( SeriesType, @@ -18,10 +18,8 @@ TimeSeriesSelector, ) -from .pushdown import VARIANT_COL, Pushdown, PushdownContext, PushdownRule - -if TYPE_CHECKING: - from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig +from .context import SeriesContext +from .pushdown import Pushdown, PushdownContext, PushdownRule _RATE = "sample_rate" _DEFAULT_RATE = 1.0 @@ -39,36 +37,18 @@ class ResamplePushdownRule(PushdownRule): Each sample is assigned the bucket ``c`` of the first grid point at or after its ``tstart``. The last sample of a bucket holds the value of grid points - ``c, c + 1, …`` up to the next non-empty bucket, so only those samples emit rows + ``c, c + 1, …`` up to the next sample's bucket, so only those samples are kept and every row carries just the grid points it covers. Values are looked up within a block; this matches the pandas evaluation unless intervals overlap. - Every series is cut into equal time slices (chunks) holding about *chunk_rows* - samples on average, and the windows over samples run per chunk, so a worker - holds one chunk of a series at a time instead of the whole series. Gaps at - chunk boundaries, the block a chunk starts in and that block's start are - resolved from one summary row per chunk. - Grid points are computed as ``t_min + k * delta`` in IEEE double arithmetic. numpy builds that fuse this into a multiply-add can differ by one ulp for fractional grids (e.g. seconds-as-double with a rate of ``0.1``); integral grids are exact. - - Parameters - ---------- - chunk_rows : int, optional - Average number of samples per chunk (default 100 000). """ name = "resample" - def __init__(self, chunk_rows: int = 100_000): - if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, numbers.Integral): - raise TypeError(f"chunk_rows must be an integer, got {chunk_rows!r}") - if chunk_rows < 1: - raise ValueError(f"chunk_rows must be positive, got {chunk_rows}") - self.chunk_rows = int(chunk_rows) - def match(self, node: TimeSeriesExpression) -> Pushdown | None: if ( type(node) is not TimeSeriesOp @@ -98,270 +78,115 @@ def supports(self, pushdown: Pushdown, context: PushdownContext) -> bool: rate = dict(pushdown.params)[_RATE] return rate.is_integer() or not any(isinstance(t, T.IntegralType) for t in time_types) - def to_spark( - self, rows: DataFrame, variants: Mapping[int, Pushdown], config: SolverConfig - ) -> DataFrame: - cid, ch = config.container_id_col, config.channel_id_col - ts, te = config.tstart_col, config.tend_col - key = [cid, ch, VARIANT_COL] - rate = F.try_element_at( - F.create_map( - *[ - item - for vid, pushdown in sorted(variants.items()) - for item in (F.lit(vid), F.lit(dict(pushdown.params)[_RATE])) - ] + def row_columns(self, ctx: SeriesContext) -> Mapping[str, Column]: + grid = _Grid(ctx) + resampled = ctx.col("resampled") + return { + # A single sample is returned as it is. + "resampled": ctx.n >= 2, + **grid.first_index("c", ctx.tstart, F.lit(False), resampled), + **grid.first_index( + "next_c", ctx.next.tstart, F.lit(False), resampled & _next_in_block(ctx) ), - F.col(VARIANT_COL), - ) - ts_d, te_d = F.col(ts).cast("double"), F.col(te).cast("double") - # The pandas cache's sample order. - order = F.struct(ts, te) - - # Series boundaries. This summary and the chunk summary below read only the - # key and time columns, so Spark serves both from one extra narrow read. - series = rows.groupBy(*key).agg( - F.min(ts_d).alias("_rs_first"), - F.max(_ignore_nan(te_d)).alias("_rs_end"), - F.count(F.lit(1)).alias("_rs_n"), - ) - df = rows.join(F.broadcast(series), on=key) + } - # Equal time slices of about chunk_rows samples; monotone in tstart, so the - # chunks of a series follow its sample order. - first, end = F.col("_rs_first"), F.col("_rs_end") - n_chunks = F.ceil(F.col("_rs_n") / F.lit(self.chunk_rows)) - df = df.withColumn( - "_rs_chunk", - F.when( - end > first, - F.least(F.floor((ts_d - first) / (end - first) * n_chunks), n_chunks - 1), - ).otherwise(F.lit(0).cast("long")), - ) + def keep(self, ctx: SeriesContext) -> Column: + # The last sample of each bucket. + return ~ctx.col("resampled") | ~_next_in_block(ctx) | (ctx.col("next_c") != ctx.col("c")) - # Pass 1, per chunk: only running and neighbour windows (linear in the chunk - # size), sharing one sort and the container partitioning. A chunk splits into - # segments at its gaps; segment 0 continues whatever block the previous chunk - # ended in, which the chunk summary resolves. - in_chunk = Window.partitionBy(*key, "_rs_chunk").orderBy(ts, te) - running = in_chunk.rowsBetween(Window.unboundedPreceding, Window.currentRow) - prev_te = F.lag(te_d).over(in_chunk) - is_break = F.col("_rs_break") == 1 - df = ( - df.withColumn( - "_rs_break", - F.when(F.row_number().over(in_chunk) == 1, F.lit(0).cast("long")).otherwise( - (prev_te.isNull() | ((ts_d - prev_te) != 0)).cast("long") - ), - ) - .withColumn("_rs_seg", F.sum("_rs_break").over(running)) - .withColumn( - "_rs_seg_t_min", F.last(F.when(is_break, ts_d), ignorenulls=True).over(running) - ) - # Largest tend so far within the segment (segments are numbered in order). - .withColumn( - "_rs_seg_t_max", - F.max(F.struct(F.col("_rs_seg"), _ignore_nan(te_d).alias("te"))) - .over(running) - .getField("te"), - ) - .withColumn("_rs_next_ts", F.lead(ts_d).over(in_chunk)) - .withColumn("_rs_next_break", F.lead("_rs_break").over(in_chunk)) - ) + def kept_columns(self, ctx: SeriesContext) -> Mapping[str, Column]: + grid = _Grid(ctx) + t_max = ctx.block_end + # The last block keeps a grid point at t_max when the series ends with a + # zero-duration sample. + inclusive = ctx.block.is_last & (ctx.last.tstart == ctx.last.tend) + length = F.greatest(F.ceil(((t_max + grid.rate) - grid.t_min) / grid.rate), F.lit(0)) + n_grid = ctx.col("n_grid") + return { + **grid.first_index("grid_end", t_max, inclusive, ctx.col("resampled")), + "n_grid": F.least(ctx.col("grid_end"), length), + # Grid points up to the next sample's bucket, or up to the block's end. + "c_stop": F.when(_next_in_block(ctx), F.least(ctx.col("next_c"), n_grid)).otherwise( + n_grid + ), + } - # One summary row per chunk, scanned in chunk order: whether the chunk starts - # a block, the block of its segment 0 and that block's start, plus the number - # of blocks and the last sample of the series. - chunks = df.groupBy(*key, "_rs_chunk").agg( - F.min_by(ts_d, order).alias("_rs_head_ts"), - F.max_by(F.struct(ts_d.alias("ts"), te_d.alias("te")), order).alias("_rs_tail"), - F.sum("_rs_break").alias("_rs_inner_breaks"), - F.max(F.when(is_break, ts_d)).alias("_rs_last_start"), - ) - by_chunk = Window.partitionBy(*key).orderBy("_rs_chunk") - all_chunks = by_chunk.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) - head_ts, head_break = F.col("_rs_head_ts"), F.col("_rs_head_break") == 1 - prev_tail_te = F.lag(F.col("_rs_tail.te")).over(by_chunk) - last_start = F.coalesce(F.col("_rs_last_start"), F.when(head_break, head_ts)) - chunks = ( - chunks.withColumn( - "_rs_head_break", - (prev_tail_te.isNull() | ((head_ts - prev_tail_te) != 0)).cast("long"), - ) - .withColumn("_rs_breaks", F.col("_rs_head_break") + F.col("_rs_inner_breaks")) - .withColumn( - "_rs_block_base", - F.sum("_rs_breaks").over( - by_chunk.rowsBetween(Window.unboundedPreceding, Window.currentRow) - ) - - F.col("_rs_inner_breaks"), - ) - .withColumn( - "_rs_head_t_min", - F.when(head_break, head_ts).otherwise( - F.last(last_start, ignorenulls=True).over( - by_chunk.rowsBetween(Window.unboundedPreceding, -1) - ) - ), - ) - .withColumn("_rs_nblk", F.sum("_rs_breaks").over(all_chunks)) - .withColumn("_rs_last", F.last("_rs_tail").over(all_chunks)) - .select(*key, "_rs_chunk", "_rs_block_base", "_rs_head_t_min", "_rs_nblk", "_rs_last") + def emit(self, ctx: SeriesContext) -> Column: + grid = _Grid(ctx) + c, c_stop, n_grid = ctx.col("c"), ctx.col("c_stop"), ctx.col("n_grid") + points = F.transform( + F.sequence(c, c_stop - 1), + lambda k: ctx.element( + grid.point(k), + F.when(k == n_grid - 1, ctx.block_end).otherwise(grid.point(k + 1)), + ctx.value, + ), ) - df = ( - df.join(F.broadcast(chunks), on=[*key, "_rs_chunk"]) - .withColumn("_rs_block", F.col("_rs_block_base") + F.col("_rs_seg")) - .withColumn("_rs_t_min", F.coalesce(F.col("_rs_seg_t_min"), F.col("_rs_head_t_min"))) + return F.when(~ctx.col("resampled"), F.array(ctx.element(value=ctx.value))).when( + c_stop > c, points ) - # numpy.arange(t_min, t_max + rate, rate): length ceil(((t_max + rate) - t_min) / rate), - # filled as t_min, t_min + rate, then t_min + k * ((t_min + rate) - t_min). - t_min = F.col("_rs_t_min") - delta = (t_min + rate) - t_min - def grid_point(k: Column) -> Column: - return ( - F.when(k == 0, t_min) - .when(k == 1, t_min + rate) - .otherwise(t_min + k.cast("double") * delta) - ) +def _next_in_block(ctx: SeriesContext) -> Column: + """The next sample exists and continues the block.""" + return F.coalesce(~ctx.block.next_starts, F.lit(False)) - resampled = F.col(VARIANT_COL).isin(sorted(variants)) & (F.col("_rs_n") >= 2) - next_in_block = F.coalesce(F.col("_rs_next_break") == 0, F.lit(False)) - df = _with_first_grid_index( - df, "_rs_c", ts_d, F.lit(False), resampled, grid_point, t_min, delta - ) - df = _with_first_grid_index( - df, - "_rs_next_sample_c", - F.col("_rs_next_ts"), - F.lit(False), - resampled & next_in_block, - grid_point, - t_min, - delta, - ) - c = F.col("_rs_c") - # Only the last sample of each bucket emits grid points. A chunk's last sample - # is kept too: its bucket may continue in the next chunk. - df = df.where(~resampled | ~next_in_block | (F.col("_rs_next_sample_c") != c)) - # Pass 2, over the remaining rows of each block (about one per grid point): the - # next non-empty bucket, and the block's t_max (each segment's last sample - # carries its maximum). A chunk's last sample whose bucket continues in the - # next chunk is followed by its own bucket and emits nothing. Rows this rule - # leaves alone stay grouped by chunk. - in_block = Window.partitionBy( - *key, "_rs_block", F.when(~resampled, F.col("_rs_chunk")) - ).orderBy("_rs_c", ts, te) - whole_block = in_block.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing) - df = df.withColumn("_rs_next_c", F.lead("_rs_c").over(in_block)).withColumn( - "_rs_t_max", F.max("_rs_seg_t_max").over(whole_block) - ) - next_c, t_max = F.col("_rs_next_c"), F.col("_rs_t_max") - block_end = next_c.isNull() - # The last block keeps a grid point at t_max when the series ends with a - # zero-duration sample. - inclusive = (F.col("_rs_block") == F.col("_rs_nblk")) & ( - F.col("_rs_last.ts") == F.col("_rs_last.te") - ) - length = F.greatest(F.ceil(((t_max + rate) - t_min) / rate), F.lit(0)) - df = _with_first_grid_index( - df, "_rs_grid_end", t_max, inclusive, resampled, grid_point, t_min, delta - ) - n_grid = F.least(F.col("_rs_grid_end"), length) - # A bucket's last sample emits the grid points up to the next non-empty bucket - # of its block, or up to the block's end. - df = df.withColumn( - "_rs_c_stop", F.when(block_end, n_grid).otherwise(F.least(next_c, n_grid)) - ) +class _Grid: + """The numpy.arange grid of the sample's block. - def element(kind: int, t: Column, tend: Column) -> Column: - return F.struct( - F.lit(kind).alias("kind"), - t.cast("double").alias("t"), - tend.cast("double").alias("tend"), - ) + ``numpy.arange(t_min, t_max + rate, rate)`` has length + ``ceil(((t_max + rate) - t_min) / rate)`` and is filled as ``t_min``, + ``t_min + rate``, then ``t_min + k * ((t_min + rate) - t_min)``. + """ - c_stop = F.col("_rs_c_stop") - grid = F.transform( - F.sequence(c, c_stop - 1), - lambda k: element( - 1, - grid_point(k), - F.when(k == n_grid - 1, t_max).otherwise(grid_point(k + 1)), - ), - ) - no_time = F.lit(None) - # Series this rule leaves alone keep their row (kind 2); of the resampled - # series only the grid points survive (kind 1). explode drops the rest. - elements = F.when(~resampled, F.array(element(2, no_time, no_time))).when(c_stop > c, grid) - df = df.withColumn("_rs_e", F.explode(elements)) + def __init__(self, ctx: SeriesContext): + self.ctx = ctx + self.rate = ctx.param(_RATE) + self.t_min = ctx.block.start + self.delta = (self.t_min + self.rate) - self.t_min - is_grid = F.col("_rs_e.kind") == 1 - derived = { - ts: F.col("_rs_e.t").cast(rows.schema[ts].dataType), - te: F.col("_rs_e.tend").cast(rows.schema[te].dataType), - } - return df.select( - *[ - ( - F.when(is_grid, derived[col]).otherwise(F.col(col)).alias(col) - if col in derived - else col - ) - for col in rows.columns - ] + def point(self, k: Column) -> Column: + return ( + F.when(k == 0, self.t_min) + .when(k == 1, self.t_min + self.rate) + .otherwise(self.t_min + k.cast("double") * self.delta) ) + def first_index( + self, name: str, x: Column, strict: Column, active: Column + ) -> dict[str, Column]: + """Columns ending in *name*: the smallest ``k >= 0`` with ``point(k) >= x`` (``> x`` if *strict*). -def _with_first_grid_index( - df: DataFrame, - name: str, - x: Column, - strict: Column, - active: Column, - grid_point: Callable[[Column], Column], - t_min: Column, - delta: Column, -) -> DataFrame: - """Add *name*: the smallest ``k >= 0`` with ``grid_point(k) >= x`` (``> x`` if *strict*). + Estimated from ``(x - t_min) / delta`` and corrected by one step against the + exact grid points; a result that still is not exact raises instead of + misplacing a sample. Only evaluated where *active*. + """ - Estimated from ``(x - t_min) / delta`` and corrected by one step against the - exact grid points; a result that still is not exact raises instead of - misplacing a sample. Only evaluated where *active*. - """ + def reached(k: Column) -> Column: + g = self.point(k) + return F.when(strict, g > x).otherwise(g >= x) - def reached(k: Column) -> Column: - g = grid_point(k) - return F.when(strict, g > x).otherwise(g >= x) - - estimate, k = F.col(f"{name}_estimate"), F.col(name) - return ( - df.withColumn( - f"{name}_estimate", - F.when( + estimate, step = self.ctx.col(f"{name}_estimate"), self.ctx.col(f"{name}_step") + return { + f"{name}_estimate": F.when( active, - F.when(delta > 0, F.greatest(F.ceil((x - t_min) / delta), F.lit(0))).otherwise( + F.when( + self.delta > 0, + F.greatest(F.ceil((x - self.t_min) / self.delta), F.lit(0)), + ).otherwise( F.raise_error( F.lit("resample: sample_rate is below the resolution of the timestamps") ) ), ), - ) - .withColumn( - name, - F.when((estimate > 0) & reached(estimate - 1), estimate - 1) + f"{name}_step": F.when((estimate > 0) & reached(estimate - 1), estimate - 1) .when(~reached(estimate), estimate + 1) .otherwise(estimate), - ) - .withColumn( - name, - F.when(~active, F.lit(None)) - .when(reached(k) & ((k == 0) | ~reached(k - 1)), k) + name: F.when(~active, F.lit(None)) + .when(reached(step) & ((step == 0) | ~reached(step - 1)), step) .otherwise(F.raise_error(F.lit("resample: grid index could not be determined"))), - ) - ) + } def _sample_rate(args: tuple, kwargs: dict[str, Any]) -> float | None: @@ -373,8 +198,3 @@ def _sample_rate(args: tuple, kwargs: dict[str, Any]) -> float | None: return None rate = float(rate) return rate if math.isfinite(rate) and rate > 0 else None - - -def _ignore_nan(col: Column) -> Column: - """Map NaN to null so ``min`` / ``max`` behave like ``numpy.nanmin`` / ``nanmax``.""" - return F.when(~F.isnan(col), col) diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index 1b4c599d..c56a856c 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -23,6 +23,7 @@ PushdownContext, PushdownPlan, QueryOptimizer, + apply_pushdowns, ) from impulse_query_engine.model.series.points_in_time_series import PointsInTimeSeries from impulse_query_engine.model.series.sample_series import SampleSeries @@ -1199,7 +1200,7 @@ def solve(self, query, channels_df, selections, dtypes) -> DataFrame: DataFrame containing results for each container. """ col_map = self.config.col_map - q, joined_df, container_count, plan = self._prepare_channels_join(query, channels_df) + q, joined_df, container_count, plan, _ = self._prepare_channels_join(query, channels_df) container_tag_cols, container_metric_cols = self._container_meta_cols(query) schema = self._build_solve_output_schema(q, selections, dtypes) @@ -1224,8 +1225,8 @@ def solve(self, query, channels_df, selections, dtypes) -> DataFrame: ) def _prepare_channels_join( - self, query, channels_df - ) -> tuple[DataFrame, DataFrame, int, PushdownPlan | None]: + self, query, channels_df, native_selections=None + ) -> tuple[DataFrame, DataFrame, int, PushdownPlan | None, dict | None]: """Shared prelude for :meth:`solve` and :meth:`solve_calculated_channels`. Applies optional per-channel unit conversion, reads and column-maps the @@ -1248,15 +1249,20 @@ def _prepare_channels_join( Query object containing database and filter information. channels_df : pyspark.sql.DataFrame Channel-match DataFrame from the filter pipeline. + native_selections : list, optional + Calculated-channel selections that may be solved without the UDF + (see :meth:`PushdownPlan.native_targets`). Returns ------- - tuple[DataFrame, DataFrame, int, PushdownPlan or None] - ``(channels_q, joined_df, container_count, plan)`` — the column-mapped - channel-data DataFrame (authoritative for output ``container_id`` / - ``channel_id`` types), the join fed to the grouped-map UDF, the - number of distinct containers, and the pushdown plan the UDF must - receive (``None`` when nothing is pushed down). + tuple[DataFrame, DataFrame, int, PushdownPlan or None, dict or None] + ``(channels_q, joined_df, container_count, plan, native_targets)`` — + the column-mapped channel-data DataFrame (authoritative for output + ``container_id`` / ``channel_id`` types), the join fed to the + grouped-map UDF, the number of distinct containers, the pushdown + plan the UDF must receive (``None`` when nothing is pushed down), + and the native targets of *native_selections* (``None`` unless + they are solved natively; ``joined_df`` then holds no raw rows). """ source_unit_col = self.config.source_unit_col target_unit_col = self.config.target_unit_col @@ -1304,14 +1310,19 @@ def _prepare_channels_join( container_count = channels_df.select(self.config.container_id_col).distinct().count() plan = self._plan_pushdowns(query, q) if container_count > 0 else None channel_data = q + targets = None if plan is not None: - channel_data = self._apply_pushdowns(q, channels_df, plan, container_count) + if native_selections is not None: + targets = plan.native_targets(native_selections) + # Raw rows are only read by the UDF, for selectors that are not pushed down. + include_raw = targets is None and bool(plan.raw_selector_ids) + channel_data = self._apply_pushdowns(q, channels_df, plan, include_raw) joined_df = channel_data.join( F.broadcast(channels_df), on=[self.config.container_id_col, self.config.channel_id_col], ) - return q, joined_df, container_count, plan + return q, joined_df, container_count, plan, targets def _plan_pushdowns(self, query, channels_q: DataFrame) -> PushdownPlan | None: """Plan the pushdowns for *query*, or ``None`` when disabled or nothing applies.""" @@ -1330,30 +1341,24 @@ def _apply_pushdowns( channels_q: DataFrame, channels_df: DataFrame, plan: PushdownPlan, - container_count: int, + include_raw: bool, ) -> DataFrame: """Replace raw channel data with the derived rows of *plan* where possible. Each matched channel is tagged with the variants it ships (:meth:`PushdownPlan.channel_variants`): its raw data (variant ``0``) while any of its selectors still needs it, otherwise the derived - variants of its selectors, computed by their rules. The result carries - a :data:`VARIANT_COL` column and keeps the channel-data columns and types. - - The channel data is read once and partitioned by container exactly like - :meth:`_apply_grouped_map` does, so the rules' per-channel windows and - the grouped-map UDF share a single shuffle. + variants of its selectors, computed by their rules (see + :func:`~impulse_query_engine.analyze.query.optimizer.apply_pushdowns`). + The result carries a :data:`VARIANT_COL` column and keeps the + channel-data columns and types; raw rows are kept only if *include_raw*. """ key = [self.config.container_id_col, self.config.channel_id_col] variants = channels_df.select( *key, F.explode(plan.channel_variants(F.col("selector_ids"))).alias(VARIANT_COL) ) - rows = channels_q.join(F.broadcast(variants), on=key).repartition( - container_count, self.config.container_id_col - ) - for rule, rule_variants in plan.variants_by_rule(): - rows = rule.to_spark(rows, rule_variants, self.config) - return rows + rows = channels_q.join(F.broadcast(variants), on=key) + return apply_pushdowns(rows, plan, self.config, include_raw=include_raw) @staticmethod def _query_contains_poi_selections(selections: Iterable[Any]) -> bool: @@ -1662,10 +1667,11 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame the derived rows and no pandas UDF runs. """ col_map = self.config.col_map - q, joined_df, container_count, plan = self._prepare_channels_join(query, channels_df) + q, joined_df, container_count, plan, targets = self._prepare_channels_join( + query, channels_df, native_selections=selections + ) self._validate_calculated_channel_id_type(q) - targets = plan.native_targets(selections) if plan is not None else None if targets is not None: logger.info("Calculated channels solved natively in Spark, without the pandas UDF") native = self._native_calculated_channel_rows(joined_df, q, targets) @@ -1801,9 +1807,8 @@ def _native_calculated_channel_rows( ) ) has_conversion = conv in joined_df.columns - # Raw rows (variant 0) have no targets, so inline drops them. Not filtering - # them explicitly keeps the filter out of the pushdown rules' summaries, which - # would otherwise stop sharing one read of the channel data. + # joined_df holds no raw rows here; any other row without targets is + # dropped by inline. rows = joined_df.select( cid, ts, diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py new file mode 100644 index 00000000..3d753cbb --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py @@ -0,0 +1,183 @@ +"""Helpers for pushdown pipeline tests: a runner and test-only rules.""" + +import contextlib +from collections.abc import Callable +from unittest import mock + +import numpy as np +import pyspark.sql.functions as F +import pyspark.sql.types as T +from pyspark.sql import Column, DataFrame + +from impulse_query_engine.analyze.metadata.time_series_expression import ( + SeriesType, + TimeSeriesOp, + TimeSeriesSelector, +) +from impulse_query_engine.analyze.query.optimizer import ( + VARIANT_COL, + Pushdown, + PushdownPlan, + PushdownRule, + SeriesContext, + apply_pushdowns, + pipeline, +) +from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig +from impulse_query_engine.model.series.sample_series import SampleSeries + +CONFIG = SolverConfig() + + +@contextlib.contextmanager +def chunk_rows(n: int | None): + """Run the pipeline with *n* samples per chunk (the default when ``None``).""" + if n is None: + yield + return + with mock.patch.object(pipeline, "CHUNK_ROWS", n): + yield + + +def channel_rows(spark, series_by_key, time_type=T.LongType()) -> DataFrame: + """Tagged channel data for ``{(cid, ch, variant): (tstarts, tends, values)}``.""" + schema = T.StructType( + [ + T.StructField("container_id", T.LongType()), + T.StructField("channel_id", T.LongType()), + T.StructField("tstart", time_type), + T.StructField("tend", time_type), + T.StructField("value", T.DoubleType()), + T.StructField(VARIANT_COL, T.IntegerType()), + ] + ) + to_time = int if isinstance(time_type, T.IntegralType) else float + rows = [ + (cid, ch, to_time(ts), to_time(te), None if v is None else float(v), vid) + for (cid, ch, vid), (tstarts, tends, values) in series_by_key.items() + for ts, te, v in zip(tstarts, tends, values, strict=True) + ] + return spark.createDataFrame(rows, schema=schema) + + +def plan_for(pushdowns: dict[int, Pushdown]) -> PushdownPlan: + """A plan computing *pushdowns* as the given variants.""" + rules = tuple(dict.fromkeys(p.rule for p in pushdowns.values())) + return PushdownPlan( + rules=rules, + variant_ids={p.variant_key: vid for vid, p in pushdowns.items()}, + pushdowns_by_variant=dict(pushdowns), + ) + + +def run( + spark, + series_by_key, + pushdowns: dict[int, Pushdown], + time_type=T.LongType(), + chunk_size: int | None = None, + include_raw: bool = True, +) -> dict[tuple, SampleSeries]: + """Run *pushdowns* on the series; ``{(cid, ch, variant): SampleSeries}`` like the cache.""" + df = channel_rows(spark, series_by_key, time_type) + with chunk_rows(chunk_size): + result = apply_pushdowns(df, plan_for(pushdowns), CONFIG, include_raw=include_raw) + assert result.schema == df.schema + pdf = result.toPandas() + out = {} + for key, group in pdf.groupby(["container_id", "channel_id", VARIANT_COL]): + group = group.sort_values(["tstart", "tend"], kind="stable") + out[tuple(int(k) for k in key)] = SampleSeries( + group["tstart"], group["tend"], group["value"].astype(np.float64) + ) + return out + + +def random_series(rng, n, nan_tend=False): + """Samples with gaps, zero-duration samples and NaN values, in series order. + + With *nan_tend*, some samples have a NaN ``tend`` (use a double time column). + """ + tstarts, tends, values = [], [], [] + t, after_point = 0, False + for _ in range(n): + if rng.random() < 0.15: + t += int(rng.integers(1, 6)) + duration = 0 if not after_point and rng.random() < 0.1 else int(rng.integers(1, 5)) + tstarts.append(t) + nan = nan_tend and duration > 0 and rng.random() < 0.1 + tends.append(np.nan if nan else t + duration) + values.append(np.nan if rng.random() < 0.1 else float(rng.integers(0, 4))) + t, after_point = t + duration, duration == 0 + return tstarts, tends, values + + +def _selector_op(node, operation: str) -> TimeSeriesSelector | None: + """The continuous selector of ``selector.()``, or ``None``.""" + if ( + type(node) is not TimeSeriesOp + or node.optype != "cls" + or node.operation != operation + or len(node.args) != 1 + or node.kwargs + or not isinstance(node.args[0], TimeSeriesSelector) + or node.args[0].series_type is not SeriesType.CONTINUOUS + ): + return None + return node.args[0] + + +class DiffRule(PushdownRule): + """Test-only ``channel.diff()``: the value change within a block, 0 at block starts.""" + + name = "diff" + + def match(self, node): + selector = _selector_op(node, "diff") + if selector is None: + return None + return Pushdown(rule=self, selector=selector, params=(), consumed=(selector,)) + + def row_columns(self, ctx: SeriesContext): + return {"diff": F.when(ctx.block.starts, F.lit(0.0)).otherwise(ctx.value - ctx.prev.value)} + + def emit(self, ctx: SeriesContext) -> Column: + return F.array(ctx.element(value=ctx.col("diff"))) + + +class OffsetRule(PushdownRule): + """Test-only per-sample rule adding its ``offset`` parameter to every value.""" + + name = "offset" + ordered = False + + def match(self, node): + return None + + def emit(self, ctx: SeriesContext) -> Column: + return F.array(ctx.element(value=ctx.value + ctx.param("offset"))) + + +class ProbeRule(PushdownRule): + """Test-only rule emitting every sample with ``probe(ctx)`` as its value.""" + + match = OffsetRule.match + + def __init__( + self, + name: str, + probe: Callable[[SeriesContext], Column], + runs: dict | None = None, + ): + self.name = name + self.probe = probe + self._runs = runs or {} + + def runs(self): + return self._runs + + def emit(self, ctx: SeriesContext) -> Column: + return F.array(ctx.element(value=self.probe(ctx))) + + def pushdown(self, selector: TimeSeriesSelector) -> Pushdown: + return Pushdown(rule=self, selector=selector, params=(), consumed=(selector,)) diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py new file mode 100644 index 00000000..327b85a9 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py @@ -0,0 +1,235 @@ +# pylint: disable=missing-function-docstring +"""The pushdown pipeline's services match a straightforward per-series reference. + +Probe rules emit one service column per sample; the series are cut into chunks of +a few samples, so every service has to carry information across chunk boundaries. +""" + +import re + +import numpy as np +import numpy.testing as nptest +import pyspark.sql.functions as F +import pyspark.sql.types as T +import pytest + +from impulse_query_engine.analyze.metadata.tag_expression import TagSelector +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.query.optimizer import VARIANT_COL, Pushdown, apply_pushdowns +from impulse_query_engine.model.series.sample_series import SampleSeries +from tests.conftest import spark +from tests.impulse_query_engine.unit.analyze.query.optimizer.pipeline_helpers import ( + CONFIG, + DiffRule, + OffsetRule, + ProbeRule, + channel_rows, + plan_for, + random_series, + run, +) + +_X = TimeSeriesSelector(TagSelector("channel_name") == "x") + + +def _same_value(prev, cur) -> bool: + # Spark compares NaN as equal to NaN. + return prev == cur or (np.isnan(prev) and np.isnan(cur)) + + +def _value_change(prev, cur): + return prev.value.isNull() | (prev.value != cur.value) + + +def _runs(breaks: list[bool], tstarts: list) -> dict[str, list]: + ids = np.cumsum(breaks).tolist() + starts, start = [], None + for brk, ts in zip(breaks, tstarts, strict=True): + start = ts if brk else start + starts.append(start) + return { + "starts": [float(b) for b in breaks], + "id": ids, + "start": starts, + "next_starts": [float(b) for b in breaks[1:]] + [None], + "is_last": [float(i == ids[-1]) for i in ids], + } + + +def _reference(tstarts, tends, values) -> dict[str, list]: + """Every probed service for one series, computed sample by sample.""" + n = len(tstarts) + block_breaks = [i == 0 or tstarts[i] - tends[i - 1] != 0 for i in range(n)] + level_breaks = [i == 0 or not _same_value(values[i - 1], values[i]) for i in range(n)] + block = _runs(block_breaks, tstarts) + block_end = [] + for i in range(n): + members = [tends[j] for j in range(n) if block["id"][j] == block["id"][i]] + finite = [t for t in members if not np.isnan(t)] + block_end.append(max(finite) if finite else None) + return { + "prev_tstart": [None, *tstarts[:-1]], + "prev_tend": [None, *tends[:-1]], + "prev_value": [None, *values[:-1]], + "next_tstart": [*tstarts[1:], None], + "next_value": [*values[1:], None], + "n": [n] * n, + "last_tstart": [tstarts[-1]] * n, + "last_tend": [tends[-1]] * n, + **{f"block_{name}": column for name, column in block.items()}, + "block_end": block_end, + **{f"level_{name}": column for name, column in _runs(level_breaks, tstarts).items()}, + } + + +_PROBES = { + "prev_tstart": lambda ctx: ctx.prev.tstart, + "prev_tend": lambda ctx: ctx.prev.tend, + "prev_value": lambda ctx: ctx.prev.value, + "next_tstart": lambda ctx: ctx.next.tstart, + "next_value": lambda ctx: ctx.next.value, + "n": lambda ctx: ctx.n, + "last_tstart": lambda ctx: ctx.last.tstart, + "last_tend": lambda ctx: ctx.last.tend, + "block_starts": lambda ctx: ctx.block.starts, + "block_id": lambda ctx: ctx.block.id, + "block_start": lambda ctx: ctx.block.start, + "block_next_starts": lambda ctx: ctx.block.next_starts, + "block_is_last": lambda ctx: ctx.block.is_last, + "block_end": lambda ctx: ctx.block_end, + "level_starts": lambda ctx: ctx.run("level").starts, + "level_id": lambda ctx: ctx.run("level").id, + "level_start": lambda ctx: ctx.run("level").start, + "level_next_starts": lambda ctx: ctx.run("level").next_starts, + "level_is_last": lambda ctx: ctx.run("level").is_last, +} + + +def _probe_pushdowns() -> dict[int, Pushdown]: + probes = {} + for vid, (name, probe) in enumerate(_PROBES.items(), start=1): + runs = {"level": _value_change} if name.startswith("level_") else None + probes[vid] = ProbeRule(f"probe_{name}", probe, runs).pushdown(_X) + return probes + + +def _as_float(column: list) -> np.ndarray: + return np.array([np.nan if v is None else float(v) for v in column]) + + +@pytest.mark.parametrize( + ("time_type", "nan_tend"), [(T.LongType(), False), (T.DoubleType(), True)] +) +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 7, None]) +def test_services_across_chunk_boundaries(spark, chunk_size, time_type, nan_tend): + rng = np.random.default_rng(17) + pushdowns = _probe_pushdowns() + series, expected = {}, {} + for cid, n in enumerate([1, 2, 5, 40, 90]): + samples = random_series(rng, n, nan_tend=nan_tend) + reference = _reference(*samples) + for vid, pushdown in pushdowns.items(): + series[(cid, 1, vid)] = samples + expected[(cid, 1, vid)] = ( + samples, + reference[pushdown.rule.name.removeprefix("probe_")], + ) + actual = run(spark, series, pushdowns, time_type=time_type, chunk_size=chunk_size) + assert actual.keys() == expected.keys() + for key, ((tstarts, tends, _), column) in expected.items(): + name = pushdowns[key[2]].rule.name + nptest.assert_array_equal(actual[key].tstarts, tstarts, err_msg=name) + nptest.assert_array_equal(actual[key].tends, tends, err_msg=name) + nptest.assert_array_equal(actual[key].values, _as_float(column), err_msg=name) + + +def test_diff_matches_sample_series(spark): + rng = np.random.default_rng(23) + series, expected = {}, {} + for cid, n in enumerate([1, 2, 30, 70]): + samples = random_series(rng, n) + series[(cid, 1, 1)] = samples + expected[(cid, 1, 1)] = SampleSeries(*samples).diff() + pushdowns = {1: DiffRule().match(_X.diff())} + for chunk_size in (1, 3, None): + actual = run(spark, series, pushdowns, chunk_size=chunk_size) + for key, exp in expected.items(): + nptest.assert_array_equal(actual[key].tstarts, exp.tstarts) + nptest.assert_array_equal(actual[key].tends, exp.tends) + nptest.assert_array_equal(actual[key].values, exp.values) + + +def _offset(value: float) -> Pushdown: + rule = OffsetRule() + return Pushdown(rule=rule, selector=_X, params=(("offset", value),), consumed=(_X,)) + + +def test_per_sample_rules_share_the_row_branch_with_raw_rows(spark): + samples = ([0, 1, 5], [1, 2, 6], [1.0, 2.0, 4.0]) + series = {(1, 1, 0): samples, (1, 1, 1): samples, (1, 1, 2): samples, (2, 1, 3): samples} + pushdowns = {1: _offset(10.0), 2: _offset(0.5), 3: DiffRule().match(_X.diff())} + actual = run(spark, series, pushdowns, chunk_size=2) + nptest.assert_array_equal(actual[(1, 1, 0)].values, [1.0, 2.0, 4.0]) + nptest.assert_array_equal(actual[(1, 1, 1)].values, [11.0, 12.0, 14.0]) + nptest.assert_array_equal(actual[(1, 1, 2)].values, [1.5, 2.5, 4.5]) + nptest.assert_array_equal(actual[(2, 1, 3)].values, [0.0, 1.0, 0.0]) + nptest.assert_array_equal(actual[(1, 1, 1)].tstarts, [0, 1, 5]) + + +def test_per_sample_rules_skip_the_series_scan(spark): + df = channel_rows(spark, {(1, 1, 1): ([0, 1], [1, 2], [1.0, 2.0])}) + result = apply_pushdowns(df, plan_for({1: _offset(1.0)}), CONFIG, include_raw=False) + plan = result._jdf.queryExecution().optimizedPlan().toString() + assert "_pd_chunk" not in plan + assert "Window" not in plan + assert "Aggregate" not in plan + + +def test_raw_rows_only_when_asked(spark): + samples = ([0, 1], [1, 2], [1.0, 2.0]) + series = {(1, 1, 0): samples, (1, 1, 1): samples} + pushdowns = {1: DiffRule().match(_X.diff())} + assert set(run(spark, series, pushdowns)) == {(1, 1, 0), (1, 1, 1)} + assert set(run(spark, series, pushdowns, include_raw=False)) == {(1, 1, 1)} + + +def test_ordered_services_are_refused_to_per_sample_rules(spark): + class _Neighbourly(OffsetRule): + name = "neighbourly" + + def emit(self, ctx): + return F.array(ctx.element(value=ctx.prev.value)) + + df = channel_rows(spark, {(1, 1, 1): ([0], [1], [1.0])}) + pushdown = Pushdown(rule=_Neighbourly(), selector=_X, params=(), consumed=(_X,)) + with pytest.raises(TypeError, match="needs an ordered rule"): + apply_pushdowns(df, plan_for({1: pushdown}), CONFIG) + + +def test_block_end_is_only_available_to_kept_samples(spark): + rule = ProbeRule("early_block_end", lambda ctx: ctx.value) + rule.row_columns = lambda ctx: {"end": ctx.block_end} + df = channel_rows(spark, {(1, 1, 1): ([0], [1], [1.0])}) + with pytest.raises(TypeError, match="kept samples"): + apply_pushdowns(df, plan_for({1: rule.pushdown(_X)}), CONFIG) + + +def test_one_shuffle_of_the_samples(spark): + rng = np.random.default_rng(3) + series = {(cid, 1, 1): random_series(rng, 30) for cid in range(3)} + df = channel_rows(spark, series) + plan = plan_for({1: DiffRule().match(_X.diff())}) + result = apply_pushdowns(df, plan, CONFIG, include_raw=False) + result.collect() + final = result._jdf.queryExecution().executedPlan().toString().split("== Initial Plan ==")[0] + # Shuffles (not broadcasts), each listed once even when reused. + exchanges = [ + re.sub(r"#\d+L?", "", e) for e in re.findall(r"(?<=- )Exchange (\w+\([^)]*\))", final) + ] + key = f"container_id, channel_id, {VARIANT_COL}" + samples = [e for e in exchanges if "_pd_chunk" in e] + assert len(samples) == 1 and samples[0].startswith(f"hashpartitioning({key}, _pd_chunk") + # The others are the per-series boundaries and chunk summaries. + assert all(e.startswith(f"hashpartitioning({key}") for e in exchanges) + # One full read, shared by the chunk pass and its summaries, and the boundaries. + assert len(re.findall(r"Scan ExistingRDD", final)) == 2 diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py index cbac91d7..c69e38ca 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -11,53 +11,34 @@ from impulse_query_engine.analyze.metadata.tag_expression import TagSelector from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector from impulse_query_engine.analyze.query.optimizer import VARIANT_COL, ResamplePushdownRule -from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig from impulse_query_engine.model.series.sample_series import SampleSeries from tests.conftest import spark +from tests.impulse_query_engine.unit.analyze.query.optimizer.pipeline_helpers import ( + CONFIG, + DiffRule, + channel_rows, + plan_for, + random_series, + run, +) +from impulse_query_engine.analyze.query.optimizer import apply_pushdowns _RULE = ResamplePushdownRule() -_CONFIG = SolverConfig() +_X = TimeSeriesSelector(TagSelector("channel_name") == "x") def _pushdown(rate): - selector = TimeSeriesSelector(TagSelector("channel_name") == "x") - return _RULE.match(selector.resample(rate)) + return _RULE.match(_X.resample(rate)) def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType(), chunk_rows=None): - """Run the rule on ``{(cid, ch, variant): (tstarts, tends, values)}``. + """Resample ``{(cid, ch, variant): (tstarts, tends, values)}`` at ``{variant: rate}``. - Returns ``{(cid, ch, variant): SampleSeries}`` built like the pandas cache does. + Variant 0 rows are raw rows. Returns ``{(cid, ch, variant): SampleSeries}`` + built like the pandas cache does. """ - schema = T.StructType( - [ - T.StructField("container_id", T.LongType()), - T.StructField("channel_id", T.LongType()), - T.StructField(VARIANT_COL, T.IntegerType()), - T.StructField("tstart", time_type), - T.StructField("tend", time_type), - T.StructField("value", T.DoubleType()), - ] - ) - to_time = int if isinstance(time_type, T.IntegralType) else float - rows = [ - (cid, ch, vid, to_time(ts), to_time(te), None if v is None else float(v)) - for (cid, ch, vid), (tstarts, tends, values) in series_by_key.items() - for ts, te, v in zip(tstarts, tends, values, strict=True) - ] - df = spark.createDataFrame(rows, schema=schema) - rule = _RULE if chunk_rows is None else ResamplePushdownRule(chunk_rows=chunk_rows) - result = rule.to_spark(df, {vid: _pushdown(rate) for vid, rate in rates.items()}, _CONFIG) - assert result.schema == df.schema - - pdf = result.toPandas() - out = {} - for key, group in pdf.groupby(["container_id", "channel_id", VARIANT_COL]): - group = group.sort_values("tstart", kind="stable") - out[tuple(int(k) for k in key)] = SampleSeries( - group["tstart"], group["tend"], group["value"].astype(np.float64) - ) - return out + pushdowns = {vid: _pushdown(rate) for vid, rate in rates.items()} + return run(spark, series_by_key, pushdowns, time_type=time_type, chunk_size=chunk_rows) def _expected(series, rate): @@ -118,7 +99,7 @@ def test_resample_matches_sample_series(spark, name): # Few samples per chunk put chunk boundaries at gaps, inside buckets and next to # zero-duration samples, with empty chunks in between. @pytest.mark.parametrize("chunk_rows", [None, 1, 2, 3]) -def test_all_cases_in_one_pass_and_other_variants_unchanged(spark, chunk_rows): +def test_all_cases_in_one_pass_and_raw_rows_unchanged(spark, chunk_rows): series = {} expected = {} for i, (tstarts, tends, values, _rate) in enumerate(_CASES.values()): @@ -126,12 +107,11 @@ def test_all_cases_in_one_pass_and_other_variants_unchanged(spark, chunk_rows): for vid, r in ((1, 1), (2, 2)): series[(i, 7, vid)] = (tstarts, tends, values) expected[(i, 7, vid)] = _expected((tstarts, tends, values), r) - # Raw rows (variant 0) and another rule's variant pass through untouched. - for vid in (0, 3): - series[(i, 7, vid)] = (tstarts, tends, values) - expected[(i, 7, vid)] = SampleSeries( - tstarts, tends, [np.nan if v is None else v for v in values] - ) + # Raw rows (variant 0) pass through untouched. + series[(i, 7, 0)] = (tstarts, tends, values) + expected[(i, 7, 0)] = SampleSeries( + tstarts, tends, [np.nan if v is None else v for v in values] + ) actual = _resample_in_spark(spark, series, {1: 1, 2: 2}, chunk_rows=chunk_rows) assert actual.keys() == expected.keys() for key, exp in expected.items(): @@ -158,35 +138,17 @@ def test_ns_epoch_longs_above_2_pow_53(spark, chunk_rows): def test_only_linear_window_frames(spark): # Spark recomputes a ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING frame for every # row (quadratic in the series length), so the rule must not use one. - schema = "container_id long, channel_id long, {} int, tstart long, tend long, value double" - df = spark.createDataFrame([(1, 1, 1, 0, 1, 1.0)], schema=schema.format(VARIANT_COL)) - result = _RULE.to_spark(df, {1: _pushdown(1)}, _CONFIG) - plan = result._jdf.queryExecution().optimizedPlan().toString() + plan = _optimized_plan(spark) assert "specifiedwindowframe(RowFrame, currentrow$(), unboundedfollowing$())" not in plan assert "specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())" in plan -def _random_series(rng, n): - """Integral samples with gaps, zero-duration samples and NaN values, in tstart order.""" - tstarts, tends, values = [], [], [] - t, after_point = 0, False - for _ in range(n): - if rng.random() < 0.15: - t += int(rng.integers(1, 6)) - duration = 0 if not after_point and rng.random() < 0.1 else int(rng.integers(1, 5)) - tstarts.append(t) - tends.append(t + duration) - values.append(np.nan if rng.random() < 0.1 else float(rng.integers(0, 100))) - t, after_point = t + duration, duration == 0 - return tstarts, tends, values - - def test_random_series_across_chunk_sizes(spark): rng = np.random.default_rng(3) rates = {1: 1, 2: 3, 3: 7} series, expected = {}, {} for cid, n in enumerate([1, 2, 40, 60, 60, 80]): - tstarts, tends, values = _random_series(rng, n) + tstarts, tends, values = random_series(rng, n) for vid, rate in rates.items(): series[(cid, 1, vid)] = (tstarts, tends, values) expected[(cid, 1, vid)] = _expected((tstarts, tends, values), rate) @@ -210,13 +172,16 @@ def _top_level_items(text): return [*items, text[start:].strip()] +def _optimized_plan(spark) -> str: + df = channel_rows(spark, {(1, 1, 1): ([0, 1], [1, 2], [1.0, 2.0])}) + result = apply_pushdowns(df, plan_for({1: _pushdown(1)}), CONFIG, include_raw=False) + return result._jdf.queryExecution().optimizedPlan().toString() + + def test_windows_over_samples_are_chunked(spark): # Spark buffers a whole window partition, so no window over the samples may be # partitioned by the series alone; only the chunk summaries are. - schema = "container_id long, channel_id long, {} int, tstart long, tend long, value double" - df = spark.createDataFrame([(1, 1, 1, 0, 1, 1.0)], schema=schema.format(VARIANT_COL)) - result = _RULE.to_spark(df, {1: _pushdown(1)}, _CONFIG) - plan = result._jdf.queryExecution().optimizedPlan().toString() + plan = _optimized_plan(spark) specs = re.findall(r"windowspecdefinition\((.*?), specifiedwindowframe", plan) series_key = ["container_id", "channel_id", VARIANT_COL] seen = set() @@ -226,18 +191,29 @@ def test_windows_over_samples_are_chunked(spark): order = [item.split(" ")[0] for item in items[len(partition) :]] seen.add((tuple(partition), tuple(order))) if partition == series_key: - assert order == ["_rs_chunk"] + assert order == ["_pd_chunk"] else: - assert partition[:3] == series_key - assert {"_rs_chunk", "_rs_block"} & set(partition) - assert any(order[-2:] == ("tstart", "tend") for _, order in seen) + assert partition == [*series_key, "_pd_chunk"] + assert any(order == ("tstart", "tend", "value") for _, order in seen) assert any(partition == tuple(series_key) for partition, _ in seen) + # Photon runs window aggregates over structs row-based. + assert not re.search(r"(max|min)\(struct\(", plan) -@pytest.mark.parametrize("chunk_rows", [0, -1, 1.5, True]) -def test_chunk_rows_must_be_a_positive_integer(chunk_rows): - with pytest.raises((TypeError, ValueError), match="chunk_rows"): - ResamplePushdownRule(chunk_rows=chunk_rows) +def test_resample_and_diff_in_one_pass(spark): + rng = np.random.default_rng(5) + series, expected = {}, {} + for cid in range(3): + samples = random_series(rng, 50) + series[(cid, 1, 1)] = series[(cid, 1, 2)] = samples + expected[(cid, 1, 1)] = _expected(samples, 3) + expected[(cid, 1, 2)] = SampleSeries(*samples).diff() + pushdowns = {1: _pushdown(3), 2: DiffRule().match(_X.diff())} + for chunk_size in (2, 7): + actual = run(spark, series, pushdowns, chunk_size=chunk_size) + assert actual.keys() == expected.keys() + for key, exp in expected.items(): + _assert_same(actual[key], exp) @pytest.mark.parametrize("chunk_rows", [None, 7]) diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py index 2a4fb035..573bc629 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py @@ -12,6 +12,7 @@ from pyspark.sql import DataFrame from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel +from impulse_query_engine.analyze.query.optimizer import ResamplePushdownRule, query_optimizer from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver from impulse_query_engine.analyze.query.solvers.solver_config import ( ChannelMappingConfig, @@ -20,9 +21,13 @@ ) from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig from tests.conftest import spark +from tests.impulse_query_engine.unit.analyze.query.optimizer.pipeline_helpers import DiffRule from tests.impulse_query_engine.unit.analyze.query.solvers.default_solver_pushdown_test import ( - CHANNEL_DATA_SCANS, + PUSHED_SCANS, + SAMPLE_SHUFFLE, channel_data_scans, + container_shuffles_of_channel_data, + shuffles, small_chunks, ) @@ -122,7 +127,35 @@ def channels(q): assert rows["channel_id"].nunique() == 2 plan = native._jdf.queryExecution().executedPlan().toString() assert "FlatMapGroupsInPandas" not in plan - assert channel_data_scans(native) == CHANNEL_DATA_SCANS + assert channel_data_scans(native) == PUSHED_SCANS + # The samples are shuffled by series and chunk, never by container. + assert len([e for e in shuffles(native) if e.startswith(SAMPLE_SHUFFLE)]) == 1 + assert container_shuffles_of_channel_data(native) == 0 + + def test_two_rules_share_the_reads(self, spark, key_value_store_db, udf_calls, monkeypatch): + monkeypatch.setattr( + query_optimizer, "default_rules", lambda: (ResamplePushdownRule(), DiffRule()) + ) + + def channels(q): + return [ + CalculatedChannel( + q.channel(channel_name="Engine RPM").resample(RATE), {"channel_name": "rpm"} + ), + CalculatedChannel( + q.channel(channel_name="Vehicle Speed Sensor").diff(), + {"channel_name": "speed_diff"}, + ), + ] + + rows, native = _solve_both( + spark, key_value_store_db, channels, udf_calls, config=_kvs_config() + ) + assert udf_calls == [1, 0] + assert rows["channel_id"].nunique() == 2 + assert rows["value"].ne(0).any() + assert channel_data_scans(native) == PUSHED_SCANS + assert len([e for e in shuffles(native) if e.startswith(SAMPLE_SHUFFLE)]) == 1 def test_same_expression_under_two_identities_and_two_rates( self, spark, key_value_store_db, udf_calls diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py index bc5c334e..20ff4e8d 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py @@ -19,7 +19,7 @@ PushdownPlan, QueryOptimizer, ResamplePushdownRule, - query_optimizer, + pipeline, ) from impulse_query_engine.analyze.query.solvers.default_solver import ( DefaultSolver, @@ -57,26 +57,59 @@ def _unit_conversion_config() -> SolverConfig: @pytest.fixture def small_chunks(monkeypatch): - """Resample in chunks of a few samples, so every channel spans many chunks.""" - monkeypatch.setattr( - query_optimizer, "default_rules", lambda: (ResamplePushdownRule(chunk_rows=7),) - ) + """Run the pushdown pipeline in chunks of a few samples, so channels span many chunks.""" + monkeypatch.setattr(pipeline, "CHUNK_ROWS", 7) + + +def _final_plan(df) -> str: + """The final physical plan of *df*, which must have been executed (e.g. collected).""" + return df._jdf.queryExecution().executedPlan().toString().split("== Initial Plan ==")[0] def channel_data_scans(df) -> list[str]: """Columns of each channels-table scan in the final plan of the executed *df*.""" - final = df._jdf.queryExecution().executedPlan().toString().split("== Initial Plan ==")[0] return sorted( re.sub(r"#\d+L?", "", columns) - for columns in re.findall(r"FileScan parquet \S+\.channels\[([^\]]*)\]", final) + for columns in re.findall(r"FileScan parquet \S+\.channels\[([^\]]*)\]", _final_plan(df)) ) -#: One full read for the rows, one narrow read shared by the resample summaries. -CHANNEL_DATA_SCANS = [ +def shuffles(df) -> list[str]: + """Partitionings of the shuffles in the final plan of the executed *df* (reuses once).""" + return [ + re.sub(r"#\d+L?", "", e) + for e in re.findall(r"(?<=- )Exchange (\w+\([^)]*\))", _final_plan(df)) + ] + + +def container_shuffles_of_channel_data(df) -> int: + """Shuffles by container alone that read the channels table, in the executed *df*.""" + lines = _final_plan(df).splitlines() + + def depth(line: str) -> int: + return len(line) - len(line.lstrip(" :+-")) + + count = 0 + for i, line in enumerate(lines): + if re.search(r"- Exchange hashpartitioning\(container_id#\d+L?, \d+\)", line): + subtree = [] + for below in lines[i + 1 :]: + if depth(below) <= depth(line): + break + subtree.append(below) + count += any(re.search(r"FileScan parquet \S+\.channels\[", b) for b in subtree) + return count + + +#: Pushed channels: one full read, and one narrow read for the series boundaries. +PUSHED_SCANS = [ "container_id,channel_id,tstart,tend", "container_id,channel_id,tstart,tend,value", ] +#: Channels also used raw add one read. +MIXED_SCANS = sorted([*PUSHED_SCANS, "container_id,channel_id,tstart,tend,value"]) +#: The one shuffle of the pushed samples: by series and chunk. +SAMPLE_SHUFFLE = f"hashpartitioning(container_id, channel_id, {VARIANT_COL}, _pd_chunk" @pytest.fixture @@ -144,17 +177,29 @@ def test_resample_only(self, spark: SparkSession, key_value_store_db, udf_rows): assert set(rows_on) == {(5, 1)} assert rows_on[(5, 1)] == _derived_rows(off, "rpm") < rows_off[(5, None)] / 5 - def test_reads_the_channel_data_once_plus_one_narrow_read( + def test_pushed_channels_are_read_twice_and_shuffled_by_chunk( self, spark: SparkSession, key_value_store_db ): + q = key_value_store_db.query + result = q.select(q.channel(channel_name="Engine RPM").resample(RATE).alias("rpm")).solve( + spark, solver=DefaultSolver(spark, config=_kvs_config(), enable_pushdown=True) + ) + assert len(result.collect()) > 0 # executes result's own plan, which the checks read + assert channel_data_scans(result) == PUSHED_SCANS + samples = [e for e in shuffles(result) if "_pd_chunk" in e] + assert len(samples) == 1 and samples[0].startswith(SAMPLE_SHUFFLE) + # The grouped-map UDF still groups the shipped rows by container. + assert container_shuffles_of_channel_data(result) == 1 + + def test_channels_used_raw_add_one_read(self, spark: SparkSession, key_value_store_db): q = key_value_store_db.query result = q.select( q.channel(channel_name="Engine RPM").resample(RATE).alias("rpm"), q.channel(channel_name="Vehicle Speed Sensor").max().alias("speed_max"), ).solve(spark, solver=DefaultSolver(spark, config=_kvs_config(), enable_pushdown=True)) - rows = result.collect() # executes result's own plan, which the check reads + rows = result.collect() assert rows and all(row["speed_max"] > 0 for row in rows) - assert channel_data_scans(result) == CHANNEL_DATA_SCANS + assert channel_data_scans(result) == MIXED_SCANS def test_raw_use_of_the_same_channel_ships_raw_rows_only( self, spark: SparkSession, key_value_store_db, udf_rows From 027bf34ce25dcaad440564ada1780bc47c7e6e35 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Thu, 15 Oct 2026 11:25:38 +0200 Subject: [PATCH 09/16] docs(query-engine): add "How a Query Is Solved" walkthrough page Add a new documentation page that walks through the five stages of a `solve()` call, from driver-side planning through metadata resolution, channel loading, Spark pushdown, and per-container pandas UDF evaluation. Include a supporting diagram and link it from the query engine index and TSAL evaluation docs. --- .../query_engine/img/solve_flow.svg | 239 ++++++++++++++++++ .../docs/references/query_engine/index.md | 3 + .../docs/references/query_engine/solving.md | 143 +++++++++++ .../query_engine/tsal/evaluation.md | 4 +- 4 files changed, 387 insertions(+), 2 deletions(-) create mode 100644 docs/impulse/docs/references/query_engine/img/solve_flow.svg create mode 100644 docs/impulse/docs/references/query_engine/solving.md diff --git a/docs/impulse/docs/references/query_engine/img/solve_flow.svg b/docs/impulse/docs/references/query_engine/img/solve_flow.svg new file mode 100644 index 00000000..86c9fec8 --- /dev/null +++ b/docs/impulse/docs/references/query_engine/img/solve_flow.svg @@ -0,0 +1,239 @@ + + How a query is solved + Five stages: plan on the driver, resolve channels from the metadata tables, load the channel data, compute pushed-down subtrees natively in Spark, and evaluate the rest of each container's expressions in a pandas UDF, which consumes the precomputed subtree results. Example: eng_rpm.where(event).histogram2d(veh_spd.where(event), bins) with event = veh_spd > 0. Spark computes the event intervals and the filtered veh_spd samples; the UDF filters the raw eng_rpm rows to the event and computes the 2D histogram. Calculated channels that are pushed down entirely skip the UDF. + + + + + + + + + + + + + + + + 1 + Plan · driver + + + query.select(e1, …, en) + .solve(spark, solver) + + + Infer result types + build each expression on an empty cache + + + Plan pushdowns (opt-in) + QueryOptimizer picks subtrees for Spark + + + + + + + + + 2 + Resolve channels · Spark, metadata tables (no sample data) + + + container_tags + container filters + + + container_metrics + project, metric filters + + + channel_tags + channel tags (EAV) + + + channel_metrics + match each selector + + + channel_mapping + aliases, unit factors + + + + + + + + channel matches: container_id · channel_id · selector_ids · conversion_factor + + + + + + 3 + Load channel data · Spark, channels table + + + channels + column mapping, filters + + + Raw mode: encode + RLE or intervals; POI union + + + Join channel matches + every row knows the selectors that read it + + + + + + + + + 4 + Compute subtrees natively · Spark pushdown (opt-in) + + + + + Split the expression at its pushed-down subtrees + eng_rpm.where(veh_spd > 0).histogram2d(veh_spd.where(veh_spd > 0), bins) + + + eng_rpm + + + ① veh_spd > 0 + + ② veh_spd.where(veh_spd > 0) + + + + + + .where(·) + + + + + + .histogram2d(…) + rest: pandas UDF + + + Rules + e.g. resample(…) or filter + conditions like veh_spd > 0 + each reads one channel; + one shared pass, + parallel per channel chunk + + + + + + + 5 + Evaluate per container · pandas UDF (GROUPED_MAP) + + + results of pushed-down subtrees (derived rows) and raw rows still needed + + + Group by container + repartition by container_id + + + TimeSeriesCache + one container's rows + + + Evaluate the rest of the expression + + eng_rpm + + ① event intervals + + ② filtered veh_spd + + + + + + .where(·) + + + + + + .histogram2d(…) + Precomputed results arrive ready-made; pandas evaluates only the rest. + + + Serialize each result + SampleSeries → binary, others → arrays or doubles + + + + + + + + + + + solve() · toPandas() + one row per container, one column per selected expression + + + solve_calculated_channels() + rows: container_id, channel_id, tstart, tend, value, identity + + + + native: every calculated channel pushed down + + + + optional table or opt-in step + + computed natively in Spark + + raw channel rows + + evaluated in pandas + + no pandas UDF + diff --git a/docs/impulse/docs/references/query_engine/index.md b/docs/impulse/docs/references/query_engine/index.md index 8156a6c4..73a1a335 100644 --- a/docs/impulse/docs/references/query_engine/index.md +++ b/docs/impulse/docs/references/query_engine/index.md @@ -12,6 +12,9 @@ silver-layer data. It has two parts: - **[Query Solvers](query_solvers.md)** — the `DefaultSolver` that knows how your silver tables are laid out, reads them, and evaluates the expressions per container. +[How a Query Is Solved](solving.md) walks through a `solve()` call stage by stage, from the +metadata tables to the per-container results. + ## Two ways to use it - **Directly, for ad-hoc analysis.** Build a query against a `MeasurementDB` and solve it diff --git a/docs/impulse/docs/references/query_engine/solving.md b/docs/impulse/docs/references/query_engine/solving.md new file mode 100644 index 00000000..1bde5ba5 --- /dev/null +++ b/docs/impulse/docs/references/query_engine/solving.md @@ -0,0 +1,143 @@ +--- +sidebar_position: 3 +title: How a Query Is Solved +--- + +# How a Query Is Solved + +One `solve()` call evaluates **every selected expression for every matching container**. The +expressions themselves are lazy ([Evaluation](tsal/evaluation.md)): nothing is computed until the +solver runs them through five stages. Query planning (stage 1) runs on the driver; query +execution (stages 2 to 5) runs in Spark. With [Spark pushdown](query_solvers.md#spark-pushdown-opt-in) +enabled, Spark computes supported subtrees natively in stage 4, and the pandas UDF in stage 5 +consumes their results instead of evaluating them. + +![How a query is solved: plan on the driver, resolve channels from the metadata tables, load the channel data, compute pushed-down subtrees natively in Spark, and evaluate the rest of each container's expressions in a pandas UDF](./img/solve_flow.svg) + +The example computes a 2D histogram of engine rpm and vehicle speed while the vehicle is +moving: + +```python +eng_rpm = db.query.channel(channel_name="Engine RPM") +veh_spd = db.query.channel(channel_name="Vehicle Speed") +event = veh_spd > 0 +result = db.query.select( + eng_rpm.where(event) + .histogram2d(veh_spd.where(event), [0, 1000, 2000, 3000, 4000], [0, 50, 100, 150, 200]) + .alias("rpm_vs_speed"), +).solve(spark, solver=DefaultSolver(spark, enable_pushdown=True)) +``` + +## 1. Plan (driver) + +- `select(...)` stores the expressions; `solve()` starts by **building each expression once + against an empty cache**. This infers every result's type, and with it the schema of the output + DataFrame, before any data is read (see + [Planning: type inference before solve](tsal/evaluation.md#planning-type-inference-before-solve)). +- With pushdown enabled, the optimizer also picks the subtrees Spark computes natively in + stage 4 (see stage 4 for the example). + +## 2. Resolve channels (metadata tables) + +The solver finds the physical channels each expression reads, using only the metadata tables. No +sample data is read in this stage. + +1. **Containers.** `container_tags` (optional) and `container_metrics` select the containers of + the configured project that match the query's container filters. Incremental report runs pass + their pre-filtered containers here. +2. **Channels.** Every selector of every expression, e.g. `channel(channel_name="Engine RPM")`, is + matched against the channel attributes on `channel_tags` (optional, EAV) or `channel_metrics`. + A selector resolves to at most one channel per container, the one with the lowest `channel_id` + (see [Channel selection](tsal/defining_expressions.md#channel-selection)). +3. **Aliases.** Selectors written with `channel_with_alias(...)` are resolved through + `channel_mapping` (optional), by alias priority. A configured `unit_conversion` table adds each + channel's conversion factor. +4. **Container metadata.** Container tags and metrics that expressions read are attached. + +The result is a small table of channel matches: `container_id`, `channel_id`, the `selector_ids` +of the selectors that read the channel, and its `conversion_factor`. + +## 3. Load channel data (Spark) + +- The `channels` table is read with its column mapping and filters, keeping only `container_id`, + `channel_id`, `tstart`, `tend` and `value`. Raw data is first encoded into RLE or interval + samples; points-in-time data from `poi_channels` is added when an expression selects POI + channels. +- The samples are joined with the channel matches from stage 2, so every row knows which + selectors read it. + +## 4. Compute subtrees natively (Spark pushdown, opt-in) + +Pushed-down subtrees are computed **in Spark, before any pandas code runs**, and only the rest of +each expression is left for the pandas UDF. The optimizer is rule-based: each rule computes one +kind of subtree on the rows of one channel, e.g. `resample(...)` or filter conditions such as +`veh_spd > 0`. All rules share one pass over the channel data, which runs in parallel per channel +chunk (see +[How Spark runs pushed-down subtrees](query_solvers.md#how-spark-runs-pushed-down-subtrees)). + +The example expression forms this tree (`event` is used twice, so it appears twice): + +``` +histogram2d(rpm_bins, speed_bins) pandas UDF +├── where pandas UDF +│ ├── eng_rpm raw rows +│ └── veh_spd > 0 ① Spark: condition → Intervals +└── veh_spd.where(veh_spd > 0) ② Spark: filter → SampleSeries +``` + +The optimizer cuts out a subtree when + +1. **a rule matches it.** It takes the largest matching subtree: ② as a whole, not just the event + inside it. +2. **it reads a single channel.** `eng_rpm.where(veh_spd > 0)` and `histogram2d` combine two + channels and stay in pandas. +3. **every use of its channel is cut out.** All three uses of `veh_spd` lie in ① or ②, so its raw + rows are no longer needed. `eng_rpm` is read by pandas directly, so its raw rows are shipped as + before. + +The results of cut-out subtrees are **derived rows** that replace the raw rows of their channel. +Here, instead of the full `veh_spd` channel, the UDF receives the event intervals ① (usually a few +rows) and the samples with `veh_spd > 0` ②. A channel can ship several derived results like this. + +:::note Rules available today +The optimizer currently implements the `resample` rule; rules for further operations such as +filter conditions follow the same flow. Until a rule covers a subtree, it is evaluated in stage 5 +as before. +::: + +## 5. Evaluate per container (pandas UDF) + +- The rows are **grouped by container**. Each container's rows reach one call of a pandas + `GROUPED_MAP` UDF, so one task holds one container's rows in memory. +- A `TimeSeriesCache` serves the container's channels, its container tags and metrics, and the + results of pushed-down subtrees. +- Each expression is **evaluated bottom-up, except for its pushed-down subtrees**: those arrive as + finished results from stage 4. In the example, the UDF clips the raw `eng_rpm` rows to the + precomputed event intervals ① and computes the 2D histogram against the precomputed `veh_spd` + samples ②. Every other leaf loads its channel from the cache as a `SampleSeries` (or + `PointsInTimeSeries` for POI channels), multiplied by the conversion factor for aliased + selectors, and each operation then works on its children's results. +- Each result is serialized into its output column: `SampleSeries` as binary, `Intervals` and + points as arrays, scalars as doubles (see + [From expression to result](tsal/evaluation.md#from-expression-to-result)). + +## Results + +| Call | Output | +|------|--------| +| `solve()` | One row per container, one column per selected expression, named by its alias. | +| `toPandas()` | The same, as a pandas DataFrame with `SampleSeries` objects deserialized. | +| `solve_calculated_channels()` | Silver-shaped rows `container_id, channel_id, tstart, tend, value` plus the channel's `identity`. | + +When **every** calculated channel of a `solve_calculated_channels()` call is pushed down as a +whole, e.g. `CalculatedChannel(veh_spd.resample(1e6), {...})`, its rows come straight from stage 4 +and no pandas UDF runs (see +[Calculated channels without the pandas UDF](query_solvers.md#calculated-channels-without-the-pandas-udf)). + +:::note Where the work happens +Stages 2 to 4 are distributed Spark jobs; stage 4 even splits each channel into chunks that run +in parallel. Stage 5 runs in parallel across containers, but each container is evaluated in a +single task: the more channels and samples a container's expressions read, the more memory that +task needs. Every pushed-down subtree moves work from stage 5 into stage 4 and shrinks the rows +that reach stage 5. +::: diff --git a/docs/impulse/docs/references/query_engine/tsal/evaluation.md b/docs/impulse/docs/references/query_engine/tsal/evaluation.md index 36becb3d..99345aeb 100644 --- a/docs/impulse/docs/references/query_engine/tsal/evaluation.md +++ b/docs/impulse/docs/references/query_engine/tsal/evaluation.md @@ -70,5 +70,5 @@ subtrees in pandas. The tree itself is unchanged, and so is every result. `toPandas()` deserializes the binary `SampleSeries` columns back into objects; the array-backed types are returned as nested lists. See the [Core Data Model](core_data_model.md) for the semantics of each -result class, and [Query Solvers](../query_solvers.md) for the full solver pipeline and the -`DefaultSolver` that reads the silver layer. +result class, [How a Query Is Solved](../solving.md) for the full flow of a `solve()` call, and +[Query Solvers](../query_solvers.md) for the `DefaultSolver` that reads the silver layer. From 4bebf1ce541a6f4ed783297eff61a99ec9c44c14 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Thu, 15 Oct 2026 11:54:34 +0200 Subject: [PATCH 10/16] docs(query-engine): align solving docs and diagram with latest solver behavior Update the "How a Query Is Solved" page and flow diagram to reflect that: - `solve_calculated_channels()` skips driver-side type inference and emits sample rows. - Alias resolution and unit conversion are now separate stages in the channel-match pipeline. - Stage 4 pushdown rules share channel-data reads, with ordered rules running per chunk. - Output encoding differs between `solve()` (serialized columns) and `solve_calculated_channels()` (rows). Also fix the `query_solvers.md` example output schema to list `identity` instead of `channel_name`/`data_key`. --- .../query_engine/img/solve_flow.svg | 43 +++++++++++-------- .../references/query_engine/query_solvers.md | 2 +- .../docs/references/query_engine/solving.md | 42 ++++++++++-------- 3 files changed, 50 insertions(+), 37 deletions(-) diff --git a/docs/impulse/docs/references/query_engine/img/solve_flow.svg b/docs/impulse/docs/references/query_engine/img/solve_flow.svg index 86c9fec8..9f6b5b79 100644 --- a/docs/impulse/docs/references/query_engine/img/solve_flow.svg +++ b/docs/impulse/docs/references/query_engine/img/solve_flow.svg @@ -52,7 +52,7 @@ Infer result types - build each expression on an empty cache + solve(): build on an empty cache Plan pushdowns (opt-in) @@ -86,7 +86,7 @@ channel_mapping - aliases, unit factors + alias resolution @@ -94,7 +94,7 @@ - channel matches: container_id · channel_id · selector_ids · conversion_factor + channel matches: container_id · channel_id · selector_ids @@ -103,21 +103,26 @@ 3 Load channel data · Spark, channels table - - channels - column mapping, filters + + channels + column mapping, filters - - Raw mode: encode - RLE or intervals; POI union + + Raw mode: encode + RLE or intervals - - Join channel matches - every row knows the selectors that read it + + Add POI channels + when POI channels are selected - - - + + Join channel matches + keep matched channels + + + + + @@ -157,8 +162,8 @@ e.g. resample(…) or filter conditions like veh_spd > 0 each reads one channel; - one shared pass, - parallel per channel chunk + all share the reads + of the channel data @@ -202,8 +207,8 @@ Precomputed results arrive ready-made; pandas evaluates only the rest. - Serialize each result - SampleSeries → binary, others → arrays or doubles + Encode each result + solve(): serialized value · calculated channels: rows diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index 099b3ee3..77b79eb5 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -157,7 +157,7 @@ cc = CalculatedChannel( {"channel_name": "speed_kmh", "data_key": "CALC"}, ) df = db.query.select(cc).solve_calculated_channels(spark, solver=DefaultSolver(spark)) -# df: [container_id, channel_id, tstart, tend, value, channel_name, data_key] +# df: [container_id, channel_id, tstart, tend, value, identity] ``` The reporting layer builds on this stage to persist derived channels to the gold layer — see diff --git a/docs/impulse/docs/references/query_engine/solving.md b/docs/impulse/docs/references/query_engine/solving.md index 1bde5ba5..d1a4276b 100644 --- a/docs/impulse/docs/references/query_engine/solving.md +++ b/docs/impulse/docs/references/query_engine/solving.md @@ -34,13 +34,16 @@ result = db.query.select( against an empty cache**. This infers every result's type, and with it the schema of the output DataFrame, before any data is read (see [Planning: type inference before solve](tsal/evaluation.md#planning-type-inference-before-solve)). + `solve_calculated_channels()` skips this step: its results are always sample rows. - With pushdown enabled, the optimizer also picks the subtrees Spark computes natively in stage 4 (see stage 4 for the example). ## 2. Resolve channels (metadata tables) The solver finds the physical channels each expression reads, using only the metadata tables. No -sample data is read in this stage. +sample data is read in this stage. It covers the first five stages of the `QuerySolver` pipeline +(container tags, container metrics, channel tags, channel metrics, alias resolution); the sixth, +`solve`, covers stages 3 to 5 here. 1. **Containers.** `container_tags` (optional) and `container_metrics` select the containers of the configured project that match the query's container filters. Incremental report runs pass @@ -49,13 +52,13 @@ sample data is read in this stage. matched against the channel attributes on `channel_tags` (optional, EAV) or `channel_metrics`. A selector resolves to at most one channel per container, the one with the lowest `channel_id` (see [Channel selection](tsal/defining_expressions.md#channel-selection)). -3. **Aliases.** Selectors written with `channel_with_alias(...)` are resolved through - `channel_mapping` (optional), by alias priority. A configured `unit_conversion` table adds each - channel's conversion factor. +3. **Aliases.** Alongside the direct selectors, selectors written with `channel_with_alias(...)` + are resolved through `channel_mapping` (optional), by alias priority. 4. **Container metadata.** Container tags and metrics that expressions read are attached. -The result is a small table of channel matches: `container_id`, `channel_id`, the `selector_ids` -of the selectors that read the channel, and its `conversion_factor`. +The result is a small table of channel matches: `container_id`, `channel_id` and the +`selector_ids` of the selectors that read the channel, plus the container metadata and, for +aliased channels, their source and target units. ## 3. Load channel data (Spark) @@ -63,16 +66,19 @@ of the selectors that read the channel, and its `conversion_factor`. `channel_id`, `tstart`, `tend` and `value`. Raw data is first encoded into RLE or interval samples; points-in-time data from `poi_channels` is added when an expression selects POI channels. -- The samples are joined with the channel matches from stage 2, so every row knows which - selectors read it. +- Conversion factors for aliased channels are looked up in `unit_conversion` (optional). +- The samples are joined with the channel matches from stage 2, which keeps only the matched + channels and attaches their selector ids, conversion factors and container metadata. With + pushdown, the rows first carry only which results they feed; the rest is attached after + stage 4. ## 4. Compute subtrees natively (Spark pushdown, opt-in) Pushed-down subtrees are computed **in Spark, before any pandas code runs**, and only the rest of each expression is left for the pandas UDF. The optimizer is rule-based: each rule computes one kind of subtree on the rows of one channel, e.g. `resample(...)` or filter conditions such as -`veh_spd > 0`. All rules share one pass over the channel data, which runs in parallel per channel -chunk (see +`veh_spd > 0`. All rules share the same reads of the channel data; rules that need a channel's +samples in order run in parallel per chunk of the channel (see [How Spark runs pushed-down subtrees](query_solvers.md#how-spark-runs-pushed-down-subtrees)). The example expression forms this tree (`event` is used twice, so it appears twice): @@ -117,9 +123,10 @@ as before. samples ②. Every other leaf loads its channel from the cache as a `SampleSeries` (or `PointsInTimeSeries` for POI channels), multiplied by the conversion factor for aliased selectors, and each operation then works on its children's results. -- Each result is serialized into its output column: `SampleSeries` as binary, `Intervals` and - points as arrays, scalars as doubles (see +- Each result is encoded for the output. `solve()` serializes it into its column: `SampleSeries` + as binary, `Intervals` and points as arrays, histograms as structs, scalars as doubles (see [From expression to result](tsal/evaluation.md#from-expression-to-result)). + `solve_calculated_channels()` emits the samples of each calculated channel as rows. ## Results @@ -135,9 +142,10 @@ and no pandas UDF runs (see [Calculated channels without the pandas UDF](query_solvers.md#calculated-channels-without-the-pandas-udf)). :::note Where the work happens -Stages 2 to 4 are distributed Spark jobs; stage 4 even splits each channel into chunks that run -in parallel. Stage 5 runs in parallel across containers, but each container is evaluated in a -single task: the more channels and samples a container's expressions read, the more memory that -task needs. Every pushed-down subtree moves work from stage 5 into stage 4 and shrinks the rows -that reach stage 5. +Stages 2 to 4 are distributed Spark jobs; for rules that need samples in order, stage 4 even +splits each channel into chunks that run in parallel. Stage 5 runs in parallel across containers, +but each container is evaluated in a single task: the more channels and samples a container's +expressions read, the more memory that task needs. Every pushed-down subtree moves work from +stage 5 into stage 4 and usually reduces the rows that reach stage 5 (resampling to a finer rate +than the channel's own adds rows). ::: From 86f43dcb050e374776e0f37a9b0f9e9003ff61d1 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Thu, 15 Oct 2026 13:26:10 +0200 Subject: [PATCH 11/16] feat(query-engine): harden pushdown pipeline against name collisions and non-op nodes, fix alias selector ids and NaN conversion factors - Enforce lowercase rule/column/run names and unique rule names to prevent pipeline column collisions; add `check_name` helper and validation in `QueryOptimizer` and `PushdownPlan`. - Change pipeline column separators to `__` so rule names cannot clash with internal columns. - Restrict pushdown matching to `TimeSeriesOp` nodes via `match_node`; aggregations, UDFs, and events now correctly keep raw rows. - Fix `TimeSeriesSelector.selector_id` for aliased non-CONTINUOUS selections so aliased and direct selectors on the same tags remain distinct. - Treat NaN unit-conversion factors like null factors, leaving values unconverted instead of nulling them. - Handle null/NaN `tstart` and all-NaN `tend` in chunking so series ordering and chunk assignment stay consistent. - Add unit tests for name validation, collision avoidance, non-op node exclusion, selector id distinctness, and NaN conversion factors. --- .../analyze/query/optimizer/plan.md | 4 ++ .../analyze/query/optimizer/pushdown.md | 20 +++++- .../metadata/time_series_expression.py | 11 ++- .../analyze/query/optimizer/context.py | 4 +- .../analyze/query/optimizer/pipeline.py | 23 ++++--- .../analyze/query/optimizer/plan.py | 27 +++++--- .../analyze/query/optimizer/pushdown.py | 30 ++++++++- .../query/optimizer/query_optimizer.py | 6 +- .../analyze/query/optimizer/traversal.py | 9 ++- .../analyze/query/solvers/default_solver.py | 6 +- .../query/optimizer/pipeline_helpers.py | 6 +- .../analyze/query/optimizer/pipeline_test.py | 67 ++++++++++++++++++- .../query/optimizer/query_optimizer_test.py | 61 +++++++++++++++++ .../default_solver_channel_resolution_test.py | 14 +++- ..._solver_native_calculated_channels_test.py | 41 ++++++++++++ 15 files changed, 288 insertions(+), 41 deletions(-) diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md index c549f9f6..87252e14 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md @@ -66,6 +66,10 @@ def variants_by_rule() -> list[tuple[PushdownRule, dict[int, Pushdown]]] Group the variants by the rule that computes them. +**Raises**: + +- `ValueError`: If two different rules share a name: their pipeline columns would +collide. #### lookup diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md index 3cef9462..a986fc8f 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -16,6 +16,19 @@ Column tagging each channel-data row with the series variant it belongs to. Variant id of the unmodified (raw) channel data. +#### check\_name + +```python +def check_name(kind: str, name: str) -> str +``` + +Return *name* if it is a valid rule, rule column or run name, else raise. + +Names are lowercase letters, digits and single underscores, starting with a +letter. Pipeline columns join them with ``__``, which a name never contains, +so names of different rules cannot collide. + + ## PushdownContext ```python @@ -86,7 +99,9 @@ Planning and decoding run in Python: :meth:`match` recognizes the subtree, #### name -Unique rule name, part of every :attr:`Pushdown.variant_key`. +Unique rule name, part of every :attr:`Pushdown.variant_key`; lowercase + +letters, digits and single underscores (see :func:`check_name`). #### ordered @@ -145,6 +160,9 @@ def row_columns(ctx: SeriesContext) -> Mapping[str, Column] Columns computed for every sample, in order; refer to them via ``ctx.col``. +Column names follow :func:`check_name` and only need to be unique within +the rule. + #### keep diff --git a/src/impulse_query_engine/analyze/metadata/time_series_expression.py b/src/impulse_query_engine/analyze/metadata/time_series_expression.py index f6c87c78..dd4d481a 100644 --- a/src/impulse_query_engine/analyze/metadata/time_series_expression.py +++ b/src/impulse_query_engine/analyze/metadata/time_series_expression.py @@ -797,13 +797,12 @@ def value_type(self) -> SeriesValueType: def selector_id(self) -> int: # series_type / value_type keep distinct-typed selections of the same tag # apart, and an aliased selection resolves differently from a direct one on - # the same tags; a direct CONTINUOUS selection uses the bare ``str(expr)`` - # hash for back-compat. + # the same tags; direct selections keep their original hashes (a CONTINUOUS + # one the bare ``str(expr)``) for back-compat. + alias = "alias|" if self._uses_alias else "" if self._series_type is SeriesType.CONTINUOUS: - if self._uses_alias: - return zlib.crc32(f"alias|{self._expr}".encode()) - return zlib.crc32(str(self._expr).encode()) - return zlib.crc32(f"{self._series_type}|{self._expr}|{self._value_type}".encode()) + return zlib.crc32(f"{alias}{self._expr}".encode()) + return zlib.crc32(f"{alias}{self._series_type}|{self._expr}|{self._value_type}".encode()) def dtype(self) -> T.DataType: """ diff --git a/src/impulse_query_engine/analyze/query/optimizer/context.py b/src/impulse_query_engine/analyze/query/optimizer/context.py index f270a3ab..60c8c9e1 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/context.py +++ b/src/impulse_query_engine/analyze/query/optimizer/context.py @@ -47,7 +47,7 @@ def run_col(run: str, name: str) -> str: """Name of the pipeline column *name* of the run *run*.""" - return f"_pd_run_{run}_{name}" + return f"_pd_run__{run}__{name}" @dataclass(frozen=True) @@ -159,7 +159,7 @@ def active(self) -> Column: def name(self, column: str) -> str: """Pipeline name of the rule's column *column*.""" - return f"_pd_{self._rule.name}_{column}" + return f"_pd_rule__{self._rule.name}__{column}" def col(self, column: str) -> Column: """The rule's column *column*, defined by an earlier hook entry.""" diff --git a/src/impulse_query_engine/analyze/query/optimizer/pipeline.py b/src/impulse_query_engine/analyze/query/optimizer/pipeline.py index ab94b687..aa186750 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/pipeline.py +++ b/src/impulse_query_engine/analyze/query/optimizer/pipeline.py @@ -63,7 +63,7 @@ run_col, ) from .plan import PushdownPlan -from .pushdown import RAW_VARIANT, VARIANT_COL, Pushdown, PushdownRule +from .pushdown import RAW_VARIANT, VARIANT_COL, Pushdown, PushdownRule, check_name #: Average number of samples per chunk of a series. CHUNK_ROWS = 100_000 @@ -143,7 +143,7 @@ def __init__(self, rules: Sequence[tuple[PushdownRule, Mapping[int, Pushdown]]], for rule, variants in self.rules: names = {} for name, breaks in rule.runs().items(): - names[name] = f"{rule.name}_{name}" + names[name] = f"{rule.name}__{check_name('run', name)}" self.runs.append(_Run(names[name], breaks, tuple(sorted(variants)))) self.rule_runs[rule.name] = names @@ -172,18 +172,21 @@ def _chunked(self, rows: DataFrame) -> DataFrame: """Series boundaries and each sample's chunk, projected for the shuffle.""" ts, te, value = self.order ts_d, te_d = F.col(ts).cast("double"), F.col(te).cast("double") + # The time range spans the finite times; a series without a finite tend is + # still cut by its tstart range. series = rows.groupBy(*self.key).agg( - F.min(ts_d).alias("_pd_first"), - F.max(_nan_free(te_d)).alias("_pd_end"), + F.min(_nan_free(ts_d)).alias("_pd_first"), + F.greatest(F.max(_nan_free(te_d)), F.max(_nan_free(ts_d))).alias("_pd_end"), F.count(F.lit(1)).alias("_pd_count"), ) first, end = F.col("_pd_first"), F.col("_pd_end") n_chunks = F.ceil(F.col("_pd_count") / F.lit(CHUNK_ROWS)) + position = F.least(F.floor((ts_d - first) / (end - first) * n_chunks), n_chunks - 1) + # Null and NaN tstarts go where the series order puts them: first and last. chunk = F.coalesce( - F.when( - end > first, - F.least(F.floor((ts_d - first) / (end - first) * n_chunks), n_chunks - 1), - ), + F.when(ts_d.isNull(), F.lit(0).cast("long")) + .when(F.isnan(ts_d), n_chunks - 1) + .when(end > first, position), F.lit(0).cast("long"), ) return rows.join(F.broadcast(series), on=self.key).select( @@ -453,12 +456,12 @@ def _row_branch( def _with_rule_columns(df: DataFrame, ctx: SeriesContext, columns: Mapping[str, Column]): """Add a rule's columns, defined for the rule's samples only.""" for name, column in columns.items(): - df = df.withColumn(ctx.name(name), F.when(ctx.active, column)) + df = df.withColumn(ctx.name(check_name("column", name)), F.when(ctx.active, column)) return df def _keep_col(rule: PushdownRule) -> str: - return f"_pd_keep_{rule.name}" + return f"_pd_keep__{rule.name}" def _keep(df: DataFrame, contexts, extra: Column) -> DataFrame: diff --git a/src/impulse_query_engine/analyze/query/optimizer/plan.py b/src/impulse_query_engine/analyze/query/optimizer/plan.py index 90df8048..6ceaf137 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/plan.py +++ b/src/impulse_query_engine/analyze/query/optimizer/plan.py @@ -12,6 +12,7 @@ from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesExpression from .pushdown import RAW_VARIANT, Pushdown, PushdownRule +from .traversal import match_node @dataclass(frozen=True, eq=False) @@ -88,23 +89,31 @@ def pushed_selector_ids(self) -> frozenset[int]: return frozenset(self.variants_by_selector) def variants_by_rule(self) -> list[tuple[PushdownRule, dict[int, Pushdown]]]: - """Group the variants by the rule that computes them.""" + """Group the variants by the rule that computes them. + + Raises + ------ + ValueError + If two different rules share a name: their pipeline columns would + collide. + """ groups: dict[str, tuple[PushdownRule, dict[int, Pushdown]]] = {} for variant, pushdown in self.pushdowns_by_variant.items(): rule = pushdown.rule - groups.setdefault(rule.name, (rule, {}))[1][variant] = pushdown + group = groups.setdefault(rule.name, (rule, {})) + if group[0] is not rule: + raise ValueError(f"two different pushdown rules are named {rule.name!r}") + group[1][variant] = pushdown return list(groups.values()) def lookup(self, node: TimeSeriesExpression) -> tuple[Pushdown, int] | None: """Return ``(pushdown, variant)`` when *node* is a pushed-down subtree.""" - for rule in self.rules: - pushdown = rule.match(node) - if pushdown is None: - continue - variant = self.variant_ids.get(pushdown.variant_key) - if variant in self.variants_by_selector.get(pushdown.selector.selector_id, ()): - return pushdown, variant + pushdown = match_node(node, self.rules) + if pushdown is None: return None + variant = self.variant_ids.get(pushdown.variant_key) + if variant in self.variants_by_selector.get(pushdown.selector.selector_id, ()): + return pushdown, variant return None def native_targets( diff --git a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py index 4da2f34e..5119acfc 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -3,6 +3,7 @@ from __future__ import annotations import abc +import re from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar @@ -30,6 +31,23 @@ #: Variant id of the unmodified (raw) channel data. RAW_VARIANT = 0 +_NAME = re.compile(r"[a-z][a-z0-9]*(_[a-z0-9]+)*") + + +def check_name(kind: str, name: str) -> str: + """Return *name* if it is a valid rule, rule column or run name, else raise. + + Names are lowercase letters, digits and single underscores, starting with a + letter. Pipeline columns join them with ``__``, which a name never contains, + so names of different rules cannot collide. + """ + if not isinstance(name, str) or not _NAME.fullmatch(name): + raise ValueError( + f"{kind} name {name!r} must consist of lowercase letters, digits and single " + "underscores, starting with a letter" + ) + return name + @dataclass(frozen=True) class PushdownContext: @@ -110,7 +128,8 @@ class PushdownRule(abc.ABC): ``(container_id, channel_id, variant, tstart, tend, value)``. """ - #: Unique rule name, part of every :attr:`Pushdown.variant_key`. + #: Unique rule name, part of every :attr:`Pushdown.variant_key`; lowercase + #: letters, digits and single underscores (see :func:`check_name`). name: ClassVar[str] #: Whether the rule needs its series in order (neighbours, runs, series #: facts). Per-sample rules set ``False`` and skip the series scan. @@ -136,12 +155,17 @@ def runs(self) -> Mapping[str, RunBreak]: """Run definitions beyond the built-in continuous blocks, by name. Each maps ``(prev, cur)`` samples to whether ``cur`` starts a new run; - :meth:`SeriesContext.run` then describes the sample's run. + :meth:`SeriesContext.run` then describes the sample's run. Names follow + :func:`check_name` and only need to be unique within the rule. """ return {} def row_columns(self, ctx: SeriesContext) -> Mapping[str, Column]: - """Columns computed for every sample, in order; refer to them via ``ctx.col``.""" + """Columns computed for every sample, in order; refer to them via ``ctx.col``. + + Column names follow :func:`check_name` and only need to be unique within + the rule. + """ return {} def keep(self, ctx: SeriesContext) -> Column: diff --git a/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py index 810153e3..fa9a2448 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py +++ b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py @@ -9,7 +9,7 @@ from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesExpression from .plan import PushdownPlan -from .pushdown import Pushdown, PushdownContext, PushdownRule +from .pushdown import Pushdown, PushdownContext, PushdownRule, check_name from .resample_rule import ResamplePushdownRule from .traversal import discover @@ -39,6 +39,10 @@ def __init__(self, rules: Sequence[PushdownRule] | None = None): self.rules: tuple[PushdownRule, ...] = ( tuple(rules) if rules is not None else default_rules() ) + names = [check_name("rule", rule.name) for rule in self.rules] + duplicates = sorted({name for name in names if names.count(name) > 1}) + if duplicates: + raise ValueError(f"pushdown rule names must be unique: {duplicates}") def optimize(self, selections: Iterable[Any], context: PushdownContext) -> PushdownPlan: """Build the pushdown plan for the *selections* of one solve call. diff --git a/src/impulse_query_engine/analyze/query/optimizer/traversal.py b/src/impulse_query_engine/analyze/query/optimizer/traversal.py index cf6d2f72..e4d1cde7 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/traversal.py +++ b/src/impulse_query_engine/analyze/query/optimizer/traversal.py @@ -45,7 +45,14 @@ def _expressions_in(values: Iterable) -> Iterator[TimeSeriesExpression]: def match_node(node: TimeSeriesExpression, rules: Sequence[PushdownRule]) -> Pushdown | None: - """Return the first rule's pushdown for *node*, or ``None``.""" + """Return the first rule's pushdown for *node*, or ``None``. + + Only plain :class:`TimeSeriesOp` nodes can be pushed down: their ``build`` + serves precomputed results from the cache, other node types (UDFs, + aggregations, events, …) would evaluate on missing raw rows instead. + """ + if type(node) is not TimeSeriesOp: + return None for rule in rules: pushdown = rule.match(node) if pushdown is not None: diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index c56a856c..a5734459 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -1821,9 +1821,9 @@ def _native_calculated_channel_rows( value = F.col(val).cast("double") if has_conversion: factor = F.col(conv) - value = F.when(F.col("_cc_uses_alias") & factor.isNotNull(), value * factor).otherwise( - value - ) + # Like TimeSeriesCache.to_sample_series: a null or NaN factor leaves values as they are. + applies = F.col("_cc_uses_alias") & factor.isNotNull() & ~F.isnan(factor) + value = F.when(applies, value * factor).otherwise(value) def as_udf_output(col: Column, dtype: T.DataType) -> Column: # The UDF emits float64 arrays; Arrow casts them back and masks NaN as null. diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py index 3d753cbb..f27b5907 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py @@ -51,7 +51,11 @@ def channel_rows(spark, series_by_key, time_type=T.LongType()) -> DataFrame: T.StructField(VARIANT_COL, T.IntegerType()), ] ) - to_time = int if isinstance(time_type, T.IntegralType) else float + cast = int if isinstance(time_type, T.IntegralType) else float + + def to_time(t): + return None if t is None else cast(t) + rows = [ (cid, ch, to_time(ts), to_time(te), None if v is None else float(v), vid) for (cid, ch, vid), (tstarts, tends, values) in series_by_key.items() diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py index 327b85a9..32e7f7bb 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py @@ -15,7 +15,12 @@ from impulse_query_engine.analyze.metadata.tag_expression import TagSelector from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector -from impulse_query_engine.analyze.query.optimizer import VARIANT_COL, Pushdown, apply_pushdowns +from impulse_query_engine.analyze.query.optimizer import ( + VARIANT_COL, + Pushdown, + SeriesPipeline, + apply_pushdowns, +) from impulse_query_engine.model.series.sample_series import SampleSeries from tests.conftest import spark from tests.impulse_query_engine.unit.analyze.query.optimizer.pipeline_helpers import ( @@ -24,6 +29,7 @@ OffsetRule, ProbeRule, channel_rows, + chunk_rows, plan_for, random_series, run, @@ -159,9 +165,11 @@ def test_diff_matches_sample_series(spark): nptest.assert_array_equal(actual[key].values, exp.values) +_OFFSET = OffsetRule() + + def _offset(value: float) -> Pushdown: - rule = OffsetRule() - return Pushdown(rule=rule, selector=_X, params=(("offset", value),), consumed=(_X,)) + return Pushdown(rule=_OFFSET, selector=_X, params=(("offset", value),), consumed=(_X,)) def test_per_sample_rules_share_the_row_branch_with_raw_rows(spark): @@ -233,3 +241,56 @@ def test_one_shuffle_of_the_samples(spark): assert all(e.startswith(f"hashpartitioning({key}") for e in exchanges) # One full read, shared by the chunk pass and its summaries, and the boundaries. assert len(re.findall(r"Scan ExistingRDD", final)) == 2 + + +def _values_by_tend(df) -> dict: + return {row["tend"]: row["value"] for row in df.collect()} + + +@pytest.mark.parametrize("chunk_size", [1, 2, None]) +def test_null_and_nan_tstarts_keep_their_place_in_the_series(spark, chunk_size): + # Spark orders a null tstart first and a NaN tstart last; their chunks must too. + nan = float("nan") + samples = ([None, 0.0, 1.0, 2.0, 3.0, 4.0, nan], [0.5, 1, 2, 3, 4, 5, 5.5], [1.0] * 7) + df = channel_rows(spark, {(1, 1, 1): samples, (1, 1, 2): samples}, time_type=T.DoubleType()) + prev = ProbeRule("prev_start", lambda ctx: ctx.prev.tstart).pushdown(_X) + nxt = ProbeRule("next_start", lambda ctx: ctx.next.tstart).pushdown(_X) + with chunk_rows(chunk_size): + result = apply_pushdowns(df, plan_for({1: prev, 2: nxt}), CONFIG, include_raw=False) + by_variant = {v: _values_by_tend(result.where(F.col(VARIANT_COL) == v)) for v in (1, 2)} + expected_prev = {0.5: None, 1.0: None, 2.0: 0.0, 3.0: 1.0, 4.0: 2.0, 5.0: 3.0, 5.5: 4.0} + expected_next = {0.5: 0.0, 1.0: 1.0, 2.0: 2.0, 3.0: 3.0, 4.0: 4.0, 5.0: nan, 5.5: None} + for actual, expected in ((by_variant[1], expected_prev), (by_variant[2], expected_next)): + assert actual.keys() == expected.keys() + nptest.assert_array_equal(_as_float(actual.values()), _as_float(expected.values())) + + +def test_series_without_a_finite_tend_is_still_chunked(spark): + nan = float("nan") + samples = ([0.0, 1.0, 2.0, 3.0, 4.0, 5.0], [nan] * 6, [1.0] * 6) + df = channel_rows(spark, {(1, 1, 1): samples}, time_type=T.DoubleType()) + rule = DiffRule() + with chunk_rows(2): + chunks = SeriesPipeline([(rule, {1: rule.match(_X.diff())})], CONFIG)._chunked(df) + assert chunks.select("_pd_chunk").distinct().count() == 3 + + +def test_rule_names_cannot_collide_with_pipeline_columns(spark): + # The rule column "ts" of a rule named "prev" once overwrote the pipeline's prev tstart. + rule = ProbeRule("prev", lambda ctx: ctx.prev.tstart) + rule.row_columns = lambda ctx: {"ts": ctx.value * 100} + samples = ([0, 1, 2], [1, 2, 3], [1.0, 2.0, 3.0]) + actual = run(spark, {(1, 1, 1): samples}, {1: rule.pushdown(_X)}, chunk_size=1) + nptest.assert_array_equal(actual[(1, 1, 1)].values, [np.nan, 0.0, 1.0]) + + +@pytest.mark.parametrize("hook", ["row_columns", "runs"]) +def test_invalid_rule_column_and_run_names_are_rejected(spark, hook): + rule = ProbeRule("bad_names", lambda ctx: ctx.value) + if hook == "row_columns": + rule.row_columns = lambda ctx: {"Bad__name": ctx.value} + else: + rule._runs = {"Bad__name": _value_change} + df = channel_rows(spark, {(1, 1, 1): ([0], [1], [1.0])}) + with pytest.raises(ValueError, match="name 'Bad__name'"): + apply_pushdowns(df, plan_for({1: rule.pushdown(_X)}), CONFIG) diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py index a50b2f89..2445ca40 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py @@ -21,10 +21,13 @@ SequenceOfEventsExpression, ) from impulse_query_engine.analyze.query.optimizer import ( + Pushdown, PushdownContext, + PushdownPlan, QueryOptimizer, ResamplePushdownRule, ) +from impulse_query_engine.analyze.query.optimizer.traversal import discover from impulse_query_engine.model.series.intervals import Intervals _LONG = PushdownContext(T.LongType(), T.LongType()) @@ -130,6 +133,64 @@ def test_walk_disagreeing_with_get_selectors_keeps_raw_data(self, rpm): assert _plan(over_counted, hidden_raw_use).is_empty +class _MatchesAnyNode(ResamplePushdownRule): + """Claims every non-selector node that reads exactly one selector, of any type.""" + + name = "any_node" + + def match(self, node): + if isinstance(node, TimeSeriesSelector) or not isinstance(node, TimeSeriesExpression): + return None + selectors = node.get_selectors() + if len(selectors) != 1: + return None + return Pushdown(rule=self, selector=selectors[0], params=(), consumed=tuple(selectors)) + + def supports(self, pushdown, context): + return True + + +class TestOnlyOperationsArePushedDown: + # Only TimeSeriesOp.build serves precomputed results; other node types must + # keep their raw rows. + def test_aggregation_node_is_not_a_candidate(self, rpm): + found, leaves = discover(HistogramDuration(rpm, bins=[0, 1000]), [_MatchesAnyNode()]) + assert found == [] + assert [s.selector_id for s in leaves] == [rpm.selector_id] + + def test_udf_node_is_not_a_candidate_but_its_arguments_are(self, rpm): + found, _ = discover(rpm.resample(1e6).apply(lambda s: s), [_MatchesAnyNode()]) + assert [node.operation for node, _ in found] == ["resample"] + + def test_lookup_ignores_other_node_types(self, rpm): + rule = _MatchesAnyNode() + plan = QueryOptimizer([rule]).optimize([rpm.resample(1e6)], _LONG) + assert plan.lookup(rpm.resample(1e6)) is not None + assert plan.lookup(HistogramDuration(rpm, bins=[0, 1000])) is None + + +class TestRuleNames: + @pytest.mark.parametrize("name", ["Bad", "a__b", "_x", "x_", "1x", ""]) + def test_invalid_rule_names_are_rejected(self, name): + rule = type("_Named", (ResamplePushdownRule,), {"name": name})() + with pytest.raises(ValueError, match="rule name"): + QueryOptimizer([rule]) + + def test_rule_names_must_be_unique(self): + with pytest.raises(ValueError, match="unique"): + QueryOptimizer([ResamplePushdownRule(), ResamplePushdownRule()]) + + def test_plan_rejects_two_rules_sharing_a_name(self, rpm, spd): + first, second = ResamplePushdownRule(), ResamplePushdownRule() + plan = PushdownPlan.build( + (first, second), + [first.match(rpm.resample(1e6)), second.match(spd.resample(2e6))], + [rpm.selector_id, spd.selector_id], + ) + with pytest.raises(ValueError, match="two different pushdown rules"): + plan.variants_by_rule() + + class _Opaque(TimeSeriesExpression): """Node whose attributes disagree with the selectors it reports.""" diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py index d2be8d8a..258ad11a 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py @@ -14,7 +14,10 @@ import pytest from impulse_query_engine.analyze.metadata.tag_expression import TagSelector -from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.metadata.time_series_expression import ( + SeriesType, + TimeSeriesSelector, +) from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver from impulse_query_engine.analyze.query.solvers.solver_config import ( ChannelMappingConfig, @@ -181,3 +184,12 @@ def test_aliased_and_direct_selectors_on_the_same_tags_have_distinct_ids(): aliased = TimeSeriesSelector(expr, uses_alias=True) assert direct.selector_id == zlib.crc32(str(expr).encode()) assert aliased.selector_id != direct.selector_id + + +def test_aliased_and_direct_poi_selectors_on_the_same_tags_have_distinct_ids(): + expr = TagSelector("channel_name") == "Fault Code" + direct = TimeSeriesSelector(expr, series_type=SeriesType.POINTS_IN_TIME) + aliased = TimeSeriesSelector(expr, uses_alias=True, series_type=SeriesType.POINTS_IN_TIME) + expected_direct = f"{SeriesType.POINTS_IN_TIME}|{expr}|{direct._value_type}" + assert direct.selector_id == zlib.crc32(expected_direct.encode()) + assert aliased.selector_id != direct.selector_id diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py index 573bc629..30b2ff60 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py @@ -211,6 +211,47 @@ def channels(q): by_channel = rows.groupby("channel_id")["value"].sum() assert len(by_channel) == 2 and by_channel.iloc[0] != by_channel.iloc[1] + def test_nan_conversion_factor_leaves_values_unconverted( + self, spark, unit_conversion_dataframes, mock_workspace_client, udf_calls + ): + # vehicle_speed converts km/h to m/s; a NaN km/h factor must be skipped like + # the pandas evaluation does, not turn every value into null. + tables = dict(unit_conversion_dataframes) + tables["unit_conversion"] = tables["unit_conversion"].withColumn( + "conversion_factor", + F.when(F.col("unit") == "km/h", F.lit(float("nan"))).otherwise( + F.col("conversion_factor").cast("double") + ), + ) + cfg = MeasurementDBConfig.for_debug(tables) + cfg.channel_mapping_table = "channel_mapping" + cfg.unit_conversion_table = "unit_conversion" + db = MeasurementDB(cfg, ws=mock_workspace_client) + + def channels(q): + return [ + CalculatedChannel( + q.channel_with_alias(channel_alias="vehicle_speed").resample(RATE), + {"channel_name": "speed_aliased"}, + ), + CalculatedChannel( + q.channel(channel_name="Vehicle Speed Sensor").resample(RATE), + {"channel_name": "speed_raw"}, + ), + ] + + rows, _ = _solve_both(spark, db, channels, udf_calls, config=_unit_conversion_config()) + assert udf_calls == [1, 0] + aliased, raw = ( + rows[rows["identity"] == (("channel_name", name),)] + for name in ("speed_aliased", "speed_raw") + ) + same_samples = aliased.merge(raw, on=["container_id", "tstart", "tend"]) + assert len(same_samples) > 0 and same_samples["value_x"].notna().all() + pd.testing.assert_series_equal( + same_samples["value_x"], same_samples["value_y"], check_names=False + ) + def test_raw_data(self, spark, setup_raw_channels_db, mock_workspace_client, udf_calls): tables = { name: spark.read.table(f"spark_catalog.silver_raw.{name}") From 67cbced0019906889b3dc07a58ec982a66b88ed6 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Thu, 15 Oct 2026 15:54:53 +0200 Subject: [PATCH 12/16] feat(query-engine): classify pushdowns bottom-up and support chained operands Introduce `Classifier` to memoize bottom-up classification of `TimeSeriesExpression` nodes into pushdowns, replacing the top-down `match_node` helper. Rules now receive an `Operand` callback that exposes whether each operand is a raw selector, a pushed subexpression, or evaluated in pandas, enabling rules to absorb or chain pushed operands. - Add `Classifier` and `Operand` to the optimizer public API and docs. - Update `PushdownRule.match` signature to accept `operand`. - Add `Pushdown.source` for pushdowns that compute on another pushdown's result; reject planned chained pushdowns in `PushdownPlan.build`. - Update `QueryOptimizer` to accept only single-selector, non-chained pushdowns and share one `Classifier` across discovery. - Update `ResamplePushdownRule` and test helpers to the new `match` interface. - Cache one `Classifier` per container in `TimeSeriesCache` so UDF node lookups classify each node once. - Add unit tests for bottom-up classification, chained/absorbed operands, unsupported parent fallback, and shared-node memoization. --- .gitignore | 5 +- .../analyze/query/optimizer/plan.md | 11 +- .../analyze/query/optimizer/pushdown.md | 41 +++++- .../query/optimizer/query_optimizer.md | 4 + .../analyze/query/optimizer/__init__.py | 12 +- .../analyze/query/optimizer/plan.py | 27 +++- .../analyze/query/optimizer/pushdown.py | 39 +++-- .../query/optimizer/query_optimizer.py | 29 ++-- .../analyze/query/optimizer/resample_rule.py | 9 +- .../analyze/query/optimizer/traversal.py | 92 +++++++++--- .../analyze/query/solvers/default_solver.py | 5 +- .../query/optimizer/pipeline_helpers.py | 15 +- .../analyze/query/optimizer/pipeline_test.py | 11 +- .../query/optimizer/query_optimizer_test.py | 133 ++++++++++++++++-- .../query/optimizer/resample_rule_test.py | 10 +- .../solvers/default_solver_pushdown_test.py | 19 ++- 16 files changed, 385 insertions(+), 77 deletions(-) diff --git a/.gitignore b/.gitignore index e9222661..033db734 100644 --- a/.gitignore +++ b/.gitignore @@ -85,4 +85,7 @@ yarn-error.log* # Claude Code .claude/ .isaac/ -CLAUDE.md \ No newline at end of file +CLAUDE.md + +# Local design docs +/docs/design/ \ No newline at end of file diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md index 87252e14..12a0adcf 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md @@ -40,6 +40,11 @@ Assign variant ids to *pushdowns*, in discovery order. - `pushdowns` (`Iterable[Pushdown]`): Accepted pushdowns; each one's selector is pushed down entirely. - `selector_ids` (`Iterable[int]`): Ids of all selectors in the solve. +**Raises**: + +- `ValueError`: If a pushdown computes on another pushdown's result (``source``), +which the pipeline cannot compute. + #### is\_empty ```python @@ -74,11 +79,15 @@ collide. #### lookup ```python -def lookup(node: TimeSeriesExpression) -> tuple[Pushdown, int] | None +def lookup(node: TimeSeriesExpression, + classify: Classifier | None = None) -> tuple[Pushdown, int] | None ``` Return ``(pushdown, variant)`` when *node* is a pushed-down subtree. +Pass a *classify* shared by the lookups of one evaluation so that every +node is classified once. + #### native\_targets diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md index a986fc8f..7a905ff3 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -61,12 +61,18 @@ A subtree that a :class:`PushdownRule` computes natively in Spark. Expression nodes overload ``==``, so pushdowns compare by identity; use :attr:`variant_key` to identify equivalent pushdowns. +A pushdown reads the samples of its channel, possibly absorbing pushed +operands into that pass, or computes on the result of another pushdown, its +:attr:`source`. The pipeline computes the former only. + **Arguments**: - `rule` (`PushdownRule`): The rule that matched the subtree. - `selector` (`TimeSeriesSelector`): The physical channel selector the subtree reads. - `params` (`tuple of (str, Any)`): Canonical, hashable operation parameters (e.g. ``(("sample_rate", 1e6),)``). - `consumed` (`tuple of TimeSeriesSelector`): Every selector occurrence inside the subtree (its ``get_selectors()``). +- `source` (`Pushdown`): The pushdown whose result this one computes on; ``None`` when it reads +the channel's samples. #### variant\_key @@ -86,6 +92,14 @@ def describe() -> str Human-readable ``rule(param=value, ...)`` label. +#### Operand + +Classification of an operand of the node a rule matches: the selector of a + +channel, the :class:`Pushdown` computing the operand, or ``None`` when the +operand is evaluated in pandas or is not an expression. + + ## PushdownRule ```python @@ -94,7 +108,24 @@ class PushdownRule(abc.ABC) Recognizes a subtree and computes it in Spark ahead of the pandas UDF. -Planning and decoding run in Python: :meth:`match` recognizes the subtree, +Planning and decoding run in Python: :meth:`match` recognizes a node from +its operands' classification, :meth:`decode` turns one +container/channel/variant slice of the derived rows back into the object the +subtree evaluates to. The result must equal evaluating the subtree in pandas +on the same channel data. + +The Spark side is a set of hooks returning columns, evaluated by the shared +pushdown pipeline (:mod:`.pipeline`) for the samples of the rule's variants, +one series (channel and variant) at a time, with a :class:`SeriesContext`: + +1. :meth:`row_columns` per sample; +2. :meth:`keep` selects the samples that go on; +3. :meth:`kept_columns` per kept sample; +4. :meth:`emit` the derived rows of a kept sample. + +Rules never handle DataFrames, so every rule shares the pipeline's reads and +shuffles. Derived rows have the channel-data shape +``(container_id, channel_id, variant, tstart, tend, value)``. #### name @@ -114,11 +145,17 @@ facts). Per-sample rules set ``False`` and skip the series scan. #### match ```python -def match(node: TimeSeriesExpression) -> Pushdown | None +def match(node: TimeSeriesExpression, operand: Operand) -> Pushdown | None ``` Return the pushdown for *node*, or ``None`` if the rule does not apply. +Nodes are classified bottom-up: *operand* returns the classification of +any operand of *node*. A pushdown reading the channel's samples sets no +``source``; when it absorbs a pushed operand into that pass, its +``consumed`` includes the operand's and :meth:`supports` covers both. A +pushdown computing on a pushed operand's result sets ``source`` to it. + Purely syntactic: it runs on the driver while planning and again inside the UDF to recognize the node, so it must not depend on data or schemas. diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md index aceaecb0..31d21612 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md @@ -23,6 +23,10 @@ class QueryOptimizer() Plans which subtrees of a solve's selections are computed natively in Spark. +Nodes are classified bottom-up (:class:`~.traversal.Classifier`), and the +topmost classified nodes whose pushdown the channel data supports are the +candidates; a node that is not supported leaves its operands as candidates. + All selections of a solve share the raw data of their physical channels, so a subtree is pushed down only when **every** occurrence of its input selector, across all selections, sits inside a pushed-down subtree. Then the raw rows diff --git a/src/impulse_query_engine/analyze/query/optimizer/__init__.py b/src/impulse_query_engine/analyze/query/optimizer/__init__.py index a3a218a3..75824ad8 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/__init__.py +++ b/src/impulse_query_engine/analyze/query/optimizer/__init__.py @@ -1,14 +1,24 @@ from .context import ELEMENT_TYPE, Run, RunBreak, Sample, SeriesContext from .pipeline import SeriesPipeline, apply_pushdowns from .plan import PushdownPlan -from .pushdown import RAW_VARIANT, VARIANT_COL, Pushdown, PushdownContext, PushdownRule +from .pushdown import ( + RAW_VARIANT, + VARIANT_COL, + Operand, + Pushdown, + PushdownContext, + PushdownRule, +) from .query_optimizer import QueryOptimizer, default_rules from .resample_rule import ResamplePushdownRule +from .traversal import Classifier __all__ = [ "ELEMENT_TYPE", "RAW_VARIANT", "VARIANT_COL", + "Classifier", + "Operand", "Pushdown", "PushdownContext", "PushdownPlan", diff --git a/src/impulse_query_engine/analyze/query/optimizer/plan.py b/src/impulse_query_engine/analyze/query/optimizer/plan.py index 6ceaf137..f7181ac9 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/plan.py +++ b/src/impulse_query_engine/analyze/query/optimizer/plan.py @@ -12,7 +12,7 @@ from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesExpression from .pushdown import RAW_VARIANT, Pushdown, PushdownRule -from .traversal import match_node +from .traversal import Classifier @dataclass(frozen=True, eq=False) @@ -58,11 +58,19 @@ def build( Accepted pushdowns; each one's selector is pushed down entirely. selector_ids : Iterable[int] Ids of all selectors in the solve. + + Raises + ------ + ValueError + If a pushdown computes on another pushdown's result (``source``), + which the pipeline cannot compute. """ variant_ids: dict[tuple, int] = {} pushdowns_by_variant: dict[int, Pushdown] = {} variants_by_selector: dict[int, list[int]] = {} for pushdown in pushdowns: + if pushdown.source is not None: + raise ValueError(f"pushdown {pushdown.describe()} has a source") key = pushdown.variant_key if key not in variant_ids: variant_ids[key] = len(variant_ids) + 1 @@ -106,9 +114,17 @@ def variants_by_rule(self) -> list[tuple[PushdownRule, dict[int, Pushdown]]]: group[1][variant] = pushdown return list(groups.values()) - def lookup(self, node: TimeSeriesExpression) -> tuple[Pushdown, int] | None: - """Return ``(pushdown, variant)`` when *node* is a pushed-down subtree.""" - pushdown = match_node(node, self.rules) + def lookup( + self, node: TimeSeriesExpression, classify: Classifier | None = None + ) -> tuple[Pushdown, int] | None: + """Return ``(pushdown, variant)`` when *node* is a pushed-down subtree. + + Pass a *classify* shared by the lookups of one evaluation so that every + node is classified once. + """ + if classify is None: + classify = Classifier(self.rules) + pushdown = classify(node) if pushdown is None: return None variant = self.variant_ids.get(pushdown.variant_key) @@ -143,10 +159,11 @@ def native_targets( if not selections: return None targets: dict[int, dict[int, list[tuple[int, bool]]]] = {} + classify = Classifier(self.rules) for selection in selections: if not isinstance(selection, CalculatedChannel): return None - hit = self.lookup(selection.expr) + hit = self.lookup(selection.expr, classify) if hit is None or not hit[0].rule.emits_result_rows(hit[0]): return None pushdown, variant = hit diff --git a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py index 5119acfc..337a0d35 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/pushdown.py +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -4,7 +4,7 @@ import abc import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar @@ -78,6 +78,10 @@ class Pushdown: Expression nodes overload ``==``, so pushdowns compare by identity; use :attr:`variant_key` to identify equivalent pushdowns. + A pushdown reads the samples of its channel, possibly absorbing pushed + operands into that pass, or computes on the result of another pushdown, its + :attr:`source`. The pipeline computes the former only. + Parameters ---------- rule : PushdownRule @@ -88,31 +92,44 @@ class Pushdown: Canonical, hashable operation parameters (e.g. ``(("sample_rate", 1e6),)``). consumed : tuple of TimeSeriesSelector Every selector occurrence inside the subtree (its ``get_selectors()``). + source : Pushdown, optional + The pushdown whose result this one computes on; ``None`` when it reads + the channel's samples. """ rule: PushdownRule selector: TimeSeriesSelector params: tuple[tuple[str, Any], ...] consumed: tuple[TimeSeriesSelector, ...] + source: Pushdown | None = None @property def variant_key(self) -> tuple: """Identity of the derived series, independent of the input channel.""" - return (self.rule.name, self.params) + key = (self.rule.name, self.params) + return key if self.source is None else (*key, self.source.variant_key) def describe(self) -> str: """Human-readable ``rule(param=value, ...)`` label.""" params = ", ".join(f"{k}={v}" for k, v in self.params) - return f"{self.rule.name}({params})" + label = f"{self.rule.name}({params})" + return label if self.source is None else f"{label} on {self.source.describe()}" + + +#: Classification of an operand of the node a rule matches: the selector of a +#: channel, the :class:`Pushdown` computing the operand, or ``None`` when the +#: operand is evaluated in pandas or is not an expression. +Operand = Callable[[Any], "TimeSeriesSelector | Pushdown | None"] class PushdownRule(abc.ABC): """Recognizes a subtree and computes it in Spark ahead of the pandas UDF. - Planning and decoding run in Python: :meth:`match` recognizes the subtree, - :meth:`decode` turns one container/channel/variant slice of the derived rows - back into the object the subtree evaluates to. The result must equal - evaluating the subtree in pandas on the same channel data. + Planning and decoding run in Python: :meth:`match` recognizes a node from + its operands' classification, :meth:`decode` turns one + container/channel/variant slice of the derived rows back into the object the + subtree evaluates to. The result must equal evaluating the subtree in pandas + on the same channel data. The Spark side is a set of hooks returning columns, evaluated by the shared pushdown pipeline (:mod:`.pipeline`) for the samples of the rule's variants, @@ -136,9 +153,15 @@ class PushdownRule(abc.ABC): ordered: ClassVar[bool] = True @abc.abstractmethod - def match(self, node: TimeSeriesExpression) -> Pushdown | None: + def match(self, node: TimeSeriesExpression, operand: Operand) -> Pushdown | None: """Return the pushdown for *node*, or ``None`` if the rule does not apply. + Nodes are classified bottom-up: *operand* returns the classification of + any operand of *node*. A pushdown reading the channel's samples sets no + ``source``; when it absorbs a pushed operand into that pass, its + ``consumed`` includes the operand's and :meth:`supports` covers both. A + pushdown computing on a pushed operand's result sets ``source`` to it. + Purely syntactic: it runs on the driver while planning and again inside the UDF to recognize the node, so it must not depend on data or schemas. """ diff --git a/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py index fa9a2448..009d837d 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py +++ b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py @@ -4,6 +4,7 @@ from collections import Counter from collections.abc import Iterable, Sequence +from functools import partial from typing import Any from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesExpression @@ -11,7 +12,7 @@ from .plan import PushdownPlan from .pushdown import Pushdown, PushdownContext, PushdownRule, check_name from .resample_rule import ResamplePushdownRule -from .traversal import discover +from .traversal import Classifier, discover def default_rules() -> tuple[PushdownRule, ...]: @@ -22,6 +23,10 @@ def default_rules() -> tuple[PushdownRule, ...]: class QueryOptimizer: """Plans which subtrees of a solve's selections are computed natively in Spark. + Nodes are classified bottom-up (:class:`~.traversal.Classifier`), and the + topmost classified nodes whose pushdown the channel data supports are the + candidates; a node that is not supported leaves its operands as candidates. + All selections of a solve share the raw data of their physical channels, so a subtree is pushed down only when **every** occurrence of its input selector, across all selections, sits inside a pushed-down subtree. Then the raw rows @@ -63,12 +68,14 @@ def optimize(self, selections: Iterable[Any], context: PushdownContext) -> Pushd if not all(isinstance(s, TimeSeriesExpression) for s in selections): return PushdownPlan(rules=self.rules) + classify = Classifier(self.rules) + accept = partial(_accepted, context=context) total: Counter[int] = Counter() candidates: list[tuple[TimeSeriesExpression, Pushdown]] = [] for selection in selections: expected = Counter(s.selector_id for s in selection.get_selectors()) total.update(expected) - found, leaves = discover(selection, self.rules) + found, leaves = discover(selection, classify, accept) walked = Counter(s.selector_id for s in leaves) walked.update(s.selector_id for _, p in found for s in p.consumed) # A walk that disagrees with get_selectors() missed or double-counted part @@ -76,12 +83,8 @@ def optimize(self, selections: Iterable[Any], context: PushdownContext) -> Pushd if walked == expected: candidates.extend(found) - candidates = [ - (node, pushdown) - for node, pushdown in candidates - if pushdown.rule.supports(pushdown, context) and _has_result_type(node, pushdown) - ] - + # Candidates read a single selector, so the operands of a dropped candidate + # read one whose raw rows are shipped anyway and need not be reconsidered. while True: consumed = Counter(s.selector_id for _, p in candidates for s in p.consumed) fully_pushed = {sid for sid, count in consumed.items() if count == total[sid]} @@ -97,6 +100,16 @@ def optimize(self, selections: Iterable[Any], context: PushdownContext) -> Pushd return PushdownPlan.build(self.rules, [p for _, p in candidates], total) +def _accepted(node: TimeSeriesExpression, pushdown: Pushdown, context: PushdownContext) -> bool: + """Whether the pipeline can compute *pushdown* for *node* on the channel data.""" + if pushdown.source is not None: + # Computing on another pushdown's result needs a pipeline stage per rule. + return False + if any(s.selector_id != pushdown.selector.selector_id for s in pushdown.consumed): + return False + return pushdown.rule.supports(pushdown, context) and _has_result_type(node, pushdown) + + def _has_result_type(node: TimeSeriesExpression, pushdown: Pushdown) -> bool: """Whether *node* evaluates to the type the rule decodes into.""" try: diff --git a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py index 0ad2cf6f..d71f1020 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -19,7 +19,7 @@ ) from .context import SeriesContext -from .pushdown import Pushdown, PushdownContext, PushdownRule +from .pushdown import Operand, Pushdown, PushdownContext, PushdownRule _RATE = "sample_rate" _DEFAULT_RATE = 1.0 @@ -49,17 +49,18 @@ class ResamplePushdownRule(PushdownRule): name = "resample" - def match(self, node: TimeSeriesExpression) -> Pushdown | None: + def match(self, node: TimeSeriesExpression, operand: Operand) -> Pushdown | None: if ( type(node) is not TimeSeriesOp or node.optype != "cls" or not isinstance(node.operation, str) or node.operation != "resample" + or not node.args ): return None - if not node.args or not isinstance(node.args[0], TimeSeriesSelector): + selector = operand(node.args[0]) + if not isinstance(selector, TimeSeriesSelector): return None - selector = node.args[0] if selector.series_type is not SeriesType.CONTINUOUS: return None rate = _sample_rate(node.args[1:], node.kwargs) diff --git a/src/impulse_query_engine/analyze/query/optimizer/traversal.py b/src/impulse_query_engine/analyze/query/optimizer/traversal.py index e4d1cde7..9db5d358 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/traversal.py +++ b/src/impulse_query_engine/analyze/query/optimizer/traversal.py @@ -1,8 +1,9 @@ -"""Generic walk over TSAL expression trees.""" +"""Bottom-up classification of TSAL expression trees into pushdowns.""" from __future__ import annotations -from collections.abc import Iterable, Iterator, Sequence +from collections.abc import Callable, Iterable, Iterator, Sequence +from typing import Any from impulse_query_engine.analyze.metadata.time_series_expression import ( TimeSeriesAliasSelector, @@ -13,6 +14,9 @@ from .pushdown import Pushdown, PushdownRule +#: Plan-time check whether a classified node's pushdown is planned. +Accept = Callable[[TimeSeriesExpression, Pushdown], bool] + def child_expressions(node: TimeSeriesExpression) -> Iterator[TimeSeriesExpression]: """Yield the direct sub-expressions of *node*. @@ -44,34 +48,82 @@ def _expressions_in(values: Iterable) -> Iterator[TimeSeriesExpression]: yield from _expressions_in(value.values()) -def match_node(node: TimeSeriesExpression, rules: Sequence[PushdownRule]) -> Pushdown | None: - """Return the first rule's pushdown for *node*, or ``None``. +class Classifier: + """Classifies expression nodes bottom-up into the pushdowns computing them. + + A node is classified after its operands: the first rule whose + :meth:`~.pushdown.PushdownRule.match` returns a pushdown for the node, given + its operands' classification, computes it. Only plain :class:`TimeSeriesOp` + nodes can be pushed down: their ``build`` serves precomputed results from the + cache, other node types (UDFs, aggregations, events, …) would evaluate on + missing raw rows instead. - Only plain :class:`TimeSeriesOp` nodes can be pushed down: their ``build`` - serves precomputed results from the cache, other node types (UDFs, - aggregations, events, …) would evaluate on missing raw rows instead. + Classification is syntactic, so the driver and the UDF classify alike. + Results are memoized per node object: use one classifier per optimization or + per container, and never store it on a plan or a node. + + Parameters + ---------- + rules : Sequence[PushdownRule] + Rules in priority order. """ - if type(node) is not TimeSeriesOp: + + def __init__(self, rules: Sequence[PushdownRule]): + self.rules = tuple(rules) + # Nodes are not hashable; holding each node keeps its id from being reused. + self._memo: dict[int, tuple[TimeSeriesOp, Pushdown | None]] = {} + + def __call__(self, node: TimeSeriesExpression) -> Pushdown | None: + """The pushdown computing *node*, or ``None``.""" + if type(node) is not TimeSeriesOp: + return None + stack: list[tuple[TimeSeriesOp, bool]] = [(node, False)] + while stack: + current, operands_classified = stack.pop() + if id(current) in self._memo: + continue + if operands_classified: + self._memo[id(current)] = (current, self._match(current)) + else: + stack.append((current, True)) + stack.extend( + (child, False) + for child in child_expressions(current) + if type(child) is TimeSeriesOp + ) + return self._memo[id(node)][1] + + def operand(self, value: Any) -> TimeSeriesSelector | Pushdown | None: + """Classification of an operand: its selector for a channel, else its pushdown.""" + if isinstance(value, TimeSeriesSelector): + return value + if isinstance(value, TimeSeriesExpression): + return self(value) + return None + + def _match(self, node: TimeSeriesOp) -> Pushdown | None: + for rule in self.rules: + pushdown = rule.match(node, self.operand) + if pushdown is not None: + return pushdown return None - for rule in rules: - pushdown = rule.match(node) - if pushdown is not None: - return pushdown - return None def discover( - root: TimeSeriesExpression, rules: Sequence[PushdownRule] + root: TimeSeriesExpression, classify: Classifier, accept: Accept | None = None ) -> tuple[list[tuple[TimeSeriesExpression, Pushdown]], list[TimeSeriesSelector]]: - """Find the maximal subtrees of *root* that a rule can push down. + """Find the topmost subtrees of *root* that are pushed down. - The walk does not descend into a matched subtree and keeps no visited set, so - a node referenced twice is counted twice, like ``get_selectors()`` does. + A classified node is a candidate when *accept* (by default every pushdown) + accepts its pushdown; otherwise the walk continues into its operands, whose + pushdowns can still be candidates. The walk does not descend into a + candidate and keeps no visited set, so a node referenced twice is counted + twice, like ``get_selectors()`` does. Returns ------- tuple - ``(candidates, leaves)``: the matched ``(node, pushdown)`` pairs and the + ``(candidates, leaves)``: the ``(node, pushdown)`` candidates and the selectors reached outside of them. """ candidates: list[tuple[TimeSeriesExpression, Pushdown]] = [] @@ -79,8 +131,8 @@ def discover( stack: list[TimeSeriesExpression] = [root] while stack: node = stack.pop() - pushdown = match_node(node, rules) - if pushdown is not None: + pushdown = classify(node) + if pushdown is not None and (accept is None or accept(node, pushdown)): candidates.append((node, pushdown)) elif isinstance(node, TimeSeriesSelector): leaves.append(node) diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index a5734459..254ce019 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -20,6 +20,7 @@ from impulse_query_engine.analyze.query.optimizer import ( RAW_VARIANT, VARIANT_COL, + Classifier, PushdownContext, PushdownPlan, QueryOptimizer, @@ -96,6 +97,8 @@ def __init__( self._container_tag_cols = list(container_tag_cols or []) self._container_metric_cols = list(container_metric_cols or []) self._plan = plan + # Classifies each node once for all lookups of this container. + self._classify = Classifier(plan.rules) if plan is not None else None has_variants = VARIANT_COL in pdf.columns # *pdf* holds channel data for a whole container, so avoid creating unnecessary copies of the data. @@ -225,7 +228,7 @@ def materialized(self, expression): """ if self._plan is None: return None - hit = self._plan.lookup(expression) + hit = self._plan.lookup(expression, self._classify) if hit is None: return None pushdown, variant = hit diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py index f27b5907..3aca63a4 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py @@ -116,7 +116,7 @@ def random_series(rng, n, nan_tend=False): return tstarts, tends, values -def _selector_op(node, operation: str) -> TimeSeriesSelector | None: +def _selector_op(node, operand, operation: str) -> TimeSeriesSelector | None: """The continuous selector of ``selector.()``, or ``None``.""" if ( type(node) is not TimeSeriesOp @@ -124,11 +124,12 @@ def _selector_op(node, operation: str) -> TimeSeriesSelector | None: or node.operation != operation or len(node.args) != 1 or node.kwargs - or not isinstance(node.args[0], TimeSeriesSelector) - or node.args[0].series_type is not SeriesType.CONTINUOUS ): return None - return node.args[0] + selector = operand(node.args[0]) + if not isinstance(selector, TimeSeriesSelector): + return None + return selector if selector.series_type is SeriesType.CONTINUOUS else None class DiffRule(PushdownRule): @@ -136,8 +137,8 @@ class DiffRule(PushdownRule): name = "diff" - def match(self, node): - selector = _selector_op(node, "diff") + def match(self, node, operand): + selector = _selector_op(node, operand, "diff") if selector is None: return None return Pushdown(rule=self, selector=selector, params=(), consumed=(selector,)) @@ -155,7 +156,7 @@ class OffsetRule(PushdownRule): name = "offset" ordered = False - def match(self, node): + def match(self, node, operand): return None def emit(self, ctx: SeriesContext) -> Column: diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py index 32e7f7bb..21849de8 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py @@ -17,6 +17,7 @@ from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector from impulse_query_engine.analyze.query.optimizer import ( VARIANT_COL, + Classifier, Pushdown, SeriesPipeline, apply_pushdowns, @@ -156,7 +157,7 @@ def test_diff_matches_sample_series(spark): samples = random_series(rng, n) series[(cid, 1, 1)] = samples expected[(cid, 1, 1)] = SampleSeries(*samples).diff() - pushdowns = {1: DiffRule().match(_X.diff())} + pushdowns = {1: Classifier([DiffRule()])(_X.diff())} for chunk_size in (1, 3, None): actual = run(spark, series, pushdowns, chunk_size=chunk_size) for key, exp in expected.items(): @@ -175,7 +176,7 @@ def _offset(value: float) -> Pushdown: def test_per_sample_rules_share_the_row_branch_with_raw_rows(spark): samples = ([0, 1, 5], [1, 2, 6], [1.0, 2.0, 4.0]) series = {(1, 1, 0): samples, (1, 1, 1): samples, (1, 1, 2): samples, (2, 1, 3): samples} - pushdowns = {1: _offset(10.0), 2: _offset(0.5), 3: DiffRule().match(_X.diff())} + pushdowns = {1: _offset(10.0), 2: _offset(0.5), 3: Classifier([DiffRule()])(_X.diff())} actual = run(spark, series, pushdowns, chunk_size=2) nptest.assert_array_equal(actual[(1, 1, 0)].values, [1.0, 2.0, 4.0]) nptest.assert_array_equal(actual[(1, 1, 1)].values, [11.0, 12.0, 14.0]) @@ -196,7 +197,7 @@ def test_per_sample_rules_skip_the_series_scan(spark): def test_raw_rows_only_when_asked(spark): samples = ([0, 1], [1, 2], [1.0, 2.0]) series = {(1, 1, 0): samples, (1, 1, 1): samples} - pushdowns = {1: DiffRule().match(_X.diff())} + pushdowns = {1: Classifier([DiffRule()])(_X.diff())} assert set(run(spark, series, pushdowns)) == {(1, 1, 0), (1, 1, 1)} assert set(run(spark, series, pushdowns, include_raw=False)) == {(1, 1, 1)} @@ -226,7 +227,7 @@ def test_one_shuffle_of_the_samples(spark): rng = np.random.default_rng(3) series = {(cid, 1, 1): random_series(rng, 30) for cid in range(3)} df = channel_rows(spark, series) - plan = plan_for({1: DiffRule().match(_X.diff())}) + plan = plan_for({1: Classifier([DiffRule()])(_X.diff())}) result = apply_pushdowns(df, plan, CONFIG, include_raw=False) result.collect() final = result._jdf.queryExecution().executedPlan().toString().split("== Initial Plan ==")[0] @@ -271,7 +272,7 @@ def test_series_without_a_finite_tend_is_still_chunked(spark): df = channel_rows(spark, {(1, 1, 1): samples}, time_type=T.DoubleType()) rule = DiffRule() with chunk_rows(2): - chunks = SeriesPipeline([(rule, {1: rule.match(_X.diff())})], CONFIG)._chunked(df) + chunks = SeriesPipeline([(rule, {1: Classifier([rule])(_X.diff())})], CONFIG)._chunked(df) assert chunks.select("_pd_chunk").distinct().count() == 3 diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py index 2445ca40..79db1458 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py @@ -12,6 +12,7 @@ from impulse_query_engine.analyze.metadata.time_series_expression import ( SeriesType, TimeSeriesExpression, + TimeSeriesOp, TimeSeriesSelector, ) from impulse_query_engine.analyze.query.aggregations.histogram import HistogramDuration @@ -27,7 +28,7 @@ QueryOptimizer, ResamplePushdownRule, ) -from impulse_query_engine.analyze.query.optimizer.traversal import discover +from impulse_query_engine.analyze.query.optimizer.traversal import Classifier, discover from impulse_query_engine.model.series.intervals import Intervals _LONG = PushdownContext(T.LongType(), T.LongType()) @@ -138,7 +139,7 @@ class _MatchesAnyNode(ResamplePushdownRule): name = "any_node" - def match(self, node): + def match(self, node, operand): if isinstance(node, TimeSeriesSelector) or not isinstance(node, TimeSeriesExpression): return None selectors = node.get_selectors() @@ -154,12 +155,14 @@ class TestOnlyOperationsArePushedDown: # Only TimeSeriesOp.build serves precomputed results; other node types must # keep their raw rows. def test_aggregation_node_is_not_a_candidate(self, rpm): - found, leaves = discover(HistogramDuration(rpm, bins=[0, 1000]), [_MatchesAnyNode()]) + found, leaves = discover( + HistogramDuration(rpm, bins=[0, 1000]), Classifier([_MatchesAnyNode()]) + ) assert found == [] assert [s.selector_id for s in leaves] == [rpm.selector_id] def test_udf_node_is_not_a_candidate_but_its_arguments_are(self, rpm): - found, _ = discover(rpm.resample(1e6).apply(lambda s: s), [_MatchesAnyNode()]) + found, _ = discover(rpm.resample(1e6).apply(lambda s: s), Classifier([_MatchesAnyNode()])) assert [node.operation for node, _ in found] == ["resample"] def test_lookup_ignores_other_node_types(self, rpm): @@ -169,6 +172,116 @@ def test_lookup_ignores_other_node_types(self, rpm): assert plan.lookup(HistogramDuration(rpm, bins=[0, 1000])) is None +class _OnPushedOperand(ResamplePushdownRule): + """Matches ``x.cumtrapz()`` on a pushed ``x``, recording the operands it sees.""" + + def __init__(self, name="on_pushed", chained=False, supported=True): + self.name = name + self.chained = chained + self.supported = supported + self.seen = [] + + def match(self, node, operand): + if type(node) is not TimeSeriesOp or node.operation != "cumtrapz": + return None + inner = operand(node.args[0]) + self.seen.append(inner) + if not isinstance(inner, Pushdown): + return None + return Pushdown( + rule=self, + selector=inner.selector, + params=(("inner", inner.variant_key),), + consumed=inner.consumed, + source=inner if self.chained else None, + ) + + def supports(self, pushdown, context): + return self.supported + + +_RESAMPLE = "resample(sample_rate=1000000.0)" + + +def _described(plan): + return [p.describe() for p in plan.pushdowns_by_variant.values()] + + +class TestBottomUpClassification: + def test_rules_see_their_operands_classification(self, rpm): + rule = _OnPushedOperand() + classify = Classifier([rule, ResamplePushdownRule()]) + for node in (rpm.cumtrapz(), rpm.resample(1e6).cumtrapz(), (rpm * 2).cumtrapz()): + classify(node) + channel, pushed, evaluated_in_pandas = rule.seen + assert channel is rpm + assert pushed.describe() == _RESAMPLE + assert evaluated_in_pandas is None + + def test_chained_pushdown_is_not_planned_but_its_operand_is(self, rpm): + expr = rpm.resample(1e6).cumtrapz() + plan = _plan(expr, rules=[_OnPushedOperand(chained=True), ResamplePushdownRule()]) + assert _described(plan) == [_RESAMPLE] + assert plan.pushed_selector_ids == {rpm.selector_id} + assert plan.lookup(expr) is None + assert plan.lookup(expr.args[0])[1] == 1 + + def test_absorbing_rule_is_planned_as_one_variant(self, rpm): + expr = rpm.resample(1e6).cumtrapz() + plan = _plan(expr, rules=[_OnPushedOperand(), ResamplePushdownRule()]) + (pushdown,) = plan.pushdowns_by_variant.values() + assert pushdown.rule.name == "on_pushed" and pushdown.source is None + assert pushdown.consumed == (rpm,) + hit, variant = plan.lookup(expr) + assert hit.variant_key == pushdown.variant_key and variant == 1 + assert plan.lookup(expr.args[0]) is None + + def test_unsupported_parent_falls_back_to_its_pushed_operand(self, rpm): + expr = rpm.resample(1e6).cumtrapz() + plan = _plan(expr, rules=[_OnPushedOperand(supported=False), ResamplePushdownRule()]) + assert _described(plan) == [_RESAMPLE] + assert plan.lookup(expr) is None + assert plan.lookup(expr.args[0])[1] == 1 + + def test_first_matching_rule_decides(self, rpm): + # The UDF classifies without knowing what is supported, so an unsupported + # first rule must not hand the node to the next one. + expr = rpm.resample(1e6).cumtrapz() + rules = [ + _OnPushedOperand("unsupported", supported=False), + _OnPushedOperand("supported"), + ResamplePushdownRule(), + ] + plan = _plan(expr, rules=rules) + assert _described(plan) == [_RESAMPLE] + assert plan.lookup(expr) is None + assert plan.lookup(expr.args[0])[1] == 1 + + def test_shared_nodes_are_matched_once(self, rpm): + matched = [] + + class _Recording(ResamplePushdownRule): + def match(self, node, operand): + matched.append(node) + return super().match(node, operand) + + resampled = rpm.resample(1e6) + plan = _plan(resampled + resampled, resampled.cumtrapz(), rules=[_Recording()]) + assert len(matched) == len({id(node) for node in matched}) == 3 + assert plan.pushed_selector_ids == {rpm.selector_id} + + def test_chained_pushdowns_are_their_own_variant_and_cannot_be_planned(self, rpm): + inner = Classifier([ResamplePushdownRule()])(rpm.resample(1e6)) + chained = Pushdown( + rule=inner.rule, selector=rpm, params=inner.params, consumed=(rpm,), source=inner + ) + assert inner.variant_key == ("resample", (("sample_rate", 1e6),)) + assert chained.variant_key != inner.variant_key + assert chained.describe() == f"{_RESAMPLE} on {_RESAMPLE}" + with pytest.raises(ValueError, match="has a source"): + PushdownPlan.build([inner.rule], [chained], [rpm.selector_id]) + + class TestRuleNames: @pytest.mark.parametrize("name", ["Bad", "a__b", "_x", "x_", "1x", ""]) def test_invalid_rule_names_are_rejected(self, name): @@ -184,7 +297,7 @@ def test_plan_rejects_two_rules_sharing_a_name(self, rpm, spd): first, second = ResamplePushdownRule(), ResamplePushdownRule() plan = PushdownPlan.build( (first, second), - [first.match(rpm.resample(1e6)), second.match(spd.resample(2e6))], + [Classifier([first])(rpm.resample(1e6)), Classifier([second])(spd.resample(2e6))], [rpm.selector_id, spd.selector_id], ) with pytest.raises(ValueError, match="two different pushdown rules"): @@ -236,16 +349,16 @@ class TestResampleRule: ], ) def test_unsupported_calls_are_not_matched(self, rpm, make): - assert ResamplePushdownRule().match(make(rpm)) is None + assert Classifier([ResamplePushdownRule()])(make(rpm)) is None def test_default_rate(self, rpm): - pushdown = ResamplePushdownRule().match(rpm.resample()) + pushdown = Classifier([ResamplePushdownRule()])(rpm.resample()) assert pushdown.params == (("sample_rate", 1.0),) assert pushdown.selector is rpm def test_points_in_time_selector_is_not_matched(self): poi = _channel("dtc", series_type=SeriesType.POINTS_IN_TIME) - assert ResamplePushdownRule().match(poi.resample(1e6)) is None + assert Classifier([ResamplePushdownRule()])(poi.resample(1e6)) is None def test_fractional_rate_needs_double_time_columns(self, rpm): assert _plan(rpm.resample(0.5), context=_LONG).is_empty @@ -268,8 +381,8 @@ def result_type(self, pushdown): def test_failing_type_inference_drops_candidate(self, rpm): class _Exploding(ResamplePushdownRule): - def match(self, node): - pushdown = super().match(node) + def match(self, node, operand): + pushdown = super().match(node, operand) if pushdown is not None: node.evaluation_type = _raise return pushdown diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py index c69e38ca..68091c29 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -10,7 +10,11 @@ from impulse_query_engine.analyze.metadata.tag_expression import TagSelector from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector -from impulse_query_engine.analyze.query.optimizer import VARIANT_COL, ResamplePushdownRule +from impulse_query_engine.analyze.query.optimizer import ( + VARIANT_COL, + Classifier, + ResamplePushdownRule, +) from impulse_query_engine.model.series.sample_series import SampleSeries from tests.conftest import spark from tests.impulse_query_engine.unit.analyze.query.optimizer.pipeline_helpers import ( @@ -28,7 +32,7 @@ def _pushdown(rate): - return _RULE.match(_X.resample(rate)) + return Classifier([_RULE])(_X.resample(rate)) def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType(), chunk_rows=None): @@ -208,7 +212,7 @@ def test_resample_and_diff_in_one_pass(spark): series[(cid, 1, 1)] = series[(cid, 1, 2)] = samples expected[(cid, 1, 1)] = _expected(samples, 3) expected[(cid, 1, 2)] = SampleSeries(*samples).diff() - pushdowns = {1: _pushdown(3), 2: DiffRule().match(_X.diff())} + pushdowns = {1: _pushdown(3), 2: Classifier([DiffRule()])(_X.diff())} for chunk_size in (2, 7): actual = run(spark, series, pushdowns, chunk_size=chunk_size) assert actual.keys() == expected.keys() diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py index 20ff4e8d..ecfa563e 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py @@ -16,6 +16,7 @@ from impulse_query_engine.analyze.query.optimizer import ( RAW_VARIANT, VARIANT_COL, + Classifier, PushdownPlan, QueryOptimizer, ResamplePushdownRule, @@ -327,7 +328,7 @@ class TestTimeSeriesCacheVariants: @staticmethod def _cache(rule, node) -> TimeSeriesCache: - pushdown = rule.match(node) + pushdown = Classifier([rule])(node) plan = PushdownPlan.build([rule], [pushdown], [pushdown.selector.selector_id]) pdf = pd.DataFrame( { @@ -384,6 +385,22 @@ def decode(self, rows, pushdown, cache): assert isinstance(result, Intervals) nptest.assert_array_equal(result.tstarts, [0, 10]) + def test_each_node_is_classified_once_per_container(self): + matched = [] + + class _Recording(ResamplePushdownRule): + def match(self, node, operand): + matched.append(node) + return super().match(node, operand) + + resampled = _rpm().resample(10) + cache = self._cache(_Recording(), resampled) + matched.clear() + integral = resampled.cumtrapz().build(cache) + # The root and the served resample, each matched once across both lookups. + assert len(matched) == 2 and len({id(node) for node in matched}) == 2 + nptest.assert_array_equal(integral.tstarts, [0, 10]) + def test_optimizer_default_rules(self): assert [type(r) for r in QueryOptimizer().rules] == [ResamplePushdownRule] From bd87944aa356dde4b1aa17c4e83daa607926c3e9 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Fri, 16 Oct 2026 09:53:11 +0200 Subject: [PATCH 13/16] feat(query-engine): allow direct selectors to match multiple channels per container Remove the "lowest channel_id wins" deduplication for direct selectors in `DefaultSolver` so a selector can resolve to every matching channel per container. Aliased selectors still resolve to one channel per container (best alias priority, then lowest channel_id). Update `TimeSeriesSelector.selector_id` to no longer encode `uses_alias`, since direct and aliased selections are now handled separately. Adjust docs, generated API reference, and tests to reflect the new resolution behavior. --- .../analyze/query/solvers/default_solver.md | 4 - .../docs/references/query_engine/solving.md | 2 +- .../query_engine/tsal/defining_expressions.md | 8 +- .../metadata/time_series_expression.py | 12 +-- .../analyze/query/solvers/default_solver.py | 48 ++--------- .../default_solver_channel_resolution_test.py | 81 +++++++------------ 6 files changed, 47 insertions(+), 108 deletions(-) diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md index 3eba53ee..463e0f3c 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md @@ -172,10 +172,6 @@ frame; the selectors are applied directly to ``channel_metrics`` columns, the result is restricted to the candidate containers, and ``selector_ids`` is computed from the matching selectors. -In both modes a selector resolves to at most one channel per container: -when several channels of a container match, the lowest ``channel_id`` -is kept. - **Arguments**: - `spark` (`SparkSession`): Spark session used for query execution. diff --git a/docs/impulse/docs/references/query_engine/solving.md b/docs/impulse/docs/references/query_engine/solving.md index d1a4276b..aa6728fe 100644 --- a/docs/impulse/docs/references/query_engine/solving.md +++ b/docs/impulse/docs/references/query_engine/solving.md @@ -50,7 +50,7 @@ sample data is read in this stage. It covers the first five stages of the `Query their pre-filtered containers here. 2. **Channels.** Every selector of every expression, e.g. `channel(channel_name="Engine RPM")`, is matched against the channel attributes on `channel_tags` (optional, EAV) or `channel_metrics`. - A selector resolves to at most one channel per container, the one with the lowest `channel_id` + A selector should match at most one channel per container (see [Channel selection](tsal/defining_expressions.md#channel-selection)). 3. **Aliases.** Alongside the direct selectors, selectors written with `channel_with_alias(...)` are resolved through `channel_mapping` (optional), by alias priority. diff --git a/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md b/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md index aa4650c1..3caa0566 100644 --- a/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md +++ b/docs/impulse/docs/references/query_engine/tsal/defining_expressions.md @@ -32,11 +32,9 @@ table is configured — see [Query Solvers](../query_solvers.md#how-defaultsolve ::: :::note When a selection matches several channels -A selection resolves to at most one channel per container. If several channels of a container -match its tags (e.g. the same `channel_name` under two `data_key`s), the one with the lowest -`channel_id` is used; add tags such as `data_key` to select a specific channel. Earlier versions -picked an arbitrary match, so recompute incremental reports in full if their selections were -ambiguous. +A selection should identify at most one channel per container. If several channels of a container +match its tags (e.g. the same `channel_name` under two `data_key`s), one of them is used, and which +one is unspecified; add tags such as `data_key` to select a specific channel. ::: ### Points-in-Time (POI) channels diff --git a/src/impulse_query_engine/analyze/metadata/time_series_expression.py b/src/impulse_query_engine/analyze/metadata/time_series_expression.py index dd4d481a..8b335011 100644 --- a/src/impulse_query_engine/analyze/metadata/time_series_expression.py +++ b/src/impulse_query_engine/analyze/metadata/time_series_expression.py @@ -796,13 +796,10 @@ def value_type(self) -> SeriesValueType: @property def selector_id(self) -> int: # series_type / value_type keep distinct-typed selections of the same tag - # apart, and an aliased selection resolves differently from a direct one on - # the same tags; direct selections keep their original hashes (a CONTINUOUS - # one the bare ``str(expr)``) for back-compat. - alias = "alias|" if self._uses_alias else "" + # apart; CONTINUOUS uses the bare ``str(expr)`` hash for back-compat. if self._series_type is SeriesType.CONTINUOUS: - return zlib.crc32(f"{alias}{self._expr}".encode()) - return zlib.crc32(f"{alias}{self._series_type}|{self._expr}|{self._value_type}".encode()) + return zlib.crc32(str(self._expr).encode()) + return zlib.crc32(f"{self._series_type}|{self._expr}|{self._value_type}".encode()) def dtype(self) -> T.DataType: """ @@ -871,8 +868,7 @@ def build(self, cache: SeriesCache) -> SampleSeries | PointsInTimeSeries: if self._series_type is SeriesType.POINTS_IN_TIME: return PointsInTimeSeries.empty(value_type=self._value_type) return SampleSeries.empty() - # DefaultSolver resolves a selector to at most one channel per container; - # other solvers may pass several candidates, of which the first is used. + # TODO: select candidate mid = candidates.container_id.iloc[0] cid = candidates.channel_id.iloc[0] # The selector is the source of truth for the series type: pass it to the diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index 254ce019..d4d90ce9 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -622,10 +622,6 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: columns, the result is restricted to the candidate containers, and ``selector_ids`` is computed from the matching selectors. - In both modes a selector resolves to at most one channel per container: - when several channels of a container match, the lowest ``channel_id`` - is kept. - Parameters ---------- spark : SparkSession @@ -662,7 +658,10 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: on=[container_id_col, channel_id_col], how="inner", ) - return self._one_channel_per_selector(metrics) + metrics = metrics.withColumn("selector_ids", F.array(F.col("selector_id"))).drop( + "selector_id" + ) + return metrics.select(container_id_col, channel_id_col, "selector_ids") # Wide mode: selectors apply directly to channel_metrics columns. channel_metrics_df = channel_metrics_df.where(self._build_expr(selectors)) @@ -671,28 +670,10 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: on=[container_id_col], how="inner", ) - result = result.withColumn("selector_id", self._build_selector_id_expr(selectors)) - return self._one_channel_per_selector(result) - - def _one_channel_per_selector(self, matches: DataFrame) -> DataFrame: - """Keep the lowest matching ``channel_id`` per container and direct selector. - - *matches* carries one ``selector_id`` per ``(container_id, channel_id)`` - row (the first selector matching the channel). Returns - ``(container_id, channel_id, selector_ids)`` with one row per container - and selector. - """ - container_id_col = self.config.container_id_col - channel_id_col = self.config.channel_id_col - return ( - matches.groupBy(container_id_col, "selector_id") - .agg(F.min(channel_id_col).alias(channel_id_col)) - .select( - container_id_col, - channel_id_col, - F.array(F.col("selector_id")).alias("selector_ids"), - ) + result = result.withColumn( + "selector_ids", F.array(self._build_selector_id_expr(selectors)) ) + return result.select(container_id_col, channel_id_col, "selector_ids") def filter_aliased_channel_metrics( self, spark, db: MeasurementDB, container_df, selectors @@ -909,20 +890,7 @@ def resolve_channel_selections( ] if has_unit_cols: aliased_solve_cols.extend([source_unit_col, target_unit_col]) - # An aliased selector matching several aliases of a container resolves to - # one channel, like a direct selector: best alias priority, then lowest - # channel_id. - per_selector = Window.partitionBy( - self.config.container_id_col, F.col("selector_ids")[0] - ).orderBy( - F.col(self.config.alias_priority_col).asc_nulls_last(), - F.col(self.config.channel_id_col).asc(), - ) - aliased_for_union = ( - aliased_channel_metrics_df.withColumn("_rank", F.row_number().over(per_selector)) - .where(F.col("_rank") == 1) - .select(*aliased_solve_cols) - ) + aliased_for_union = aliased_channel_metrics_df.select(*aliased_solve_cols) merged = channel_metrics_df.unionByName( aliased_for_union, allowMissingColumns=has_unit_cols diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py index 258ad11a..f6e8ac22 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py @@ -1,12 +1,11 @@ # pylint: disable=missing-function-docstring -"""DefaultSolver resolves every selector to at most one channel per container. +"""DefaultSolver channel resolution when several channels of a container match a selection. -When several channels of a container match a selection, the lowest ``channel_id`` -wins (after the alias priority for aliased selections), independent of row order. +A direct selection keeps every matching channel. An alias resolves to one channel per +container: the best alias priority, then the lowest ``channel_id``. """ import os -import zlib import numpy as np import pandas as pd @@ -14,10 +13,7 @@ import pytest from impulse_query_engine.analyze.metadata.tag_expression import TagSelector -from impulse_query_engine.analyze.metadata.time_series_expression import ( - SeriesType, - TimeSeriesSelector, -) +from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver from impulse_query_engine.analyze.query.solvers.solver_config import ( ChannelMappingConfig, @@ -28,6 +24,9 @@ from tests.conftest import spark _ALIAS_SCHEMA = "spark_catalog.silver_key_value_store_alias" +# unit="C" matches the ambient and intake temperatures: channels 6 and 9 in containers 1 +# and 2, channels 2 and 8 in container 3. +_CELSIUS_CHANNELS = {1: [6, 9], 2: [6, 9], 3: [2, 8]} def _alias_config() -> SolverConfig: @@ -85,14 +84,12 @@ def _with_channel(channel_metrics, container_id, channel_id, **columns): class TestDirectSelectors: - def test_wide_mode_keeps_lowest_channel_id(self, spark, basic_narrow_db): - # unit="C" matches the ambient and intake temperatures: channels 6 and 9 in - # containers 1 and 2, channels 2 and 8 in container 3. + def test_wide_mode_keeps_every_matching_channel(self, spark, basic_narrow_db): q = basic_narrow_db.query celsius = q.channel(unit="C") - assert _resolved(spark, DefaultSolver(spark), q, celsius) == {1: [6], 2: [6], 3: [2]} + assert _resolved(spark, DefaultSolver(spark), q, celsius) == _CELSIUS_CHANNELS - def test_wide_mode_solves_on_the_resolved_channel(self, spark, basic_narrow_db): + def test_wide_mode_solves_on_one_of_the_matching_channels(self, spark, basic_narrow_db): q = basic_narrow_db.query result = q.select(q.channel(unit="C").max().alias("t_max")).solve( spark, solver=DefaultSolver(spark) @@ -102,18 +99,20 @@ def test_wide_mode_solves_on_the_resolved_channel(self, spark, basic_narrow_db): base_path = os.path.dirname(os.path.abspath(__file__)) base_path = base_path[: base_path.find("tests")] data = pd.read_csv(f"{base_path}/tests/unit/data/basic_narrow_csv/channel_data.csv") - expected = { - cid: np.nanmax(data[(data.container_id == cid) & (data.channel_id == ch)].value) - for cid, ch in {1: 6, 2: 6, 3: 2}.items() - } - assert actual == pytest.approx(expected) - assert all(np.isfinite(v) for v in actual.values()) - - def test_eav_mode_keeps_lowest_channel_id( + assert actual.keys() == _CELSIUS_CHANNELS.keys() + for cid, channels in _CELSIUS_CHANNELS.items(): + maxima = [ + np.nanmax(data[(data.container_id == cid) & (data.channel_id == ch)].value) + for ch in channels + ] + assert np.isfinite(actual[cid]) + assert any(actual[cid] == pytest.approx(m) for m in maxima) + + def test_eav_mode_keeps_every_matching_channel( self, spark, setup_key_value_store_alias_db, mock_workspace_client ): # Re-tag channel 9 of container 1 and channels 2 and 8 of container 3 as - # "Engine RPM": container 1 keeps 5 (the original), container 3 keeps 2. + # "Engine RPM", next to channel 5 of containers 1 and 2. tables = _alias_tables(spark) tags = tables["channel_tags"] retagged = [(1, 9), (3, 2), (3, 8)] @@ -130,9 +129,9 @@ def test_eav_mode_keeps_lowest_channel_id( q = db.query rpm = q.channel(channel_name="Engine RPM") assert _resolved(spark, DefaultSolver(spark, config=_alias_config()), q, rpm) == { - 1: [5], + 1: [5, 9], 2: [5], - 3: [2], + 3: [2, 8], } @@ -157,39 +156,21 @@ def test_priority_tie_keeps_lowest_channel_id( solver = DefaultSolver(spark, config=_alias_config()) assert _resolved(spark, solver, q, engine_speed)[1] == [expected] - def test_selector_matching_several_aliases_resolves_one_channel( + def test_selector_matching_several_aliases_resolves_each_alias( self, spark, key_value_store_alias_db ): q = key_value_store_alias_db.query solver = DefaultSolver(spark, config=_alias_config()) - per_alias = [ - _resolved(spark, solver, q, q.channel_with_alias(channel_alias=alias)) - for alias in ("engine_speed", "vehicle_speed") - ] + expected: dict = {} + for alias in ("engine_speed", "vehicle_speed"): + one = _resolved(spark, solver, q, q.channel_with_alias(channel_alias=alias)) + for cid, channels in one.items(): + expected.setdefault(cid, set()).update(channels) both = TimeSeriesSelector( (TagSelector("channel_alias") == "engine_speed") | (TagSelector("channel_alias") == "vehicle_speed"), uses_alias=True, ) resolved = _resolved(spark, solver, q, both) - assert resolved - for cid, channels in resolved.items(): - candidates = [ch for one in per_alias for ch in one.get(cid, [])] - assert channels == [min(candidates)] - - -def test_aliased_and_direct_selectors_on_the_same_tags_have_distinct_ids(): - expr = TagSelector("channel_name") == "Engine RPM" - direct = TimeSeriesSelector(expr) - aliased = TimeSeriesSelector(expr, uses_alias=True) - assert direct.selector_id == zlib.crc32(str(expr).encode()) - assert aliased.selector_id != direct.selector_id - - -def test_aliased_and_direct_poi_selectors_on_the_same_tags_have_distinct_ids(): - expr = TagSelector("channel_name") == "Fault Code" - direct = TimeSeriesSelector(expr, series_type=SeriesType.POINTS_IN_TIME) - aliased = TimeSeriesSelector(expr, uses_alias=True, series_type=SeriesType.POINTS_IN_TIME) - expected_direct = f"{SeriesType.POINTS_IN_TIME}|{expr}|{direct._value_type}" - assert direct.selector_id == zlib.crc32(expected_direct.encode()) - assert aliased.selector_id != direct.selector_id + assert resolved == {cid: sorted(channels) for cid, channels in expected.items()} + assert any(len(channels) == 2 for channels in resolved.values()) From 42d27d8091d22c808a6e557b978a2113e4c7ec06 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Fri, 16 Oct 2026 11:34:59 +0200 Subject: [PATCH 14/16] test(query-engine): add parametrized test for channel_id as int and string in pushdown solver Verify that `DefaultSolver` pushdown produces identical results when `channel_id` is cast to `int` or `string` in the key-value store tables. --- .../solvers/default_solver_pushdown_test.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py index ecfa563e..f2b9e5f0 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py @@ -7,6 +7,7 @@ import numpy as np import numpy.testing as nptest import pandas as pd +import pyspark.sql.functions as F import pytest from pyspark.sql import SparkSession @@ -298,6 +299,35 @@ def test_raw_data_with_trailing_point( rpm = on["rpm"].iloc[0] assert rpm.tstarts[-1] == rpm.tends[-1] # the series ends on a point + @pytest.mark.parametrize("channel_id_type", ["int", "string"]) + def test_channel_id_types( + self, + spark: SparkSession, + setup_key_value_store_db, + mock_workspace_client, + udf_rows, + channel_id_type, + ): + tables = { + name: spark.read.table(f"spark_catalog.silver_key_value_store.{name}") + for name in ("container_tags", "container_metrics", "channel_metrics", "channels") + } + for name in ("channel_metrics", "channels"): + tables[name] = tables[name].withColumn( + "channel_id", F.col("channel_id").cast(channel_id_type) + ) + db = MeasurementDB(MeasurementDBConfig.for_debug(tables), ws=mock_workspace_client) + off, on = _solve_both( + spark, + db, + lambda q: [q.channel(channel_name="Engine RPM").resample(RATE).alias("rpm")], + config=_kvs_config(), + ) + _assert_same(on, off) + assert all(np.isfinite(s.values).all() and len(s) > 1 for s in on["rpm"]) + _, rows_on = udf_rows + assert {variant for _, variant in rows_on} == {1} + class TestCalculatedChannels: def test_calculated_channel_with_resample( From 4ed38cd22f9a9d369f2e5b7fa44518b0f4a29215 Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Fri, 16 Oct 2026 14:33:11 +0200 Subject: [PATCH 15/16] fix(query-engine): validate channel_id type from table schema before preparing channel join Move `_validate_calculated_channel_id_type` before `_prepare_channels_join` and read the `channel_id` dtype from the raw `channels` table schema instead of the joined query. This avoids triggering a Spark job during validation and ensures the check runs early. Update the test to assert that channel data preparation is not reached when the type is invalid. --- .../analyze/query/solvers/default_solver.py | 14 +++++++++----- ...fault_solver_native_calculated_channels_test.py | 12 +++++++++++- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index d4d90ce9..57245da0 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -1637,11 +1637,11 @@ def solve_calculated_channels(self, query, channels_df, selections) -> DataFrame :meth:`PushdownPlan.native_targets`), the rows are produced in Spark from the derived rows and no pandas UDF runs. """ + self._validate_calculated_channel_id_type(query) col_map = self.config.col_map q, joined_df, container_count, plan, targets = self._prepare_channels_join( query, channels_df, native_selections=selections ) - self._validate_calculated_channel_id_type(q) if targets is not None: logger.info("Calculated channels solved natively in Spark, without the pandas UDF") @@ -1691,13 +1691,17 @@ def _attach_identity(self, rows: DataFrame, selections) -> DataFrame: ) return rows.withColumn("identity", identity_expr) - def _validate_calculated_channel_id_type(self, channels_q: DataFrame) -> None: - """Raise unless the channel-data ``channel_id`` can hold calculated-channel ids. + def _validate_calculated_channel_id_type(self, query) -> None: + """Raise unless the ``channels`` table's ``channel_id`` can hold calculated-channel ids. :attr:`CalculatedChannel.channel_id` is a 31-bit integer, which only fits - ``IntegerType`` and ``LongType`` columns. + ``IntegerType`` and ``LongType`` columns. Only the table schema is read, so + the check runs no Spark job. """ - dtype = channels_q.schema[self.config.channel_id_col].dataType + channels = self._apply_column_mapping( + query.db.channels(self.spark), self.config.channels.column_name_mapping + ) + dtype = channels.schema[self.config.channel_id_col].dataType if not isinstance(dtype, T.IntegerType | T.LongType): raise ValueError( "Calculated channels need an integer or long channel_id column, got " diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py index 30b2ff60..81e331b4 100644 --- a/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py @@ -308,8 +308,18 @@ def channels(q): @pytest.mark.parametrize("channel_id_type", ["short", "string"]) @pytest.mark.parametrize("enable_pushdown", [False, True]) def test_channel_id_must_be_integer_or_long( - spark, setup_key_value_store_db, mock_workspace_client, channel_id_type, enable_pushdown + spark, + setup_key_value_store_db, + mock_workspace_client, + monkeypatch, + channel_id_type, + enable_pushdown, ): + def _prelude(*args, **kwargs): + raise AssertionError("the channel data was prepared before the channel_id check") + + # The check reads only the schema, before the prelude runs its first Spark job. + monkeypatch.setattr(DefaultSolver, "_prepare_channels_join", _prelude) tables = { name: spark.read.table(f"{_KVS_SCHEMA}.{name}") for name in ("container_tags", "container_metrics", "channel_metrics", "channels") From b205c3b2ea25f17eeb36139a43c7d472d0e74e0e Mon Sep 17 00:00:00 2001 From: "tom.bonfert" Date: Fri, 16 Oct 2026 14:58:34 +0200 Subject: [PATCH 16/16] feat(query-engine): deduplicate channel selector ids via shared helper Add `PushdownPlan.channel_selector_ids` to centralize casting, filtering nulls, and deduplicating a channel's `selector_ids` as `long`. Use it in both `channel_variants` and `DefaultSolver` native target lookup so variant tagging and calculated-channel rows agree on which selectors a channel carries. Add unit tests covering deduplication, null handling, and empty arrays. --- .../analyze/query/optimizer/plan.md | 13 +++++++++++++ .../analyze/query/optimizer/plan.py | 19 ++++++++++++++++--- .../analyze/query/solvers/default_solver.py | 8 +------- .../unit/analyze/query/optimizer/plan_test.py | 11 +++++++++++ 4 files changed, 41 insertions(+), 10 deletions(-) diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md index 12a0adcf..06e1006e 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md @@ -123,6 +123,19 @@ A channel ships only its raw data (``[0]``) unless every one of its selectors is pushed down; then it ships the derived variants instead. +#### channel\_selector\_ids + +```python +def channel_selector_ids(selector_ids: Column | str) -> Column +``` + +A matched channel's distinct, non-null ``selector_ids``, as ``long``. + +Variant tagging (:meth:`channel_variants`) and the native calculated-channel +rows both read a channel's selectors through this, so they agree on which +selectors it carries. + + #### describe ```python diff --git a/src/impulse_query_engine/analyze/query/optimizer/plan.py b/src/impulse_query_engine/analyze/query/optimizer/plan.py index f7181ac9..295b918f 100644 --- a/src/impulse_query_engine/analyze/query/optimizer/plan.py +++ b/src/impulse_query_engine/analyze/query/optimizer/plan.py @@ -179,9 +179,7 @@ def channel_variants(self, selector_ids: Column) -> Column: A channel ships only its raw data (``[0]``) unless every one of its selectors is pushed down; then it ships the derived variants instead. """ - ids = F.filter( - F.transform(selector_ids, lambda sid: sid.cast("long")), lambda sid: sid.isNotNull() - ) + ids = PushdownPlan.channel_selector_ids(selector_ids) pushed = sorted(self.variants_by_selector) needs_raw = F.exists(ids, lambda sid: ~sid.isin(pushed)) | (F.size(ids) == 0) mapping = F.create_map( @@ -206,6 +204,21 @@ def channel_variants(self, selector_ids: Column) -> Column: derived ) + @staticmethod + def channel_selector_ids(selector_ids: Column | str) -> Column: + """A matched channel's distinct, non-null ``selector_ids``, as ``long``. + + Variant tagging (:meth:`channel_variants`) and the native calculated-channel + rows both read a channel's selectors through this, so they agree on which + selectors it carries. + """ + return F.array_distinct( + F.filter( + F.transform(selector_ids, lambda sid: sid.cast("long")), + lambda sid: sid.isNotNull(), + ) + ) + def describe(self) -> str: """One-line summary of the plan for logging.""" variants = ", ".join( diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index 57245da0..c4d53a8f 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -1768,16 +1768,10 @@ def _native_calculated_channel_rows( ] ) by_selector = F.try_element_at(by_variant, F.col(VARIANT_COL)) - selector_ids = F.array_distinct( - F.filter( - F.transform("selector_ids", lambda sid: sid.cast("long")), - lambda sid: sid.isNotNull(), - ) - ) no_targets = F.array().cast(target_type) channel_targets = F.flatten( F.transform( - selector_ids, + PushdownPlan.channel_selector_ids("selector_ids"), lambda sid: F.coalesce(F.try_element_at(by_selector, sid), no_targets), ) ) diff --git a/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py index 2466d93e..ad001174 100644 --- a/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py @@ -10,6 +10,7 @@ from impulse_query_engine.analyze.query.channels.calculated_channel import CalculatedChannel from impulse_query_engine.analyze.query.optimizer import ( PushdownContext, + PushdownPlan, QueryOptimizer, ResamplePushdownRule, ) @@ -45,6 +46,7 @@ def test_channel_variants(spark, selectors): "pushed": ([pushed.selector_id], list(pushed_variant)), "raw_and_pushed": ([raw.selector_id, pushed.selector_id], [0]), "pushed_and_null": ([pushed.selector_id, None], list(pushed_variant)), + "pushed_twice": ([pushed.selector_id, pushed.selector_id], list(pushed_variant)), "two_pushed": ( [pushed.selector_id, pushed_2.selector_id], list(pushed_variant + pushed_2_variant), @@ -63,6 +65,15 @@ def test_channel_variants(spark, selectors): assert actual == {name: expected for name, (_, expected) in cases.items()} +def test_channel_selector_ids(spark): + df = spark.createDataFrame( + [(0, [7, None, 7, 3]), (1, []), (2, None)], schema="i int, selector_ids array" + ) + result = df.select("i", PushdownPlan.channel_selector_ids("selector_ids").alias("ids")) + assert result.schema["ids"].dataType.elementType == T.LongType() + assert {r["i"]: r["ids"] for r in result.collect()} == {0: [7, 3], 1: [], 2: None} + + class TestNativeTargets: def test_whole_expression_pushdowns(self, selectors): rpm, spd = selectors["pushed"], selectors["pushed_2"]