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/config/configuration.md b/docs/impulse/docs/config/configuration.md index 3ad82f76..ccab0582 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). | | `measurement_db` | `str` | `"MeasurementDB"` | Registered name of the MeasurementDB implementation to use. The default is the built-in reader; a custom implementation can reshape a different source schema at read time. See [Custom measurement DB](../extensibility/custom_measurement_db.md). | | `measurement_db_config` | `dict` | `null` | Keyword arguments for that implementation's config class, merged with the `source` tables (e.g. extra source tables, prefilter keys). | 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/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/plan.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md new file mode 100644 index 00000000..06e1006e --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/plan.md @@ -0,0 +1,147 @@ +--- +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. + +**Raises**: + +- `ValueError`: If a pushdown computes on another pushdown's result (``source``), +which the pipeline cannot compute. + +#### 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. + +**Raises**: + +- `ValueError`: If two different rules share a name: their pipeline columns would +collide. + +#### lookup + +```python +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 + +```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 +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. + + +#### 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 +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..7a905ff3 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/pushdown.md @@ -0,0 +1,256 @@ +--- +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. + + +#### 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 +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. + +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 + +```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. + + +#### 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 +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 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 + +Unique rule name, part of every :attr:`Pushdown.variant_key`; lowercase + +letters, digits and single underscores (see :func:`check_name`). + + +#### 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 +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. + + +#### 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*. + + +#### runs + +```python +def runs() -> 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; + + +#### row\_columns + +```python +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 + +```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**: + +`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 + +```python +def decode(rows: pd.DataFrame, pushdown: Pushdown, + cache: TimeSeriesCache) -> Any +``` + +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/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..31d21612 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/query_optimizer.md @@ -0,0 +1,58 @@ +--- +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. + +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 +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..af0cdf38 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/optimizer/resample_rule.md @@ -0,0 +1,61 @@ +--- +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. + +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 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. + +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. + + +## \_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/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 5c92a6b9..c28b8873 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 @@ -211,6 +211,11 @@ project/toolbox filtering is applied. - `measurement_db` (`str, default="MeasurementDB"`): Registered name of the MeasurementDB implementation to use. - `measurement_db_config` (`dict`): Keyword arguments for that implementation's config class, merged with the ``source`` tables. +- `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..92c8d3f9 100644 --- a/docs/impulse/docs/references/api/sidebar.json +++ b/docs/impulse/docs/references/api/sidebar.json @@ -30,6 +30,18 @@ "label": "impulse_query_engine.analyze.query.channels", "type": "category" }, + { + "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", + "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/img/solve_flow.svg b/docs/impulse/docs/references/query_engine/img/solve_flow.svg new file mode 100644 index 00000000..9f6b5b79 --- /dev/null +++ b/docs/impulse/docs/references/query_engine/img/solve_flow.svg @@ -0,0 +1,244 @@ + + 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 + solve(): build 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 + alias resolution + + + + + + + + channel matches: container_id · channel_id · selector_ids + + + + + + 3 + Load channel data · Spark, channels table + + + channels + column mapping, filters + + + Raw mode: encode + RLE or intervals + + + Add POI channels + when POI channels are selected + + + Join channel matches + keep matched channels + + + + + + + + + + 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; + all share the reads + of the channel data + + + + + + + 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. + + + Encode each result + solve(): serialized value · calculated channels: rows + + + + + + + + + + + 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/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index c71cd297..77b79eb5 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"`. @@ -142,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 @@ -152,12 +157,83 @@ 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 [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, and channels whose +`[tstart, tend)` intervals overlap are not supported (the RLE and interval encoders never +produce overlaps). + +#### 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 + +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: @@ -167,3 +243,9 @@ 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; + [`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/docs/references/query_engine/solving.md b/docs/impulse/docs/references/query_engine/solving.md new file mode 100644 index 00000000..aa6728fe --- /dev/null +++ b/docs/impulse/docs/references/query_engine/solving.md @@ -0,0 +1,151 @@ +--- +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)). + `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. 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 + 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 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. +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` 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) + +- 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. +- 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 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): + +``` +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 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 + +| 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; 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). +::: 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..3caa0566 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,12 @@ 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 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 Not every channel can be interpreted as a continuous signal within its `[tstart, tend)` intervals. @@ -95,6 +101,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 @@ -177,6 +186,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..99345aeb 100644 --- a/docs/impulse/docs/references/query_engine/tsal/evaluation.md +++ b/docs/impulse/docs/references/query_engine/tsal/evaluation.md @@ -63,7 +63,12 @@ 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 -`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. diff --git a/docs/impulse/pydoc-markdown.yml b/docs/impulse/pydoc-markdown.yml index f52dda2b..7fa9bb8d 100644 --- a/docs/impulse/pydoc-markdown.yml +++ b/docs/impulse/pydoc-markdown.yml @@ -24,6 +24,12 @@ 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.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 - 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..75824ad8 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/__init__.py @@ -0,0 +1,35 @@ +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, + 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", + "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..60c8c9e1 --- /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_rule__{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..aa186750 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/pipeline.py @@ -0,0 +1,550 @@ +"""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, check_name + +#: 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}__{check_name('run', 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") + # 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(_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(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( + *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(check_name("column", 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/plan.py b/src/impulse_query_engine/analyze/query/optimizer/plan.py new file mode 100644 index 00000000..295b918f --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/plan.py @@ -0,0 +1,230 @@ +"""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 +from typing import Any + +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 +from .traversal import Classifier + + +@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. + + 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 + 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. + + 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 + 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, 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) + if variant in self.variants_by_selector.get(pushdown.selector.selector_id, ()): + return pushdown, variant + 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]]]] = {} + classify = Classifier(self.rules) + for selection in selections: + if not isinstance(selection, CalculatedChannel): + return None + 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 + 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 = 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( + *[ + 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]]), + ) + ] + ) + no_variants = F.array().cast("array") + derived = F.array_distinct( + 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 + ) + + @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( + 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..337a0d35 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/pushdown.py @@ -0,0 +1,228 @@ +"""Building blocks of Spark pushdown rules.""" + +from __future__ import annotations + +import abc +import re +from collections.abc import Callable, Mapping +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 Column + +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 + + 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. +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: + """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. + + 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 + 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, 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.""" + 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) + 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 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)``. + """ + + #: 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. + ordered: ClassVar[bool] = True + + @abc.abstractmethod + 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. + """ + + 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 + + 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. 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``. + + Column names follow :func:`check_name` and only need to be unique within + the rule. + """ + 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 emit(self, ctx: SeriesContext) -> Column: + """The derived rows of a kept sample. + + Returns + ------- + 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: + """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/optimizer/query_optimizer.py b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py new file mode 100644 index 00000000..009d837d --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/query_optimizer.py @@ -0,0 +1,120 @@ +"""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 functools import partial +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, check_name +from .resample_rule import ResamplePushdownRule +from .traversal import Classifier, 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. + + 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 + 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() + ) + 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. + + 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) + + 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, 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 + # of the tree; its candidates can't be trusted to cover all occurrences. + if walked == expected: + candidates.extend(found) + + # 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]} + 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 _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: + 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..d71f1020 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/resample_rule.py @@ -0,0 +1,201 @@ +"""Pushdown rule computing ``SampleSeries.resample`` natively in Spark.""" + +from __future__ import annotations + +import math +import numbers +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 + +from impulse_query_engine.analyze.metadata.time_series_expression import ( + SeriesType, + TimeSeriesExpression, + TimeSeriesOp, + TimeSeriesSelector, +) + +from .context import SeriesContext +from .pushdown import Operand, Pushdown, PushdownContext, PushdownRule + +_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. + + 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 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. + + 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, 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 + selector = operand(node.args[0]) + if not isinstance(selector, TimeSeriesSelector): + return None + 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 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) + ), + } + + 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")) + + 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 + ), + } + + 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, + ), + ) + return F.when(~ctx.col("resampled"), F.array(ctx.element(value=ctx.value))).when( + c_stop > c, points + ) + + +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)) + + +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)``. + """ + + 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 + + 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*). + + 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) + + estimate, step = self.ctx.col(f"{name}_estimate"), self.ctx.col(f"{name}_step") + return { + f"{name}_estimate": F.when( + active, + 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") + ) + ), + ), + f"{name}_step": F.when((estimate > 0) & reached(estimate - 1), estimate - 1) + .when(~reached(estimate), estimate + 1) + .otherwise(estimate), + 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: + """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 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..9db5d358 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/optimizer/traversal.py @@ -0,0 +1,141 @@ +"""Bottom-up classification of TSAL expression trees into pushdowns.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Iterator, Sequence +from typing import Any + +from impulse_query_engine.analyze.metadata.time_series_expression import ( + TimeSeriesAliasSelector, + TimeSeriesExpression, + TimeSeriesOp, + TimeSeriesSelector, +) + +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*. + + 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()) + + +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. + + 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. + """ + + 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 + + +def discover( + root: TimeSeriesExpression, classify: Classifier, accept: Accept | None = None +) -> tuple[list[tuple[TimeSeriesExpression, Pushdown]], list[TimeSeriesSelector]]: + """Find the topmost subtrees of *root* that are pushed down. + + 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 ``(node, pushdown)`` candidates 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 = 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) + 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..c4d53a8f 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 @@ -7,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 @@ -16,6 +17,15 @@ SeriesValueType, TimeSeriesExpression, ) +from impulse_query_engine.analyze.query.optimizer import ( + RAW_VARIANT, + VARIANT_COL, + Classifier, + 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 @@ -29,6 +39,8 @@ if TYPE_CHECKING: from impulse_query_engine.measurement_db import MeasurementDB +logger = logging.getLogger(__name__) + class TimeSeriesCache(SeriesCache): def __init__( @@ -37,6 +49,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 +65,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 +80,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 +96,34 @@ 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 + # 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. 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 +193,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 +208,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, self._classify) + 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 +339,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 +354,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) @@ -714,8 +805,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") @@ -1007,6 +1099,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 +1117,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 +1132,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 +1171,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 +1182,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 +1195,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, 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 @@ -1110,20 +1209,31 @@ 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 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] - ``(channels_q, joined_df, container_count)`` — 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. + 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 @@ -1168,12 +1278,58 @@ 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 + targets = None + if plan is not None: + 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], ) - 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, targets + + 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, + 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 (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) + return apply_pushdowns(rows, plan, self.config, include_raw=include_raw) @staticmethod def _query_contains_poi_selections(selections: Iterable[Any]) -> bool: @@ -1382,6 +1538,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 +1566,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 +1580,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"] @@ -1470,9 +1630,24 @@ 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. """ + self._validate_calculated_channel_id_type(query) col_map = self.config.col_map - q, joined_df, container_count = self._prepare_channels_join(query, channels_df) + q, joined_df, container_count, plan, targets = self._prepare_channels_join( + query, channels_df, native_selections=selections + ) + + 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) @@ -1483,6 +1658,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, @@ -1495,8 +1671,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: @@ -1507,4 +1689,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, 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. Only the table schema is read, so + the check runs no Spark job. + """ + 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 " + 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)) + no_targets = F.array().cast(target_type) + channel_targets = F.flatten( + F.transform( + PushdownPlan.channel_selector_ids("selector_ids"), + lambda sid: F.coalesce(F.try_element_at(by_selector, sid), no_targets), + ) + ) + has_conversion = conv in joined_df.columns + # joined_df holds no raw rows here; any other row without targets is + # dropped by inline. + rows = joined_df.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) + # 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. + 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/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 79b27472..a6301ca3 100644 --- a/src/impulse_reporting/config/config_parser.py +++ b/src/impulse_reporting/config/config_parser.py @@ -395,6 +395,12 @@ class QueryEngine(BaseModel): measurement_db_config : dict, optional Keyword arguments for that implementation's config class, merged with the ``source`` tables. + 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 ----- @@ -423,6 +429,7 @@ class QueryEngine(BaseModel): measurement_db_config: dict[str, Any] | 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 f8d9bdab..bd2d625d 100644 --- a/src/impulse_reporting/core/report.py +++ b/src/impulse_reporting/core/report.py @@ -362,6 +362,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/pipeline_helpers.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py new file mode 100644 index 00000000..3aca63a4 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_helpers.py @@ -0,0 +1,188 @@ +"""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()), + ] + ) + 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() + 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, operand, 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 + ): + return None + 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): + """Test-only ``channel.diff()``: the value change within a block, 0 at block starts.""" + + name = "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,)) + + 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, operand): + 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..21849de8 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/pipeline_test.py @@ -0,0 +1,297 @@ +# 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, + Classifier, + 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 ( + CONFIG, + DiffRule, + OffsetRule, + ProbeRule, + channel_rows, + chunk_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: 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(): + 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) + + +_OFFSET = OffsetRule() + + +def _offset(value: float) -> Pushdown: + return Pushdown(rule=_OFFSET, 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: 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]) + 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: 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)} + + +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: 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] + # 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 + + +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: Classifier([rule])(_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/plan_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py new file mode 100644 index 00000000..ad001174 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/plan_test.py @@ -0,0 +1,119 @@ +# 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, + PushdownPlan, + 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)), + "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), + ), + } + 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()} + + +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"] + 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/optimizer/query_optimizer_test.py b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py new file mode 100644 index 00000000..79db1458 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/query_optimizer_test.py @@ -0,0 +1,409 @@ +# 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, + TimeSeriesOp, + 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 ( + Pushdown, + PushdownContext, + PushdownPlan, + QueryOptimizer, + ResamplePushdownRule, +) +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()) +_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 _MatchesAnyNode(ResamplePushdownRule): + """Claims every non-selector node that reads exactly one selector, of any type.""" + + name = "any_node" + + def match(self, node, operand): + 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]), 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), Classifier([_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 _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): + 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), + [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"): + plan.variants_by_rule() + + +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 Classifier([ResamplePushdownRule()])(make(rpm)) is None + + def test_default_rate(self, rpm): + 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 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 + 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, operand): + pushdown = super().match(node, operand) + 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..68091c29 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/optimizer/resample_rule_test.py @@ -0,0 +1,286 @@ +# 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 +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, + 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 ( + CONFIG, + DiffRule, + channel_rows, + plan_for, + random_series, + run, +) +from impulse_query_engine.analyze.query.optimizer import apply_pushdowns + +_RULE = ResamplePushdownRule() +_X = TimeSeriesSelector(TagSelector("channel_name") == "x") + + +def _pushdown(rate): + return Classifier([_RULE])(_X.resample(rate)) + + +def _resample_in_spark(spark, series_by_key, rates, time_type=T.LongType(), chunk_rows=None): + """Resample ``{(cid, ch, variant): (tstarts, tends, values)}`` at ``{variant: rate}``. + + Variant 0 rows are raw rows. Returns ``{(cid, ch, variant): SampleSeries}`` + built like the pandas cache does. + """ + 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): + 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), + "grid_on_sample_starts": (list(range(10)), list(range(1, 11)), list(range(10)), 3), +} + + +@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)) + + +# 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_raw_rows_unchanged(spark, chunk_rows): + 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) 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(): + _assert_same(actual[key], exp) + + +@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 + 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}, chunk_rows=chunk_rows + ) + expected = _expected((tstarts, tends, values), rate) + assert len(expected) < len(tstarts) / 5 + _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. + plan = _optimized_plan(spark) + assert "specifiedwindowframe(RowFrame, currentrow$(), unboundedfollowing$())" not in plan + assert "specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())" in plan + + +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 _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. + plan = _optimized_plan(spark) + 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 == ["_pd_chunk"] + else: + 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) + + +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: Classifier([DiffRule()])(_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]) +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 + 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}, chunk_rows=chunk_rows + ) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) + + +@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(), + chunk_rows=chunk_rows, + ) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), rate)) + + +@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(), + chunk_rows=chunk_rows, + ) + _assert_same(actual[(1, 1, 1)], _expected((tstarts, tends, values), 0.75)) + + +@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. + 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(), + chunk_rows=chunk_rows, + )[(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_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..f6e8ac22 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_channel_resolution_test.py @@ -0,0 +1,176 @@ +# pylint: disable=missing-function-docstring +"""DefaultSolver channel resolution when several channels of a container match a selection. + +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 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" +# 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: + 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_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) == _CELSIUS_CHANNELS + + 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) + ) + 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") + 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", next to channel 5 of containers 1 and 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, 9], + 2: [5], + 3: [2, 8], + } + + +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_each_alias( + self, spark, key_value_store_alias_db + ): + q = key_value_store_alias_db.query + solver = DefaultSolver(spark, config=_alias_config()) + 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 == {cid: sorted(channels) for cid, channels in expected.items()} + assert any(len(channels) == 2 for channels in resolved.values()) 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..81e331b4 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_native_calculated_channels_test.py @@ -0,0 +1,338 @@ +# 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.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, + SolverConfig, + TableConfig, +) +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 ( + PUSHED_SCANS, + SAMPLE_SHUFFLE, + channel_data_scans, + container_shuffles_of_channel_data, + shuffles, + 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" + + +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 + + +@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): + 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 + 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 + ): + 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_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}") + 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, + 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") + } + 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_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..f2b9e5f0 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/default_solver_pushdown_test.py @@ -0,0 +1,439 @@ +# 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 +import numpy.testing as nptest +import pandas as pd +import pyspark.sql.functions as F +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, + Classifier, + PushdownPlan, + QueryOptimizer, + ResamplePushdownRule, + pipeline, +) +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 small_chunks(monkeypatch): + """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*.""" + return sorted( + re.sub(r"#\d+L?", "", columns) + for columns in re.findall(r"FileScan parquet \S+\.channels\[([^\]]*)\]", _final_plan(df)) + ) + + +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 +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]) + + +@pytest.mark.usefixtures("small_chunks") +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_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() + assert rows and all(row["speed_max"] > 0 for row in rows) + 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 + ): + 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 + + @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( + 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 = Classifier([rule])(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_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] + + +def _rpm(uses_alias: bool = False) -> TimeSeriesSelector: + return TimeSeriesSelector(TagSelector("channel_name") == "rpm", uses_alias=uses_alias) diff --git a/tests/impulse_reporting/integration/calculated_channel_test.py b/tests/impulse_reporting/integration/calculated_channel_test.py index 5a426e8e..44ba2eb2 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 @@ -614,3 +615,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) diff --git a/tests/impulse_reporting/unit/config/config_parser_test.py b/tests/impulse_reporting/unit/config/config_parser_test.py index b935f0d0..287dc83b 100644 --- a/tests/impulse_reporting/unit/config/config_parser_test.py +++ b/tests/impulse_reporting/unit/config/config_parser_test.py @@ -532,6 +532,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 696125d4..a60493af 100644 --- a/tests/impulse_reporting/unit/core/report_test.py +++ b/tests/impulse_reporting/unit/core/report_test.py @@ -314,6 +314,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