Repository navigation
feat(query-engine): add opt-in rule-based Spark pushdown optimizer, starting with resample - #129
Open
tombonfert wants to merge 17 commits into
Open
tombonfert wants to merge 17 commits into
tombonfert wants to merge 17 commits into
Conversation
Introduce a rule-based query optimizer that computes supported expression subtrees natively in Spark before the pandas UDF, starting with `channel.resample(rate)` on physical channels. - Add `QueryOptimizer`, `PushdownPlan`, `PushdownRule`, and `PushdownContext` in `analyze.query.optimizer`. - Implement `ResamplePushdownRule` to push `SampleSeries.resample` down to Spark, reproducing the pandas behavior block-by-block. - Wire `enable_pushdown` through `DefaultSolver`, `QueryEngine`, and report config; default is `false`. - Make `TimeSeriesOp.build` consult `SeriesCache.materialized` so pushed-down results can be injected into the evaluation tree. - Add optimizer API reference pages, update configuration docs, and document the Spark pushdown opt-in behavior in query solver and TSAL guides.
…ndas bucket semantics Update `ResamplePushdownRule` so each sample is assigned the bucket of the first grid point at or after its `tstart`, and only the last sample of a bucket emits the grid points it covers. This fixes cases where grid points land exactly on sample starts and avoids emitting duplicate or over-wide rows. Add unit tests covering grid-on-sample-start scenarios, including nanosecond-epoch timestamps that require the bucket-index correction step.
Adjust `_with_first_grid_index` to nudge the estimated bucket index up or down only when the reachability check at the adjacent index differs, falling back to the raw estimate otherwise. This tightens the correction step introduced for nanosecond-epoch timestamps without over-adjusting cases where the initial estimate already lands on the correct grid point.
…lementation details Update API reference docs for `PushdownRule.to_spark` to describe the single-pass, lineage-preserving contract for replacing variant rows with derived rows. Expand `ResamplePushdownRule` and query solver docs to explain grid bucket assignment, last-sample row emission, and the bounded memory per row.
Replace the `ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING` frame in `ResamplePushdownRule` with a two-pass approach using only linear frames. First pass computes per-block state with running/whole-partition windows; second pass derives block `t_max` and grid counts over the filtered last-sample rows. This prevents Spark from recomputing an unbounded-following frame per row. Add a unit test asserting the optimized plan contains no unbounded-following row frame.
… in Spark Add `PushdownPlan.native_targets` and `PushdownRule.emits_result_rows` so `DefaultSolver.solve_calculated_channels` can bypass the pandas UDF when every selection is a `CalculatedChannel` whose entire expression is a pushdown that emits its result rows directly. Native rows are produced from the derived variant rows and attached to the same identity map as the UDF path. Also make `DefaultSolver` resolve selectors to at most one channel per container, keeping the lowest `channel_id` (after alias priority), and give aliased and direct selections on the same tags distinct `selector_id`s. Update docs and add unit/integration tests covering native calculated channels, channel resolution deduplication, and pushdown parity.
Update `ResamplePushdownRule` to process each series in equal-time chunks of about `chunk_rows` samples (default 100,000) instead of whole series. Per-chunk windows keep worker memory bounded; chunk-boundary state (gaps, block start) is resolved from a small per-chunk summary broadcast-joined back. The summary shares one extra narrow read of the channel key/time columns, so Spark reads channel data once in full plus one narrow read. - Add `chunk_rows` parameter with validation. - Update `PushdownRule.to_spark` contract to allow broadcast-joining small summaries derived from `rows`. - Stop filtering raw variant rows before the pushdown input so the narrow summary read can be shared across rules. - Add unit tests for chunking, chunk-size validation, random series across chunk sizes, and parity with existing cases. - Update API docs and query solver docs to describe chunked resampling and the two-channel-data-read plan.
Replace per-rule `PushdownRule.to_spark` with a shared `SeriesPipeline` that evaluates all ordered rules in one pass over the channel data. Rules now declare hooks (`row_columns`, `keep`, `kept_columns`, `emit`) operating on a `SeriesContext` of Spark columns, so reads and shuffles are shared across rules. Add `context.py` with `SeriesContext`, `Sample`, `Run`, and `RunBreak`; add `pipeline.py` with chunking, per-chunk windows, chunk summaries, and bounded-memory series processing. Convert `ResamplePushdownRule` to the new hook interface and remove its `chunk_rows` parameter. Update `DefaultSolver` to use `apply_pushdowns`, regenerate API docs, and expand query solver docs to describe the shared pushdown pass.
Add a new documentation page that walks through the five stages of a `solve()` call, from driver-side planning through metadata resolution, channel loading, Spark pushdown, and per-container pandas UDF evaluation. Include a supporting diagram and link it from the query engine index and TSAL evaluation docs.
… behavior Update the "How a Query Is Solved" page and flow diagram to reflect that: - `solve_calculated_channels()` skips driver-side type inference and emits sample rows. - Alias resolution and unit conversion are now separate stages in the channel-match pipeline. - Stage 4 pushdown rules share channel-data reads, with ordered rules running per chunk. - Output encoding differs between `solve()` (serialized columns) and `solve_calculated_channels()` (rows). Also fix the `query_solvers.md` example output schema to list `identity` instead of `channel_name`/`data_key`.
…and non-op nodes, fix alias selector ids and NaN conversion factors - Enforce lowercase rule/column/run names and unique rule names to prevent pipeline column collisions; add `check_name` helper and validation in `QueryOptimizer` and `PushdownPlan`. - Change pipeline column separators to `__` so rule names cannot clash with internal columns. - Restrict pushdown matching to `TimeSeriesOp` nodes via `match_node`; aggregations, UDFs, and events now correctly keep raw rows. - Fix `TimeSeriesSelector.selector_id` for aliased non-CONTINUOUS selections so aliased and direct selectors on the same tags remain distinct. - Treat NaN unit-conversion factors like null factors, leaving values unconverted instead of nulling them. - Handle null/NaN `tstart` and all-NaN `tend` in chunking so series ordering and chunk assignment stay consistent. - Add unit tests for name validation, collision avoidance, non-op node exclusion, selector id distinctness, and NaN conversion factors.
…operands Introduce `Classifier` to memoize bottom-up classification of `TimeSeriesExpression` nodes into pushdowns, replacing the top-down `match_node` helper. Rules now receive an `Operand` callback that exposes whether each operand is a raw selector, a pushed subexpression, or evaluated in pandas, enabling rules to absorb or chain pushed operands. - Add `Classifier` and `Operand` to the optimizer public API and docs. - Update `PushdownRule.match` signature to accept `operand`. - Add `Pushdown.source` for pushdowns that compute on another pushdown's result; reject planned chained pushdowns in `PushdownPlan.build`. - Update `QueryOptimizer` to accept only single-selector, non-chained pushdowns and share one `Classifier` across discovery. - Update `ResamplePushdownRule` and test helpers to the new `match` interface. - Cache one `Classifier` per container in `TimeSeriesCache` so UDF node lookups classify each node once. - Add unit tests for bottom-up classification, chained/absorbed operands, unsupported parent fallback, and shared-node memoization.
… per container Remove the "lowest channel_id wins" deduplication for direct selectors in `DefaultSolver` so a selector can resolve to every matching channel per container. Aliased selectors still resolve to one channel per container (best alias priority, then lowest channel_id). Update `TimeSeriesSelector.selector_id` to no longer encode `uses_alias`, since direct and aliased selections are now handled separately. Adjust docs, generated API reference, and tests to reflect the new resolution behavior.
…tring in pushdown solver Verify that `DefaultSolver` pushdown produces identical results when `channel_id` is cast to `int` or `string` in the key-value store tables.
…preparing channel join Move `_validate_calculated_channel_id_type` before `_prepare_channels_join` and read the `channel_id` dtype from the raw `channels` table schema instead of the joined query. This avoids triggering a Spark job during validation and ensures the check runs early. Update the test to assert that channel data preparation is not reached when the type is invalid.
Add `PushdownPlan.channel_selector_ids` to centralize casting, filtering nulls, and deduplicating a channel's `selector_ids` as `long`. Use it in both `channel_variants` and `DefaultSolver` native target lookup so variant tagging and calculated-channel rows agree on which selectors a channel carries. Add unit tests covering deduplication, null handling, and empty arrays.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #129 +/- ##
==========================================
+ Coverage 90.30% 91.18% +0.87%
==========================================
Files 68 75 +7
Lines 6068 6818 +750
Branches 758 841 +83
==========================================
+ Hits 5480 6217 +737
- Misses 466 474 +8
- Partials 122 127 +5
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
tombonfert
commented
Oct 16, 2026
Comment on lines
+29
to
+31
| T.StructField("tstart", T.DoubleType()), | ||
| T.StructField("tend", T.DoubleType()), | ||
| T.StructField("value", T.DoubleType()), |
Collaborator
Author
There was a problem hiding this comment.
get the column names from the solver config.
tombonfert
commented
Oct 16, 2026
Comment on lines
+38
to
+45
| # 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" |
Collaborator
Author
There was a problem hiding this comment.
those column names should be added to the solver config
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a rule-based optimizer that computes supported expression subtrees natively in Spark before the pandas UDF, so fewer rows are shipped to it. Pushdown is off by default, and with
enable_pushdown=Falsethe solver works as before.What's new
DefaultSolver(spark, enable_pushdown=True), orquery_engine.enable_pushdown: truein a report config.channel.resample(sample_rate)onchannel()andchannel_with_alias()selections, with the same results asSampleSeries.resample.channelstable is read once in full and once narrow (for the chunk boundaries), and the samples are shuffled once. Channels that are also used outside a pushed-down subtree add one read for their raw rows.TimeSeriesOp.buildfirst asksSeriesCache.materialized(node)for a precomputed result, so user expressions are never rewritten.Behaviour with pushdown off
channel_idinstead of an arbitrary match.solve_calculated_channelschecks thechannel_idtype of thechannelstable before any Spark job runs, and raises unless it isintorlong. A stringchannel_idalready fails inside the UDF onmain(see Support string channel_id for calculated channels #127).Docs
Known limitations
channel_idyet (Support string channel_id for calculated channels #127).solve()with pushdown works with stringchannel_ids.