From 1205c18496ac569dc42d7e0499a8c90b4e177ad1 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 00:35:18 +0800 Subject: [PATCH 1/4] feat(promql): execute Complete aggregates in the native runtime (#9873) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(promql): execute a bound native tree node by node, over sparse evaluated series Ports the native runtime onto the sparse evaluated-series API: a chunk is addressed by a `TimeAxis` that may hold several runs, so a consumer reads positions through the axis instead of assuming `start + row`. The driver and both node kinds move to that API: - `chunk.range().len()` becomes `chunk.series.time.len()`, the compact row count, in every metric and byte estimate. - `occupied_runs` walks `time.positions()` with one field cursor, so a chunk of several runs records each run where it belongs. A run open across a hole now ends where the hole begins; leaving it open to the last row made two runs look like one and refused a series that occupies the hole. - A window keeps the pieces its evaluator decided in a queue: one push can decide several runs, and reporting only what the push returned stranded the rest until the next push. Queued pieces are charged by `retained_bytes`, and the charge is recomputed when one leaves, so a released piece is not held. - `map_chunk` answers `Option`: an operation that empties every position (`clamp` with `max < min`) yields no chunk at all, unless the chunk closes its series -- a closing chunk holding no row is handed over so a consumer can tell "nothing more here" from "no more input". `RangeEvaluator::retained_bytes` is restored with the three byte helpers it reads: a field's buffers, a series' identity and the pending batch's vectors, all by capacity, and still counted as storage rather than as a peak over a push. Tests read absolute positions from the axis (`decided` no longer adds a row number to a start), and two sparse cases cover what the old reading got wrong: a `Window -> Map -> Window` chain over a chunk of two runs, and a series that occupies the hole between another's runs. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): drop the window queue the sparse evaluator never needed `RangeEvaluator::push` and `push_series` each return `Result>`: the evaluator gathers every run it decided into one chunk and hands back at most one. So `pending.extend(decided)` inserted at most one chunk and `take_pending` removed it on the spot, leaving the queue empty after every successful push. That makes the queue dead weight: a `VecDeque` allocation per productive window, an accounting term that is always zero, a second drain site in `pull`, and doc comments claiming a problem the API does not have. The driver's "ask the input again when a window decides nothing" loop was already the whole story. Back to returning the evaluator's own answer, and drop `pending`, `take_pending` and `take_window_output`. Covers the empty-map protocol this port introduces while it is here: an operation that empties every position hands on nothing while its series is open, and one empty chunk when it closes. A non-closing empty chunk would be rejected downstream and a dropped closing one would leave the series unclosed. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): pin the empty-map output count at one The test checked that every output was empty and that exactly one closed, which a stray extra non-closing empty output would also satisfy. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): charge the memory the native evaluator actually retains Four omissions made the reservation blind to most of what a window over native histograms keeps, and the existing incremental-equals-recomputed test could not see any of them because both sides used the same formula. - `FloatHistogram` owns four heap allocations -- both bucket lists, the custom bounds, and the span layout -- that `FieldSamples::bytes` never looked at. A retained 4096-bucket sample is tens of kilobytes against a ledger that saw a struct header. `heap_bytes` counts them, and both the retained sample slice and the histograms a row boxes are charged for what they own. - Queued chunks were estimated as one float per field per row, which is not what a histogram field holds: it holds indices into a shared `Arc` and the buffers behind the Arrow arrays. The figure now walks the samples, and counts a shared table or a shared buffer slice once per pass rather than once per chunk that references it. - `history_bytes` included `size_of::()` while the map slot term charged a `History` per entry as well, so every entry paid two headers. The per-entry figure now covers only the heap the entry owns. - The primary key is charged for the allocation its `Buffer` keeps alive rather than the bytes it shows: a scan slices the key out of a decoded dictionary-values buffer, so a short key can pin a batch-sized allocation. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): keep the source queue's charge incremental The queue's reservation was recounted by walking every queued chunk on each enqueue and each pop. `push_raw` and `push_evaluated` are public and put no bound on how deep the queue gets, so a caller that queues a whole series before pulling it paid `N(N+1)/2` chunk visits to hold `N` chunks -- the same shape of defect this branch fixes in `absorb`. `QueuedCharge` keeps a reference count and the bytes behind each shared allocation: a table or buffer is charged when the first queued chunk reaches it and given back when the last one leaves, so an operation walks the one chunk it names. The count is per chunk rather than per field, because a pop re-walks the same chunk with the same dedupe and must remove exactly the references that were added. A chunk's own struct bytes are counted per chunk (`own_bytes`); only the allocations behind its fields are counted by address, since a struct has no address to dedupe on. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): append a forward-moving chunk to the emitted history `absorb` rebuilt the whole stored run array on every non-empty chunk. Legal input made that quadratic across chunks with no label collision: one series in `K` chunks, chunk `i` occupying only grid position `2*i`, stores one more run per chunk and moves `K(K+1)/2` of them -- about five billion for a hundred thousand rows. A chunk that starts at or past the end of the last stored run only ever adds to the end, so its runs are appended and the boundary run is joined where it touches. One that can interleave or reach backwards still takes the full one-pass merge, which can split a run. Both leave the runs sorted, disjoint and joined, which is what `Emitted::occupy` looks a position up by. The test records the chunks through `Emitted::occupy` itself and reads the history back from it, covering the forward case, the touching boundary, a hole that has to join both neighbours, and a position before everything stored. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): report a truncated input from a tree with no window `close` only asked windows whether a series was still open, so a tree of an input alone, or of a map over one, accepted a final chunk without `is_last` and handed the consumer a series that never closed. The input node now follows the series of the chunk it last handed over, and closing such a tree reports the truncation against it, which is the same refusal a window makes for the input it reads. `occupied_runs` read a row the axis names with `expect`, which panics on a `Series` whose fields are shorter than its axis; the fields are public, so that chunk is constructible by a caller. It now returns the misalignment as an error naming the node, the field and both lengths. The chunk protocol says one series closes once, which reads as one close per label set. State that it is per input series: a node that rewrites labels can hand two series over under one label set and both of them close. Also say in `Emitted`'s doc that the table is one node's own, so a label set costs its bytes again in every node that keeps one. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): execute an instant selection as a tree node The tree could not run `NativeExpr::Instant`: the state builder recorded `not_implemented` and refused the whole tree before reading a chunk, while the compiler already produced the node for a bare selector. A query the plan and the compiler both accept therefore had no runtime path at all. `InstantRuntime` buffers one series to its `is_last` chunk -- the selection kernel reads a complete axis -- then calls the shared `select_instants` once and gathers the rows it picked onto the binding's own grid. Positions come from the grid lattice the kernel was asked for, and a timestamp off that lattice is an error rather than a skipped row. A series that selects nothing still closes, with an empty chunk on the grid; an input that ends inside a series is reported on `finish()`. The node does not record label-set occupancy: it keeps the metric verbatim, so it cannot map two series onto one label set, and the occupancy key drops empty-valued tags -- two genuinely distinct series would be refused as a collision. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): refuse an interleaved input, and close a label set only for its owner Two series-level identity holes, both reported against `42ea525`. The truncation flag followed the last chunk handed over, which is the same as "is a series open" only while at most one series can be open. A tree with no window still accepted `Interleaved`, and `a` open, `b` closed, end of input then passed while `a` never closed. Tracking the open set would need an identity that a raw chunk in bound-label mode does not have (`metric` is `None`, the key holds it), so the ordering is refused when the state is built instead, for every tree shape and not only under a window. The flag's doc now says what it means and why that is the whole answer. `History` closed on any series' `is_last`, so with `abs` dropping the metric name a series that contributed nothing could close a label set another series was still filling: `a{job="x"}` value, `b{job="x"}` empty close, `a{job="x"}` continuation was refused with "already closed the label set". Contiguous input hands over exactly that sequence, so it was not an interleaving artefact. A history now remembers the series that filled it and closes only for that series' closing chunk; a series that produced nothing under an open set is passed over, as it already was under a closed one. The owner is charged as part of the entry. The ordering check also reads every input node rather than only the first: a tree with two inputs would otherwise have its second declaration read and dropped. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): reserve the instant gather's workspace and keep the histogram start Two findings from the review of the instant node. The gather allocates after the last chunk arrives: the concatenated axis, the descriptors that read it, the two arrays the selection returns, and the rows the builder keeps. Nothing asked the pool for any of it, so with one sample selected by every point of a million-point grid the node took tens of megabytes while its reservation fell to the buffered chunks' size -- nothing at all, on the chunk that ends the series. `gather_bytes` now derives a bound from the values in hand (rows, chunks, segments, fields, points, and the widest histogram row per field, since one row can be selected by every point), and the reservation grows by it before the gather runs. It shrinks back on both the success and the error path, so the workspace never outlives the call. A selection reads an existing sample, so it has to carry that sample's payload. `FloatHistogram` had no `start_timestamp` and `push_to` wrote `None`, which lost it for any histogram that had one: `detect_start_timestamp_reset` reads that field to spot a counter reset, so a window over this output would mis-detect resets. The field is now carried through a view and written back; the kernels that compute a histogram instead of reading one leave it `None`, which is what they produced before. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): describe the gather test by what this branch does The doc on `an_oversized_selection_is_refused_by_the_pool` said the kernel has no point limit any more. It does on this branch: the limit is removed on its own branch, which this one does not carry, so the comment described a state that has not arrived. What is true here is the order: the estimate now applies to every gather, so the pool is asked before the kernel is, and the pool's refusal is the one the caller sees. That is worth saying on its own -- a bound in bytes on the path that decides the points, rather than a count that says nothing about what they cost -- and it stays true after the limit goes. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): execute vector binary nodes in the native runtime Checkpoint the vector-vector runtime and single shared matching kernel. Keep addressed input delivery, whole-input buffering, nested composition and error cleanup together. The existing arithmetic/comparison implementation remains the only computation path. Local review checkpoint only, not ready for publication: all-targets check passes; the current run has three new fixture failures. Independent review requires current-live workspace resizing, field-local matched storage and kept-value merge budgeting before approval. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): infer the binary result fixture element type Drop a redundant complex type annotation flagged by the actual strict clippy run. The collection and literal assertions are unchanged. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): reserve binary workspace along the shared kernel Remove the separate matching prepass and reserve current live workspace at the existing kernel allocation boundaries. Release per-sample and per-field scratch after ownership ends; charge actual retained rows, values, histograms and errors. Unbudgeted public and set callers bypass accounting without another computation path. Validate with cargo fmt -p promql -- --check; cargo check -p promql --all-targets; cargo nextest run -p promql --retries 0 --no-fail-fast; cargo clippy -p promql --all-targets -- -D warnings. The public nested Binary test also fails on swapped left/right operands and passes after byte-identical restoration. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): scope binary reservations to live scans and output rows Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): charge binary buffers at actual allocation boundaries Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): skip accounting ledger for unbudgeted vector merges Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): simplify native execution memory estimates Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): match binary operands by logical series Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): cover flattened operand descriptor capacity Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): merge rewritten logical series before window evaluation Accept disjoint closed-source fragments under one output label set. Retain occupancy checks and merge affected window inputs once at EOF using shared typed sample segments. Keep direct input windows and root maps streaming, and release collected fragments on finish or error. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): clarify approximate merge workspace estimate Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): validate series boundaries before collected input merges Keep each collecting window's open metric and last position so EOF grouping cannot hide reversed runs or identity changes before closing. Preserve source and root-map piece semantics, and test malformed input release, empty closing chunks, and EOF truncation. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): supply the default binary fill in the unimplemented-node fixture Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): end a subtree before the node above it computes A window that collects its input answers only once its own input is finished, so ending the input now flushes the deepest ready window and lets the following pull carry its answer up; a binary node computes only when both operands are settled. Previously a binary node computed on an empty operand and marked itself computed, dropping what the window produced afterwards. Also stream the answer instead of collecting it: end_of_input() ends the input, and each pull runs the next stage, so the caller pulls the result and the last None closes the execution. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): reserve an instant selection for the points it decides The gather reserved the whole planned grid, while select_instants clips the range to the series' first and last sample and its lookback: a week-long query at a one-minute step reserved ~504 KB for a series that produces five points, and a small pool refused it before gathering. Reserve for the clipped candidate count, derived with the same alignment the kernel uses. Also pin the start-timestamp semantics the histogram change spread to: last_over_time and first_over_time hand back the sample they selected with its start timestamp, and limitk and limit_ratio keep it on the sample they keep, so a reader downstream can still place a counter reset. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): reserve binary memory per phase, not per value The binary kernel adjusted the pool and the metrics once per matched point, per filed row and per merged row, so a million plain float points meant roughly three million reserve calls inside a numeric loop. Reserve at three boundaries instead: the operands as chunks are retained, one figure per phase (pairing, sort, merge) with an estimate over rows/fields/segments/payload, and each result buffer when a block of it is allocated. Delete the code that modelled another component's allocator: index_column_growth and its shadow index_capacity, BinaryOperand::ends_growth, and the flatten-capacity model in assembly_bytes. Nothing pre-reserves the left-times-right product and no large allocation happens before its reservation. Tests that pinned the callback sequence or exact small-object arithmetic go; the resource properties stay: a genuinely large input or output is refused, a sparse small result is admitted, and reservations are released on finish, error, truncation and drop. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): move the histogram_quantile doc back to its own test The start-timestamp test was inserted into the middle of another test's doc comment, leaving the earlier test without one. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): compile a constant, a negation and a scalar-vector operation Three shapes reach the tree now: a constant binds as a `Scalar` on the statement's grid, a negation binds as the multiplication by `-1` it is, and a binary operation between one vector and one constant binds as a `Binary` node whose output is the vector side's own binding. The arithmetic and comparison operators all map, and `return_bool` is carried as its own field so `x > 0` (which filters) and `x > bool 0` (which reports zero or one) are one node with one field changed. The operands keep the order they were written in, because most operators are not commutative: `1 / x` is not `x / 1`, and which side the constant is on is the only thing that tells them apart. A constant is bound by the same `constant_value` fold a call's parameters already use, so `x * (1 + 2)` carries `3` rather than an addition the tree would have to know how to run. Three shapes are refused rather than approximated. An operation between two vectors needs each `Input` to name the selector it reads, which the one binding does not carry, so `a + b` and `a + a` would compile to the same tree; a set operator combines two vectors' series and has no node at all; and a statement whose value is a scalar has nothing to scan, so a constant is only bound next to a vector that gives it something to broadcast over. Each refusal names what it found. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit e672a9c5264422e1d1452d89563d8e68b3632852) (cherry picked from commit 77ea4d1f4eec7fedbdcb80841b5610dc90487ef7) * feat(promql): bind every aggregator that folds a group's values `AggregateKind` named only `Sum`, so the compiler bound `sum` and refused every other aggregator by name. Six of the ones it refused fold exactly as `sum` does — each reports one value per group per grid point — and two more (`stddev`, `stdvar`) plus `quantile` fold the same way without a kernel yet. Binding them is the same decision as binding `sum`: the compiler says what an aggregate means, and the runtime refuses a kind it cannot execute. A kernel landing later then changes no binding. `quantile` is the one of the nine with a parameter, and it carries it in the variant (`Quantile { quantile: f64 }`) as `RangeFunction::QuantileOverTime` does, because `BoundAggregate` has no field a constant could go in and a field that meant nothing for the other eight kinds would split a kind from its own parameter. That costs `AggregateKind` `Eq`, which nothing on this tree needed: the only exhaustive match is the display, and no collection is keyed on it. The five that do not fold are refused with the reason rather than mapped to a kind that would compute something else: `topk`, `bottomk`, `limitk` and `limit_ratio` select some of the input's series, so they do not report one value per group per grid point, and `count_values` writes one output label per distinct value. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit f1d1f3124ede003ae6df18fa0617aa13a56f3907) * test(promql): refuse an interleaved input under a selection when the tree is bound The merged plan validates an input ordering in the arm that reads it, so a selection over an interleaved input no longer binds at all. Assert that refusal where it happens instead of asserting a runtime rejection the tree cannot reach, and keep the map-over-a-selection case: a map above a selection does not shield the check, because the check looks at the input the selection itself reads. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): say what the start timestamp pins, and cover every path that carries it The field is GreptimeDB sample metadata, so the comments now say that instead of arguing from what Prometheus omits: a FloatHistogram there has no such field, while its rate and reset evaluation takes start timestamps through a separate channel, and this crate's native range kernel does not read the field at all. Comparison filtering and set selection were the two paths that carry the timestamp without a test. Both keep the payload of the operand they hand back, including the right side of an or at a point the left side does not cover, and a bool comparison answers with a float so no histogram reaches the output. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * style(promql): drop the blank line the appended test block left behind Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): drop the start timestamp a computed histogram never had A window holding one histogram became the result as it arrived, while two or more went through an add that clears the timestamp, so whether a summed histogram carried the payload said how many samples the window held rather than what the function did. Clear it where the fold returns, leaving the bucket layout to the caller. Also correct two comments that described behaviour this branch does not have: the selection kernel carries no point limit of its own, and the interleaved map sits above the selection rather than between it and the input. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): bill a binary operand the histogram tables it keeps alive A chunk figure leaves out the tables its segments index into, because a table is one allocation however many chunks share it. The operand then billed only the indices, so tables it kept alive cost nothing: a pool holding room for one table admitted a second and a third, each arriving from a source whose own reservation was released as it handed the chunk over. Bill each table once, to the chunk that first references it, and charge the one slot that remembers it. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): execute a complete sum over the groups its input forms A Complete aggregate now runs: it folds every evaluated chunk of its input into per-group states as they arrive, and answers one series per group only after the input as a whole ends. A chunk closing its series says nothing about the group -- the group is complete when the input is -- so the node waits for `inputs_finished` and a settled input rather than for a last chunk. The fold is the one `aggregate_float` already had, split out as `FloatAccumulator` so a chunk at a time folds into the same states a whole series would. Its indices are absolute grid positions, so the split changes no value. Only `Complete` + `Sum` run; every other kind and stage is named in the error. Groups must reach one execution whole, which the caller plans -- this node cannot check it. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(promql): refuse a modifier before folding a constant Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat(promql): bind every aggregator that folds a group's values `AggregateKind` named only `Sum`, so the compiler bound `sum` and refused every other aggregator by name. Six of the ones it refused fold exactly as `sum` does — each reports one value per group per grid point — and two more (`stddev`, `stdvar`) plus `quantile` fold the same way without a kernel yet. Binding them is the same decision as binding `sum`: the compiler says what an aggregate means, and the runtime refuses a kind it cannot execute. A kernel landing later then changes no binding. `quantile` is the one of the nine with a parameter, and it carries it in the variant (`Quantile { quantile: f64 }`) as `RangeFunction::QuantileOverTime` does, because `BoundAggregate` has no field a constant could go in and a field that meant nothing for the other eight kinds would split a kind from its own parameter. That costs `AggregateKind` `Eq`, which nothing on this tree needed: the only exhaustive match is the display, and no collection is keyed on it. The five that do not fold are refused with the reason rather than mapped to a kind that would compute something else: `topk`, `bottomk`, `limitk` and `limit_ratio` select some of the input's series, so they do not report one value per group per grid point, and `count_values` writes one output label per distinct value. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> (cherry picked from commit f1d1f3124ede003ae6df18fa0617aa13a56f3907) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> # Conflicts: # src/query/src/promql/native_compile.rs # src/query/src/promql/native_compile/tests.rs * fix(promql): admit aggregate state before allocating it Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(promql): keep aggregate byte accounting incremental Scanning every group and key on each fold made a high-cardinality aggregate quadratic in the number of chunks. Track the retained figure as groups, blocks and container backing grow, and keep a full recount behind cfg(test) to check the two against each other. Also time the finalize and handover of an aggregate node, which ran outside its self compute time. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): check the initial aggregate accounting too Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): use slices for the fixed aggregate test values Five `&vec![...]` literals where a slice does the same job; clippy fails the workspace build on them with -D warnings. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(promql): bound the aggregate containers instead of modelling them The byte model for the group index and the state B-tree was precise about bucket counts, control bytes and node layout, which pinned it to one toolchain and left the tests unable to notice it going wrong: a model that returned zero still passed. At a million groups the map is about 2% of what a fold retains and the B-tree nodes about 11%; the block states are the rest. A coarse bound that cannot under-count is therefore a small price, and it costs 45 fewer lines. Three buckets per entry (16/7 of capacity, rounded up to a power of two, plus one control byte) and one node per five entries, sized from the pair and child-pointer types rather than a field-by-field layout. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/promql/src/aggregate.rs | 517 +++++++++++- src/promql/src/aggregate/tests.rs | 317 ++++++++ src/promql/src/grid_series.rs | 150 +++- src/promql/src/native/aggregate.rs | 243 ++++++ src/promql/src/native/execute.rs | 322 +++++--- src/promql/src/native/execute/tests.rs | 793 ++++++++++++++++++- src/promql/src/native/mod.rs | 1 + src/query/src/promql/native_compile/tests.rs | 525 +++++++++++- 8 files changed, 2702 insertions(+), 166 deletions(-) create mode 100644 src/promql/src/native/aggregate.rs diff --git a/src/promql/src/aggregate.rs b/src/promql/src/aggregate.rs index 228f67ba4bc..0ab1e26a291 100644 --- a/src/promql/src/aggregate.rs +++ b/src/promql/src/aggregate.rs @@ -27,11 +27,14 @@ pub use state::{ use std::cmp::Reverse; use std::collections::hash_map::Entry; use std::collections::{BTreeMap, BinaryHeap, HashMap}; +use std::mem::size_of; use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; +use datatypes::arrow::array::Float64Array; use datatypes::series::MetricName; -use datatypes::series::samples::{SampleRef, Series, SeriesSchema}; +use datatypes::series::samples::{SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::{GridSpan, TimeAxis}; use crate::go_format::format_float_f; use crate::grid_series::{GridSeriesBuilder, external}; @@ -97,6 +100,7 @@ pub struct FloatAggregator { labels: Vec>, group_index: HashMap, groups: Vec, + retained: usize, } impl FloatAggregator { @@ -104,25 +108,78 @@ impl FloatAggregator { schema: &Arc, op: AggregateOp, grouping: Grouping<'_>, + ) -> Result { + Self::try_new_budgeted(schema, op, grouping, |_| Ok(())) + } + + pub(crate) fn try_new_budgeted( + schema: &Arc, + op: AggregateOp, + grouping: Grouping<'_>, + mut budget: impl FnMut(usize) -> Result<()>, ) -> Result { validate_input(schema, &[], grouping)?; let (by, labels) = match grouping { Grouping::By(labels) => (true, labels), Grouping::Without(labels) => (false, labels), }; + let label_storage = + grouping_bytes(grouping).saturating_add(labels_capacity_bytes(grouping)); + budget(label_storage)?; + let mut owned_labels: Vec> = Vec::new(); + owned_labels + .try_reserve_exact(labels.len()) + .map_err(|error| { + DataFusionError::Execution(format!( + "cannot reserve aggregate grouping labels: {error}" + )) + })?; + owned_labels.extend(labels.iter().map(|label| Arc::from(*label))); + let retained = owned_labels + .capacity() + .saturating_mul(size_of::>()) + .saturating_add(owned_labels.iter().fold(0usize, |bytes, label| { + bytes.saturating_add(label.len() + 2 * size_of::()) + })); Ok(Self { schema: schema.clone(), op, by, - labels: labels.iter().map(|label| Arc::from(*label)).collect(), + labels: owned_labels, group_index: HashMap::new(), groups: Vec::new(), + retained, }) } /// Folds `series` into its group. A series that fails validation leaves the state unchanged. + /// + /// A budget refusal is not a validation failure: the fold may have opened a group or filled + /// blocks before the allocation it could not admit, so a failed [`Self::update_budgeted`] + /// leaves the state half folded and its retained figure stale. The runtime stops the execution + /// and drops the aggregator on that error, which is why no rollback is kept; a caller holding + /// one must drop it rather than fold into it again. pub fn update(&mut self, series: &Series) -> Result<()> { + self.update_budgeted(series, |_| Ok(())) + } + + pub(crate) fn update_budgeted( + &mut self, + series: &Series, + mut budget: impl FnMut(usize) -> Result<()>, + ) -> Result<()> { validate_float_series(&self.schema, series)?; + let existing_bytes = self.state_bytes(); + let grouping_workspace = grouping_workspace_bytes(&series.metric); + let mut workspace_bytes = + grouping_workspace.saturating_add(pending_bytes(self.schema.fields.len())); + budget(existing_bytes.saturating_add(workspace_bytes))?; + workspace_bytes = workspace_bytes.saturating_add( + self.labels + .len() + .saturating_mul(std::mem::size_of::<&str>()), + ); + budget(existing_bytes.saturating_add(workspace_bytes))?; let labels: Vec<&str> = self.labels.iter().map(|label| &**label).collect(); let grouping = if self.by { Grouping::By(&labels) @@ -130,39 +187,210 @@ impl FloatAggregator { Grouping::Without(&labels) }; let metric = grouping.apply(&series.metric)?; - let index = match self.group_index.entry(label_key(&metric)) { + let key_bytes = metric + .tags() + .len() + .saturating_add(1) + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()); + workspace_bytes = workspace_bytes + .saturating_add(grouping_workspace_bytes(&metric)) + .saturating_add(key_bytes); + budget(existing_bytes.saturating_add(workspace_bytes))?; + let key = label_key(&metric); + let key_retained = key + .capacity() + .saturating_mul(size_of::<(Arc, Arc)>()); + let new_group = !self.group_index.contains_key(&key); + let group_capacity = self.groups.capacity(); + let map_capacity = self.group_index.capacity(); + let old_map_backing = hash_map_backing_bytes(map_capacity); + if new_group { + let next_group_capacity = if self.groups.len() == group_capacity { + group_capacity.saturating_mul(2).max(4) + } else { + group_capacity + }; + let group_realloc_bytes = if next_group_capacity > group_capacity { + next_group_capacity.saturating_mul(size_of::()) + } else { + 0 + }; + budget( + existing_bytes + .saturating_add(metric_bytes(&metric)) + .saturating_add(key_bytes) + // A realloc that cannot grow in place keeps both buffers alive while it + // copies, so a growing step admits the old buffer beside the new one -- + // the same reason the map below admits both tables. The old buffer alone + // is already in `existing_bytes`; what is added here is the new one. + .saturating_add(group_realloc_bytes) + .saturating_add(if self.group_index.len() == map_capacity { + hash_map_backing_bytes(next_hash_capacity(map_capacity)) + } else { + 0 + }) + .saturating_add(workspace_bytes), + )?; + } + let group_capacity_growth = if new_group && self.groups.len() == group_capacity { + group_capacity.saturating_mul(2).max(4) + } else { + group_capacity + }; + if new_group && group_capacity_growth > group_capacity { + self.groups.reserve(1); + } + if new_group && self.group_index.len() == self.group_index.capacity() { + self.group_index.reserve(1); + } + let index = match self.group_index.entry(key) { Entry::Occupied(entry) => *entry.get(), Entry::Vacant(entry) => { self.groups.push(Group { - metric, + metric: metric.clone(), blocks: BTreeMap::new(), }); - *entry.insert(self.groups.len() - 1) + let index = self.groups.len() - 1; + entry.insert(index); + index } }; - self.groups[index].accumulate(series, self.op); + if new_group { + self.retained = self + .retained + .saturating_add(metric_bytes(&metric)) + .saturating_add(key_retained) + .saturating_add( + self.groups + .capacity() + .saturating_sub(group_capacity) + .saturating_mul(size_of::()), + ) + .saturating_add( + hash_map_backing_bytes(self.group_index.capacity()) + .saturating_sub(old_map_backing), + ); + } + drop(labels); + drop(metric); + self.retained = self.groups[index].accumulate( + series, + self.op, + &mut budget, + self.retained, + self.schema.fields.len(), + )?; Ok(()) } + /// Number of groups currently retained by this fold. + pub(crate) fn groups(&self) -> usize { + self.groups.len() + } + + /// Bytes reserved for this fold, using bounded allowances for its maps and group payloads. + pub(crate) fn state_bytes(&self) -> usize { + self.retained + } + + #[cfg(test)] + pub(crate) fn recounted_state_bytes(&self) -> usize { + let groups = self.groups.iter().fold(0usize, |total, group| { + total + .saturating_add(metric_bytes(&group.metric)) + .saturating_add( + group + .blocks + .len() + .saturating_mul(BLOCK_POINTS) + .saturating_mul(self.schema.fields.len()) + .saturating_mul(std::mem::size_of::()), + ) + .saturating_add(btree_backing_bytes(group.blocks.len())) + }); + groups + .saturating_add( + self.labels + .capacity() + .saturating_mul(std::mem::size_of::>()), + ) + .saturating_add(hash_map_backing_bytes(self.group_index.capacity())) + .saturating_add(self.group_index.keys().fold(0usize, |bytes, key| { + bytes.saturating_add( + key.capacity() + .saturating_mul(size_of::<(Arc, Arc)>()), + ) + })) + .saturating_add( + self.labels + .iter() + .map(|label| label.len() + 2 * std::mem::size_of::()) + .sum::(), + ) + .saturating_add( + self.groups + .capacity() + .saturating_mul(std::mem::size_of::()), + ) + } + /// The result of every group, in the order the groups first appeared. pub fn finish(self) -> Result> { - finish_groups(self.schema, self.op, self.groups) + self.finish_budgeted(|_| Ok(())) + } + + pub(crate) fn finish_budgeted( + self, + budget: impl FnMut(usize) -> Result<()>, + ) -> Result> { + let state = self.state_bytes(); + finish_groups_budgeted(self.schema, self.op, self.groups, state, budget) } } -/// Turns the state of `groups` into result series, in group order. -fn finish_groups( +/// Turns the state of `groups` into result series, in group order, admitting as it builds. +fn finish_groups_budgeted( schema: Arc, op: AggregateOp, groups: Vec, + state: usize, + mut budget: impl FnMut(usize) -> Result<()>, ) -> Result> { let width = schema.fields.len(); + let group_count = groups.len(); + if group_count == 0 { + return Ok(Vec::new()); + } + let groups = groups.into_iter(); let mut builder = GridSeriesBuilder::try_new(schema)?; - let mut output = Vec::with_capacity(groups.len()); - let mut row = Vec::with_capacity(width); + let output_capacity_bytes = group_count.saturating_mul(std::mem::size_of::()); + let row_capacity_bytes = width.saturating_mul(std::mem::size_of::>()); + let base = state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes); + budget(base)?; + let mut output = Vec::new(); + output.try_reserve_exact(group_count).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve aggregate outputs: {error}")) + })?; + let mut row = Vec::new(); + row.try_reserve_exact(width).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve aggregate row: {error}")) + })?; + let mut completed_output_bytes = 0usize; for group in groups { // One builder serves every group, so each result series starts explicitly. builder.begin_series(); + let (rows, runs) = group_rows_runs(&group, width); + let finish_bytes = finish_working_bytes(rows, runs, width); + budget( + state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes) + .saturating_add(completed_output_bytes) + .saturating_add(finish_bytes), + )?; + builder.reserve_rows(rows)?; for (&block, states) in &group.blocks { // The last block can extend past the grid; only points with a sample are pushed. for (offset, point) in states.chunks_exact(width).enumerate() { @@ -174,11 +402,89 @@ fn finish_groups( builder.push(block * BLOCK_POINTS + offset, &row)?; } } - output.extend(builder.finish(group.metric)?); + let finished = builder.finish(group.metric)?; + completed_output_bytes = completed_output_bytes.saturating_add( + finished + .iter() + .map(crate::grid_series::series_bytes) + .sum::(), + ); + output.extend(finished); + budget( + state + .saturating_add(output_capacity_bytes) + .saturating_add(row_capacity_bytes) + .saturating_add(completed_output_bytes), + )?; } Ok(output) } +/// The result of every group, without a budget: the entry the partial-state merger takes, which +/// does not hold a query pool to admit against. +fn finish_groups( + schema: Arc, + op: AggregateOp, + groups: Vec, +) -> Result> { + finish_groups_budgeted(schema, op, groups, 0, |_| Ok(())) +} + +fn group_rows_runs(group: &Group, width: usize) -> (usize, usize) { + if width == 0 { + return (0, 0); + } + let mut rows = 0usize; + let mut runs = 0usize; + let mut previous: Option = None; + for (&block, states) in &group.blocks { + for (offset, point) in states.chunks_exact(width).enumerate() { + if point.iter().any(|state| state.seen) { + let index = block.saturating_mul(BLOCK_POINTS).saturating_add(offset); + rows = rows.saturating_add(1); + if previous.is_none_or(|previous| previous.checked_add(1) != Some(index)) { + runs = runs.saturating_add(1); + } + previous = Some(index); + } + } + } + (rows, runs) +} + +fn finish_working_bytes(rows: usize, runs: usize, width: usize) -> usize { + let values = rows + .saturating_mul(width) + .saturating_mul(std::mem::size_of::()); + let validity_bytes = rows + .saturating_add(7) + .checked_div(8) + .unwrap_or(0) + .div_ceil(64) + .saturating_mul(64); + let fields = crate::grid_series::field_builder_bytes(width).saturating_add( + width.saturating_mul( + std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(std::mem::size_of::()), + ), + ); + let indices = rows.saturating_mul(std::mem::size_of::()); + let spans = runs.saturating_mul(std::mem::size_of::()); + let axis_conversion = runs + .saturating_mul( + 2 * std::mem::size_of::() + std::mem::size_of::>(), + ) + .saturating_add(rows.saturating_mul(std::mem::size_of::())); + values + .saturating_add(width.saturating_mul(validity_bytes)) + .saturating_add(fields) + .saturating_add(indices) + .saturating_add(spans) + .saturating_add(axis_conversion) + .saturating_add(std::mem::size_of::() + std::mem::size_of::()) +} + /// `quantile(q, v)`, with `q` read at each grid point from an unlabelled single-field scalar on /// the same grid, as Prometheus evaluates the parameter at every step. A point where `q` has no /// value has no result; `q` outside `[0, 1]` or NaN gives what Prometheus' `quantile` gives. @@ -438,6 +744,105 @@ impl Grouping<'_> { type LabelKey = Vec<(Arc, Arc)>; +fn metric_bytes(metric: &MetricName) -> usize { + metric + .name() + .map_or(0, str::len) + .saturating_add( + metric + .tags() + .iter() + .map(|tag| tag.name.len() + tag.value.len()) + .sum::(), + ) + .saturating_add( + metric + .tags() + .len() + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()), + ) + .saturating_add(std::mem::size_of::()) +} + +fn grouping_workspace_bytes(metric: &MetricName) -> usize { + metric_bytes(metric) + .saturating_add( + metric + .tags() + .len() + .saturating_add(1) + .saturating_mul(std::mem::size_of::<(Arc, Arc)>()), + ) + .saturating_add(std::mem::size_of::()) +} + +/// Conservative backing estimate for Rust 1.98.1's insert-only `HashMap` (hashbrown 0.17.1). +/// Bucket count follows its 7/8 load factor; entries, one control byte per bucket and the +/// 16-byte control-group tail are counted. This models requested table storage, not allocator slack. +/// An upper bound on the bytes the group index holds for `capacity` entries. +/// +/// The table keeps its entries in a power-of-two number of buckets below `16/7` of the usable +/// capacity, with one control byte beside each bucket. Three buckets per entry therefore covers the +/// table whatever exact growth the map uses, which is what keeps this a bound rather than a model of +/// one implementation: a toolchain that changes how the map grows cannot make it under-count. +fn hash_map_backing_bytes(capacity: usize) -> usize { + capacity + .saturating_mul(3) + .saturating_mul(size_of::<(LabelKey, usize)>() + 1) +} + +/// An upper bound on the capacity a map holding `capacity` entries grows to. +fn next_hash_capacity(capacity: usize) -> usize { + capacity.saturating_mul(2).max(7) +} + +/// An upper bound on the bytes one B-tree node holds. +/// +/// A node keeps at most eleven key/value pairs and twelve child pointers, plus a header of a few +/// words; the leaf form is smaller than that. The header allowance is generous on purpose, so the +/// bound does not depend on how a toolchain lays the node out. +fn btree_node_bytes() -> usize { + 11 * size_of::<(usize, Box<[FloatState]>)>() + 16 * size_of::() +} + +/// An upper bound on the bytes a B-tree of `entries` holds. +/// +/// Every node but the root keeps at least five entries, so `entries / 5` nodes is enough to hold +/// them; a tree of one entry is one node. +fn btree_backing_bytes(entries: usize) -> usize { + if entries == 0 { + return 0; + } + entries.div_ceil(5).saturating_mul(btree_node_bytes()) +} + +fn pending_bytes(fields: usize) -> usize { + working_set_bytes(fields, BLOCK_POINTS) +} + +fn working_set_bytes(fields: usize, pending_capacity: usize) -> usize { + pending_capacity + .saturating_mul(std::mem::size_of::<(usize, f64)>()) + .saturating_add(fields.saturating_mul(std::mem::size_of::>() + 1)) +} + +fn labels_capacity_bytes(grouping: Grouping<'_>) -> usize { + let labels = match grouping { + Grouping::By(labels) | Grouping::Without(labels) => labels, + }; + labels.len().saturating_mul(std::mem::size_of::>()) +} + +fn grouping_bytes(grouping: Grouping<'_>) -> usize { + let labels = match grouping { + Grouping::By(labels) | Grouping::Without(labels) => labels, + }; + labels + .iter() + .map(|label| label.len() + std::mem::size_of::>()) + .sum::() +} + fn label_key(metric: &MetricName) -> LabelKey { let mut labels = Vec::with_capacity(metric.tags().len() + 1); if let Some(name) = metric.name() { @@ -463,9 +868,29 @@ struct Group { blocks: BTreeMap>, } +/// Which field of which layout a fold is writing, and how it folds it. +/// +/// One value carried through the per-block work instead of three separate arguments: the three +/// never vary independently within one series, and the block loop is the only place that reads +/// them. +#[derive(Clone, Copy)] +struct FieldFold { + width: usize, + field_index: usize, + op: AggregateOp, +} + impl Group { - fn accumulate(&mut self, series: &Series, op: AggregateOp) { - let width = series.fields.len(); + fn accumulate( + &mut self, + series: &Series, + op: AggregateOp, + budget: &mut impl FnMut(usize) -> Result<()>, + mut retained: usize, + width: usize, + ) -> Result { + let pending_capacity_bytes = BLOCK_POINTS.saturating_mul(size_of::<(usize, f64)>()); + budget(retained.saturating_add(pending_capacity_bytes))?; let mut pending = Vec::with_capacity(BLOCK_POINTS); for (field_index, field) in series.fields.iter().enumerate() { // The axis says which absolute positions the field holds and the field iterator @@ -475,35 +900,81 @@ impl Group { let mut block = usize::MAX; for ((index, _), sample) in series.time.positions().zip(field.iter()) { if index / BLOCK_POINTS != block { - self.update_block(block, width, field_index, op, &mut pending); + retained = self.update_block( + block, + FieldFold { + width, + field_index, + op, + }, + &mut pending, + budget, + retained, + )?; block = index / BLOCK_POINTS; } if let SampleRef::Float(sample) = sample { pending.push((index % BLOCK_POINTS, sample)); } } - self.update_block(block, width, field_index, op, &mut pending); + retained = self.update_block( + block, + FieldFold { + width, + field_index, + op, + }, + &mut pending, + budget, + retained, + )?; } + Ok(retained) } /// Applies one field's samples within a block, so a block costs one lookup per field run. fn update_block( &mut self, block: usize, - width: usize, - field_index: usize, - op: AggregateOp, + fold: FieldFold, pending: &mut Vec<(usize, f64)>, - ) { + budget: &mut impl FnMut(usize) -> Result<()>, + retained: usize, + ) -> Result { if pending.is_empty() { - return; + return Ok(retained); + } + let new_block = !self.blocks.contains_key(&block); + let block_bytes = if new_block { + BLOCK_POINTS + .saturating_mul(fold.width) + .saturating_mul(std::mem::size_of::()) + } else { + 0 + }; + let backing_delta = if new_block { + btree_backing_bytes(self.blocks.len() + 1) + .saturating_sub(btree_backing_bytes(self.blocks.len())) + } else { + 0 + }; + if new_block { + budget( + retained + .saturating_add(backing_delta) + .saturating_add(block_bytes) + .saturating_add(pending.capacity().saturating_mul(size_of::<(usize, f64)>())), + )?; } let states = self.blocks.entry(block).or_insert_with(|| { - vec![FloatState::default(); BLOCK_POINTS * width].into_boxed_slice() + vec![FloatState::default(); BLOCK_POINTS * fold.width].into_boxed_slice() }); for (offset, sample) in pending.drain(..) { - states[offset * width + field_index].update(op, sample); + states[offset * fold.width + fold.field_index].update(fold.op, sample); } + Ok(retained + .saturating_add(backing_delta) + .saturating_add(block_bytes)) } } diff --git a/src/promql/src/aggregate/tests.rs b/src/promql/src/aggregate/tests.rs index 31082bd784a..0fd86561fdd 100644 --- a/src/promql/src/aggregate/tests.rs +++ b/src/promql/src/aggregate/tests.rs @@ -14,6 +14,7 @@ use std::collections::BTreeMap; +use datafusion::error::DataFusionError; use datatypes::arrow::array::{Float64Array, UInt32Array}; use datatypes::series::Label; use datatypes::series::histogram::HistogramBatchBuilder; @@ -345,6 +346,27 @@ fn avg_compensation_matches_prometheus_bits() { } } +#[test] +fn one_field_missing_and_another_nan_keep_row_nullability() { + let schema = schema(&["missing", "value"]); + let series = input( + &schema, + metric("m", &[]), + 0, + &[ + vec![None, None, None], + vec![Some(1.0), Some(f64::NAN), Some(2.0)], + ], + ); + let output = aggregate_float(&schema, &[series], AggregateOp::Sum, Grouping::By(&[])).unwrap(); + assert_eq!(output.len(), 1); + assert_eq!(values(&output[0], 0).1, vec![None, None, None]); + let values = values(&output[0], 1).1; + assert_eq!(values[0], Some(1.0)); + assert!(values[1].unwrap().is_nan()); + assert_eq!(values[2], Some(2.0)); +} + #[test] fn fields_are_aggregated_independently() { // A field without input values at a row is null there, even for count and group; the row @@ -476,6 +498,301 @@ fn fields_are_aggregated_independently() { } } +#[test] +fn incremental_retained_bytes_match_full_recount_after_updates() { + let schema = schema_with_grid_len(&["value", "other"], 256); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["zone", "rack"])) + .unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "initial state" + ); + let updates = [ + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 0, + &[vec![Some(1.0)], vec![Some(2.0)]], + ), + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 64, + &[vec![Some(3.0)], vec![Some(4.0)]], + ), + input( + &schema, + metric("m", &[("zone", "a"), ("rack", "1")]), + 0, + &[vec![Some(5.0)], vec![Some(6.0)]], + ), + input( + &schema, + metric("m", &[("zone", "b"), ("rack", "2")]), + 128, + &[vec![Some(7.0)], vec![Some(8.0)]], + ), + ]; + for (index, series) in updates.iter().enumerate() { + let previous = aggregator.state_bytes(); + aggregator.update(series).unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "update {index}" + ); + if index == 2 { + assert_eq!(aggregator.state_bytes(), previous); + } + } +} + +#[test] +fn incremental_retained_bytes_match_full_recount_for_many_groups() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["id"])).unwrap(); + for id in 0..512 { + let id = id.to_string(); + let series = input( + &schema, + metric("m", &[("id", &id)]), + id.parse::().unwrap() % 64, + &[vec![Some(1.0)]], + ); + aggregator.update(&series).unwrap(); + assert_eq!( + aggregator.state_bytes(), + aggregator.recounted_state_bytes(), + "group {id}" + ); + } +} + +/// The unbudgeted entry point is the budgeted one with a no-op budget, so this compares one path +/// with itself: it is a smoke check that the wrapper still runs, not evidence that the kernels are +/// right. The numbers are checked against Prometheus in `aggregates_match_prometheus` and +/// `stat_aggregates_match_prometheus`. +#[test] +fn every_supported_kind_runs_through_the_unbudgeted_entry_point() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0), Some(3.0)]]); + for op in [ + AggregateOp::Sum, + AggregateOp::Count, + AggregateOp::Avg, + AggregateOp::Min, + AggregateOp::Max, + AggregateOp::Group, + AggregateOp::Stddev, + AggregateOp::Stdvar, + ] { + let mut plain = FloatAggregator::try_new(&schema, op, Grouping::By(&[])).unwrap(); + plain.update(&series).unwrap(); + let plain = plain.finish().unwrap(); + let mut budgeted = + FloatAggregator::try_new_budgeted(&schema, op, Grouping::By(&[]), |_| Ok(())).unwrap(); + budgeted.update_budgeted(&series, |_| Ok(())).unwrap(); + let budgeted = budgeted.finish_budgeted(|_| Ok(())).unwrap(); + assert_eq!( + values_of(&plain, 0, 128), + values_of(&budgeted, 0, 128), + "{op:?}" + ); + } +} + +#[test] +fn repeated_updates_in_one_group_and_block_keep_one_block_charge() { + let schema = schema(&["value"]); + let first = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0)]]); + let second = input(&schema, metric("m", &[]), 1, &[vec![Some(3.0)]]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + aggregator.update(&first).unwrap(); + let retained = aggregator.state_bytes(); + aggregator.update(&second).unwrap(); + assert_eq!(aggregator.state_bytes(), retained); +} + +#[test] +fn adjacent_blocks_are_admitted_and_accounted_separately() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 63, &[vec![Some(1.0), Some(2.0)]]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + let mut admissions = Vec::new(); + aggregator + .update_budgeted(&series, |bytes| { + admissions.push(bytes); + Ok(()) + }) + .unwrap(); + assert_eq!(aggregator.groups(), 1); + assert_eq!(aggregator.groups[0].blocks.len(), 2); + assert_eq!( + admissions.last().copied(), + Some( + aggregator + .state_bytes() + .saturating_add(std::mem::size_of::<(usize, f64)>() * BLOCK_POINTS,) + ) + ); +} + +#[test] +fn a_rejected_block_allocation_does_not_insert_a_block() { + let schema = schema(&["value"]); + let series = input(&schema, metric("m", &[]), 0, &[vec![Some(1.0)]]); + // The figure a new block's admission adds: the block's states and its pending buffer, on + // top of whatever the fold already holds. Located by the value, not by which call reaches it, + // so reordering the earlier admissions cannot make this test pass without testing anything. + let block_states = BLOCK_POINTS + .saturating_mul(schema.fields.len()) + .saturating_mul(std::mem::size_of::()); + let pending = BLOCK_POINTS * std::mem::size_of::<(usize, f64)>(); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&[])).unwrap(); + let mut refused = None; + let result = aggregator.update_budgeted(&series, |bytes| { + // Everything up to the block's own admission is admitted; only the block is refused. + if bytes >= block_states.saturating_add(pending) { + refused = Some(bytes); + return Err(DataFusionError::Execution("block allocation budget".into())); + } + Ok(()) + }); + assert!( + result.is_err(), + "the block was admitted with no budget left" + ); + let refused = refused.expect("the refusal did not come from a block admission"); + assert!( + refused >= block_states.saturating_add(pending), + "the refusal came from an admission smaller than a block" + ); + assert_eq!(aggregator.groups(), 1); + assert!( + aggregator.groups[0].blocks.is_empty(), + "a refused block was inserted anyway" + ); +} + +#[test] +fn a_map_growth_rejection_does_not_insert_the_new_group() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["i"])).unwrap(); + let first = input(&schema, metric("m", &[("i", "0")]), 0, &[vec![Some(1.0)]]); + aggregator.update(&first).unwrap(); + let mut index = 1; + while aggregator.group_index.len() < aggregator.group_index.capacity() { + let label = index.to_string(); + let series = input( + &schema, + metric("m", &[("i", &label)]), + 0, + &[vec![Some(1.0)]], + ); + aggregator.update(&series).unwrap(); + index += 1; + } + let group_count = aggregator.groups(); + let old_state = aggregator.state_bytes(); + assert_eq!(aggregator.group_index.capacity(), group_count); + + // A resize keeps the old table while it fills the new one, so the figure the growth has to + // admit is the whole new table on top of the old state -- not the difference between them. + // That is the property asserted here: the refusal threshold is the state plus a whole new + // table, so the test fails if the admission ever drops to the delta. + let new_table = hash_map_backing_bytes(next_hash_capacity(group_count)); + let threshold = old_state.saturating_add(new_table); + let label = index.to_string(); + let series = input( + &schema, + metric("m", &[("i", &label)]), + 0, + &[vec![Some(1.0)]], + ); + let mut peak = 0; + let result = aggregator.update_budgeted(&series, |bytes| { + peak = peak.max(bytes); + if bytes >= threshold { + Err(DataFusionError::Execution("old+new map budget".into())) + } else { + Ok(()) + } + }); + assert!( + result.is_err(), + "the map grew with no budget for both tables" + ); + assert!( + peak >= threshold, + "the growth admitted {peak}, less than the old state plus a new table ({threshold})" + ); + assert_eq!(aggregator.group_index.capacity(), group_count); + assert_eq!(aggregator.groups(), group_count); + assert_eq!( + aggregator + .groups + .iter() + .map(|group| group.blocks.len()) + .sum::(), + group_count + ); +} + +#[test] +fn finish_rejection_after_first_group_is_budgeted_before_the_next() { + let schema = schema(&["value"]); + let mut aggregator = + FloatAggregator::try_new(&schema, AggregateOp::Sum, Grouping::By(&["i"])).unwrap(); + for (label, value) in [("a", 1.0), ("b", 2.0)] { + let series = input( + &schema, + metric("m", &[("i", label)]), + 0, + &[vec![Some(value)]], + ); + aggregator.update(&series).unwrap(); + } + let mut calls = Vec::new(); + let result = aggregator.finish_budgeted(|bytes| { + calls.push(bytes); + if calls.len() == 4 { + Err(DataFusionError::Execution("budget".into())) + } else { + Ok(()) + } + }); + assert!(result.is_err()); + assert_eq!(calls.len(), 4); + assert!(calls[1] > calls[0], "first builder peak"); + assert!(calls[2] < calls[1], "first group settle"); + assert!(calls[3] > calls[2], "second builder peak"); + + // The base already includes the output Vec backing; later callbacks add only the + // builder workspace or completed Series allocation. + let builder_workspace = finish_working_bytes(1, 1, 1); + let expected_output = aggregate_float( + &schema, + &[ + input(&schema, metric("m", &[("i", "a")]), 0, &[vec![Some(1.0)]]), + input(&schema, metric("m", &[("i", "b")]), 0, &[vec![Some(2.0)]]), + ], + AggregateOp::Sum, + Grouping::By(&["i"]), + ) + .unwrap(); + let completed_series = crate::grid_series::series_bytes(&expected_output[0]); + assert_eq!(calls[1] - calls[0], builder_workspace); + assert_eq!(calls[2] - calls[0], completed_series); + assert_eq!(calls[3] - calls[2], builder_workspace); +} + #[test] fn state_blocks_keep_absolute_indices_across_boundaries_and_gaps() { // A 512-point grid spans eight state blocks; inputs straddle a block edge, leave whole diff --git a/src/promql/src/grid_series.rs b/src/promql/src/grid_series.rs index c9d8bc19d1f..90cfa0f2e62 100644 --- a/src/promql/src/grid_series.rs +++ b/src/promql/src/grid_series.rs @@ -19,12 +19,12 @@ use std::collections::hash_map::Entry; use std::sync::Arc; use datafusion::arrow::array::builder::NullBufferBuilder; -use datafusion::arrow::array::{Float64Array, UInt32Array}; +use datafusion::arrow::array::{Array, Float64Array, UInt32Array}; use datafusion::error::{DataFusionError, Result}; use datatypes::series::MetricName; -use datatypes::series::histogram::HistogramBatchBuilder; +use datatypes::series::histogram::{HistogramBatch, HistogramBatchBuilder}; use datatypes::series::samples::{ - HistogramSamples, SampleRef, Samples, Series, SeriesFields, SeriesSchema, + DenseSamplesRef, HistogramSamples, SampleRef, Samples, Series, SeriesFields, SeriesSchema, }; use datatypes::series::time::{EvalGrid, GridSpan, TimeAxis}; @@ -35,6 +35,70 @@ pub(crate) mod reference; #[cfg(test)] mod tests; +pub(crate) fn series_bytes(series: &Series) -> usize { + let runs = series.time.runs().count(); + let segments: usize = series + .fields + .iter() + .map(|field| field.segments().len()) + .sum(); + std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(runs.saturating_mul(std::mem::size_of::<(usize, usize)>())) + .saturating_add( + segments.saturating_mul(std::mem::size_of::() + std::mem::size_of::()), + ) + .saturating_add(fields_bytes(&series.fields)) + .saturating_add(series.metric.name().map_or(0, str::len)) + .saturating_add(std::mem::size_of_val(series.metric.tags())) + .saturating_add( + series + .metric + .tags() + .iter() + .map(|tag| tag.name.len() + tag.value.len()) + .sum::(), + ) +} + +pub(crate) fn fields_bytes(fields: &SeriesFields) -> usize { + fields + .iter() + .map(|field| { + field.segments().iter().fold(0usize, |total, segment| { + total.saturating_add(segment_bytes(segment)) + }) + }) + .fold(0usize, usize::saturating_add) +} + +/// The bytes one dense segment of samples holds. +pub(crate) fn segment_bytes(segment: &Samples) -> usize { + match segment.dense() { + Some(DenseSamplesRef::Float(values)) => values.get_array_memory_size(), + Some(DenseSamplesRef::Histogram(histograms)) => histograms + .indices() + .get_array_memory_size() + .saturating_add(table_bytes(histograms.batch())), + Some(DenseSamplesRef::Mixed(mixed)) => mixed + .floats() + .get_array_memory_size() + .saturating_add(mixed.histograms().indices().get_array_memory_size()) + .saturating_add(table_bytes(mixed.histograms().batch())), + // `segments` hands out dense segments only, so a nested chunk never reaches here. + None => 0, + } +} + +/// The bytes one histogram table holds: the batch struct and the arrays its columns view. +pub(crate) fn table_bytes(table: &Arc) -> usize { + std::mem::size_of::().saturating_add(table.as_arrow().get_array_memory_size()) +} + +pub(crate) fn field_builder_bytes(width: usize) -> usize { + width.saturating_mul(std::mem::size_of::()) +} + /// Builds one grid series from rows pushed at absolute grid indices. /// /// A push names one position and the value each field holds there. Rows are accumulated @@ -114,6 +178,36 @@ impl GridSeriesBuilder { }) } + /// Reserves capacity for the exact number of valued rows in the next series. + /// + /// Call after `begin_series` and before its first push. Existing public callers may continue + /// to rely on normal growth; callers with a budget can reserve the concrete row capacity first. + pub(crate) fn reserve_rows(&mut self, rows: usize) -> Result<()> { + self.indices + .try_reserve_exact(rows.saturating_sub(self.indices.len())) + .map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid row indices: {error}")) + })?; + if rows == 0 { + return Ok(()); + } + if self.fields.is_empty() { + self.fields = (0..self.schema.fields.len()) + .map(|_| FieldBuilder::new(rows)) + .collect(); + } + for field in &mut self.fields { + field + .values + .try_reserve_exact(rows.saturating_sub(field.values.len())) + .map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid values: {error}")) + })?; + field.nulls = NullBufferBuilder::new(rows); + } + Ok(()) + } + /// Pushes the field values at absolute grid `index`; `None` is a missing value. /// /// Indices must be strictly increasing; skipped indices are positions that hold no value. @@ -176,7 +270,7 @@ impl GridSeriesBuilder { return Ok(false); } if self.fields.is_empty() { - self.fields = (0..len).map(|_| FieldBuilder::new()).collect(); + self.fields = (0..len).map(|_| FieldBuilder::new(0)).collect(); } self.indices.push(index); Ok(true) @@ -239,24 +333,28 @@ impl GridSeriesBuilder { } let fields = std::mem::take(&mut self.fields); let indices = std::mem::take(&mut self.indices); - let samples = fields - .into_iter() - .map(|mut field| { - let floats = Float64Array::new(field.values.into(), field.nulls.finish()); - let Some(column) = field.histograms else { - return Ok(Samples::from_float(floats)); - }; + let mut samples = Vec::new(); + samples.try_reserve_exact(fields.len()).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid field descriptions: {error}")) + })?; + for mut field in fields { + let floats = Float64Array::new(field.values.into(), field.nulls.finish()); + let sample = if let Some(column) = field.histograms { let batch = column.batch.finish().map_err(external)?; let histograms = HistogramSamples::try_new(UInt32Array::from(column.indices), batch) .map_err(external)?; - Samples::try_from_mixed(floats, histograms).map_err(external) - }) - .collect::>>()?; + Samples::try_from_mixed(floats, histograms).map_err(external)? + } else { + Samples::from_float(floats) + }; + samples.push(sample); + } let fields = SeriesFields::try_from_vec(samples).map_err(external)?; // The runs are all the axis has to remember about the gaps, and the rows are already // compact, so no position holding no value is visited or stored. - let spans = spans_of(&indices); + let run_count = count_runs(&indices); + let spans = spans_of(&indices, run_count)?; let time = TimeAxis::from_grid_spans(self.grid.clone(), spans).map_err(external)?; Series::try_new(&self.schema, metric, time, fields) .map(Some) @@ -265,8 +363,19 @@ impl GridSeriesBuilder { } /// Cuts the retained indices into maximal runs of consecutive positions. -fn spans_of(indices: &[usize]) -> Vec { +fn count_runs(indices: &[usize]) -> usize { + indices + .windows(2) + .filter(|pair| pair[1] != pair[0] + 1) + .count() + .saturating_add(usize::from(!indices.is_empty())) +} + +fn spans_of(indices: &[usize], runs: usize) -> Result> { let mut spans: Vec = Vec::new(); + spans.try_reserve_exact(runs).map_err(|error| { + DataFusionError::Execution(format!("cannot reserve grid spans: {error}")) + })?; let mut start = 0; while let Some(&position) = indices.get(start) { let mut len = 1; @@ -276,14 +385,14 @@ fn spans_of(indices: &[usize]) -> Vec { spans.push(GridSpan::new(position, len)); start += len; } - spans + Ok(spans) } impl FieldBuilder { - fn new() -> Self { + fn new(rows: usize) -> Self { Self { values: Vec::new(), - nulls: NullBufferBuilder::new(0), + nulls: NullBufferBuilder::new(rows), histograms: None, } } @@ -454,7 +563,8 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R // The runs are all the axis remembers about the holes, so a sparse merge stores one span // per run and no row for the positions between them. let time = - TimeAxis::from_grid_spans(grid.clone(), spans_of(&positions)).map_err(external)?; + TimeAxis::from_grid_spans(grid.clone(), spans_of(&positions, count_runs(&positions))?) + .map_err(external)?; let fields = SeriesFields::try_from_vec(fields).map_err(external)?; // The first fragment names the output, as it always has: a label set ignores empty values, // so the metric of the first member is the one the merged series keeps. diff --git a/src/promql/src/native/aggregate.rs b/src/promql/src/native/aggregate.rs new file mode 100644 index 00000000000..975dfb6c3c2 --- /dev/null +++ b/src/promql/src/native/aggregate.rs @@ -0,0 +1,243 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The run-time state of an aggregate node. +//! +//! An aggregate folds a whole vector into one series per group, so its node is the one that cannot +//! answer as its input arrives: a group's value at a grid point is the fold of every value every +//! member of the group has there, and the last member to arrive is the one that decides it. What the +//! node keeps between chunks is therefore the fold's states and nothing else -- a chunk is folded as +//! it arrives and dropped -- and it produces once, when the input ends. +//! +//! That is also why this runtime cannot tell an incomplete group from a complete one: it folds what +//! it is handed. A [`AggregateStage::Complete`] aggregate is correct only if every contribution to +//! each group lies within the streams this execution reads, which is the physical plan's job -- the +//! caller that schedules scans and exchanges guarantees it, and a plan that ran a `Complete` +//! aggregate once per partition and then combined the answers would be combining values that were +//! each missing the other partitions' contributions. The stages that exist for that shape are +//! [`AggregateStage::Partial`] and [`AggregateStage::Final`], which this runtime does not execute +//! yet. + +use std::sync::Arc; + +use datafusion::error::{DataFusionError, Result}; +use datatypes::series::samples::{Series, SeriesSchema}; + +use crate::aggregate::{AggregateOp, FloatAggregator, Grouping}; +use crate::evaluated::EvaluatedChunk; +use crate::grid_series::series_bytes; +use crate::map::validate_float_series; +use crate::native::execute::{NodeBuffer, evaluated_bytes}; +use crate::native::metrics::NodePath; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, Grouping as BoundGrouping, +}; + +/// One aggregate node's states, from the first chunk it is handed to the series it finishes with. +pub(crate) struct AggregateRuntime { + /// What the node computes, as the tree bound it: the kind it folds by and the stage it reports. + op: BoundAggregate, + /// The layout its input and output share: the fields every chunk has and the grid they are + /// addressed on. A `Complete` aggregate neither adds nor drops a field, and it evaluates on its + /// input's grid, so one schema describes both sides. + schema: Arc, + /// The fold. `None` once [`Self::compute`] has taken it, which is also what says the node has + /// nothing left to fold into. + accumulator: Option, + /// Whether the fold has been finished and its series handed over. A node computes once: a second + /// pull finds the series it queued rather than a fold that ran again. + /// + /// The driver reads this to decide whether the node settled: an aggregate answers only once its + /// fold has run, so a child with nothing left to hand over is not an answer. + pub(crate) computed: bool, +} + +impl AggregateRuntime { + /// What this node computes, for a rejection that has to say what it refused. + /// + /// It names the kind this node computes, which distinguishes it from another node in the tree. + fn describe(&self) -> String { + format!("{:?}", self.op.kind) + } + + /// The states of a `Complete` aggregate over `schema`. + /// + /// A stage other than `Complete` and a kind without a supported kernel are refused here, before a single + /// chunk is folded: a node this runtime cannot compute must not start, because a partial answer + /// would be one the caller never asked for. The rejection names what was found, so a reader can + /// tell a tree this runtime cannot run from one it cannot run correctly. + pub(crate) fn new( + op: BoundAggregate, + schema: Arc, + buffer: &mut NodeBuffer, + ) -> Result { + if op.stage != AggregateStage::Complete { + return Err(DataFusionError::NotImplemented(format!( + "the native runtime does not execute a `{:?}` aggregate yet: a `Partial` aggregate \ + reports the state a later stage merges and a `Final` aggregate merges that state, \ + and only `Complete` -- every contribution inside one execution -- is computed here", + op.stage + ))); + } + let kind = match op.kind { + AggregateKind::Sum => AggregateOp::Sum, + AggregateKind::Count => AggregateOp::Count, + AggregateKind::Avg => AggregateOp::Avg, + AggregateKind::Min => AggregateOp::Min, + AggregateKind::Max => AggregateOp::Max, + AggregateKind::Group => AggregateOp::Group, + AggregateKind::Stddev => AggregateOp::Stddev, + AggregateKind::Stdvar => AggregateOp::Stdvar, + other => { + return Err(DataFusionError::NotImplemented(format!( + "the native runtime does not execute the `{other:?}` aggregate yet: only the supported `Complete` aggregate kinds are computed here" + ))); + } + }; + let name_count = match &op.group { + BoundGrouping::By(labels) | BoundGrouping::Without(labels) => labels.len(), + }; + let name_bytes = name_count.saturating_mul(std::mem::size_of::<&str>()); + buffer.grow(name_bytes)?; + let names: Vec<&str> = match &op.group { + BoundGrouping::By(labels) | BoundGrouping::Without(labels) => { + labels.iter().map(|label| label.as_ref()).collect() + } + }; + let grouping = match &op.group { + BoundGrouping::By(_) => Grouping::By(&names), + BoundGrouping::Without(_) => Grouping::Without(&names), + }; + let name_bytes = names.len().saturating_mul(std::mem::size_of::<&str>()); + let initial = FloatAggregator::try_new_budgeted(&schema, kind, grouping, |bytes| { + buffer.resize(bytes.saturating_add(name_bytes)) + })?; + drop(names); + buffer.resize(initial.state_bytes())?; + Ok(Self { + op, + schema, + accumulator: Some(initial), + computed: false, + }) + } + + /// What folding `chunk` reserves while it is folded. + /// + /// The chunk is not retained: it is folded into the states and dropped at the end of the call, + /// so what this node holds afterwards is the fold and none of the input. The charge is therefore + /// the chunk's own bytes -- the figure [`evaluated_bytes`] reports, which is one pass over the + /// chunk and no state between chunks -- and not the `operand_chunk_bytes` figure a binary + /// operand keeps, which exists because such an operand holds a series for the whole execution + /// and so must count the slices of one buffer instead of the buffer once per chunk. + pub(crate) fn pending(&self, chunk: &EvaluatedChunk) -> usize { + evaluated_bytes(chunk) + } + + /// Folds one chunk into its group. + /// + /// A chunk that closes its series closes nothing here: a group is not a series, so the last + /// contribution of one member says nothing about the group. The fold keeps its states, which is + /// what makes a group's answer its whole group's answer once the input ends. + pub(crate) fn accumulate( + &mut self, + path: &NodePath, + chunk: EvaluatedChunk, + buffer: &mut NodeBuffer, + pending: usize, + ) -> Result<()> { + // `Series` keeps its fields public, so a chunk can be built by struct literal against + // another schema; the states are indexed by absolute position on this node's grid, so a + // chunk of another field count or another grid would be folded into the wrong places. + validate_float_series(&self.schema, &chunk.series) + .map_err(|error| folded_chunk(path, &self.describe(), error))?; + let accumulator = self.accumulator.as_mut().ok_or_else(|| no_fold(path))?; + // Keep the chunk covered while the fold decides whether it needs a new group or block. + accumulator.update_budgeted(&chunk.series, |fold_bytes| { + buffer.resize(pending.saturating_add(fold_bytes)) + })?; + drop(chunk); + buffer.resize(accumulator.state_bytes()) + } + + /// The number of retained groups, for budgeting the output queue backing storage. + pub(crate) fn groups(&self) -> usize { + self.accumulator.as_ref().map_or(0, FloatAggregator::groups) + } + + /// Finishes the fold and returns one series per group, or `None` if the fold already ran. + /// + /// `None` is what a node reports after it has queued its series: the caller reads them from the + /// queue, exactly as it reads the chunks a binary node queued, so a second pull does not finish + /// a fold whose states are gone. + pub(crate) fn compute( + &mut self, + path: &NodePath, + buffer: &mut NodeBuffer, + queue_bytes: usize, + ) -> Result>> { + if self.computed { + return Ok(None); + } + if self.accumulator.is_none() { + return Err(no_fold(path)); + } + self.computed = true; + let output = self + .accumulator + .take() + .ok_or_else(|| no_fold(path))? + .finish_budgeted(|total| buffer.resize(total.saturating_add(queue_bytes)))?; + let output_bytes = output.iter().fold(0usize, |bytes, series| { + bytes.saturating_add(series_bytes(series)) + }); + let output_vec_bytes = output + .capacity() + .saturating_mul(std::mem::size_of::()); + buffer.resize( + output_bytes + .saturating_add(output_vec_bytes) + .saturating_add(queue_bytes), + )?; + Ok(Some(output)) + } +} + +/// The rejection of a chunk that is not the layout this node folds. +/// +/// The fold indexes its states by absolute grid position and by field, so a chunk of another field +/// count or another grid would be folded into states that mean something else. `error` is what the +/// series check reported, which says which of the two disagreed. +/// Wraps the error of a chunk this node cannot fold. +/// +/// Only a chunk that does not fit the schema is the node's own refusal; anything else is the +/// kernel saying it does not compute this shape at all -- a native histogram in a float aggregate, +/// for instance -- and a caller that falls back to another planner on `NotImplemented` still can. +/// Wrapping that one in an `Execution` would deny it. +fn folded_chunk(path: &NodePath, aggregate: &str, error: DataFusionError) -> DataFusionError { + match error { + DataFusionError::NotImplemented(_) => error, + error => DataFusionError::Execution(format!( + "the `{aggregate}` aggregate at {path} folds chunks bound to its own schema, and this \ + one is not: {error}" + )), + } +} + +/// The error for a fold that has already been finished. +fn no_fold(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!( + "the aggregate at {path} already finished its fold, so it takes no more input" + )) +} diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index 600a9b61d08..73fe2e94ee5 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -91,9 +91,9 @@ use std::sync::Arc; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datatypes::arrow::array::{Array, Float64Array}; +use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; -use datatypes::series::histogram::{HistogramBatch, HistogramView}; +use datatypes::series::histogram::HistogramView; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{ DenseSamplesRef, RowCursor, SampleRef, Samples, Series, SeriesFields, SeriesSchema, @@ -102,7 +102,8 @@ use datatypes::series::time::{EvalGrid, GridRange, GridSpan, GridSpec, TimeAxis} use datatypes::series::{MetricName, SeriesOrdering}; use crate::grid_series::{ - GridSample, GridSeriesBuilder, LabelSet, label_set, merge_same_labelset, owned, + GridSample, GridSeriesBuilder, LabelSet, fields_bytes, label_set, merge_same_labelset, owned, + series_bytes, table_bytes, }; use crate::histogram_eval::evaluate_histograms; use crate::instant_eval::select_instants; @@ -110,6 +111,7 @@ use crate::map::{ ArithmeticOp, Comparison, ComparisonOp, MathOp, ScalarOperand, drop_metric_name, map_arithmetic, map_compare, map_math, }; +use crate::native::aggregate::AggregateRuntime; use crate::native::batch::EvaluatedChunk; use crate::native::metrics::{NodeMetrics, NodePath}; use crate::native::plan::{ @@ -189,34 +191,11 @@ fn raw_bytes(chunk: &RawSeriesChunk) -> usize { /// The bytes an evaluated chunk holds: its struct and the Arrow buffers behind its fields. /// /// The grid is shared by every chunk addressed on it, so it is not counted per chunk. -fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { - std::mem::size_of::().saturating_add(fields_bytes(&chunk.series.fields)) -} - -/// The bytes one evaluated series holds: its header, the axis that addresses it, the identity it -/// carries and the sample buffers behind its fields. /// -/// It is the figure a fragment of an affected window is charged with as it arrives. The axis's runs -/// and the segments of its fields are counted in one pass over the series, so the estimate costs -/// the fragment rather than a walk over everything the node holds. A payload a fragment shares with -/// another — a histogram table behind two slices of it, the grid behind every axis — is counted -/// once per fragment: that over-count is accepted here rather than a deduplication that would have -/// to keep state across fragments. -fn series_bytes(series: &Series) -> usize { - let runs = series.time.runs().count(); - let segments: usize = series - .fields - .iter() - .map(|field| field.segments().len()) - .sum(); - std::mem::size_of::() - .saturating_add(std::mem::size_of::()) - .saturating_add(runs.saturating_mul(std::mem::size_of::<(usize, usize)>())) - .saturating_add( - segments.saturating_mul(std::mem::size_of::() + std::mem::size_of::()), - ) - .saturating_add(fields_bytes(&series.fields)) - .saturating_add(metric_bytes(&series.metric)) +/// It is what an aggregate node charges for a chunk it is folding: that node hands the chunk over +/// to its fold and drops it, so the chunk alone is what it holds for as long as it holds anything. +pub(super) fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { + std::mem::size_of::().saturating_add(fields_bytes(&chunk.series.fields)) } /// The bytes one chunk a binary operand retains holds on its own. @@ -342,42 +321,6 @@ fn merge_workspace_bytes(rows: usize, fragments: usize, fields: usize) -> usize .saturating_add(output) } -/// The bytes the fields of a chunk hold: the sample buffers of every segment, and the histogram -/// tables those segments reference. -fn fields_bytes(fields: &SeriesFields) -> usize { - fields - .iter() - .map(|field| { - field.segments().iter().fold(0usize, |total, segment| { - total.saturating_add(segment_bytes(segment)) - }) - }) - .fold(0usize, usize::saturating_add) -} - -/// The bytes one dense segment of samples holds. -fn segment_bytes(segment: &Samples) -> usize { - match segment.dense() { - Some(DenseSamplesRef::Float(values)) => values.get_array_memory_size(), - Some(DenseSamplesRef::Histogram(histograms)) => histograms - .indices() - .get_array_memory_size() - .saturating_add(table_bytes(histograms.batch())), - Some(DenseSamplesRef::Mixed(mixed)) => mixed - .floats() - .get_array_memory_size() - .saturating_add(mixed.histograms().indices().get_array_memory_size()) - .saturating_add(table_bytes(mixed.histograms().batch())), - // `segments` hands out dense segments only, so a nested chunk never reaches here. - None => 0, - } -} - -/// The bytes one histogram table holds: the batch struct and the arrays its columns view. -fn table_bytes(table: &Arc) -> usize { - std::mem::size_of::().saturating_add(table.as_arrow().get_array_memory_size()) -} - /// The bytes the allocation `buffer` is a slice of. /// /// The capacity is the whole allocation, which is what a slice keeps alive: a key of a few bytes @@ -408,6 +351,8 @@ enum NodeTag { /// Matches two evaluated vectors and applies one operation to every matched pair of series, or /// broadcasts one constant over the chunks of one evaluated vector. Binary, + /// Folds every series of the input into one series per group, when the input ends. + Aggregate, /// Keeps some samples of its whole input vector at each grid point. Vector, } @@ -422,6 +367,7 @@ impl NodeTag { Self::Window => "window", Self::Map => "map", Self::Binary => "binary", + Self::Aggregate => "aggregate", Self::Vector => "vector operation", } } @@ -576,7 +522,9 @@ impl MergeInput { /// in the query; the metrics report the same figure so a slow node can be told apart from a wide /// one. The reservation is also the figure itself: a node holds nothing beside what it reserved, so /// a gauge reads its size rather than a copy that could drift from it. -struct NodeBuffer { +// Visible to the aggregate node so that its fold can reserve for what it is about to allocate; the +// reservation belongs to a node, not to a module. +pub(super) struct NodeBuffer { metrics: Arc, reservation: MemoryReservation, } @@ -592,7 +540,7 @@ impl NodeBuffer { } /// Reserves `bytes` beside what this node already holds. - fn grow(&mut self, bytes: usize) -> Result<()> { + pub(super) fn grow(&mut self, bytes: usize) -> Result<()> { self.reservation.try_grow(bytes)?; self.record(); Ok(()) @@ -603,7 +551,7 @@ impl NodeBuffer { /// What a node that consumes its input chunk by chunk holds afterwards is not the sum of what /// it consumed, so the reservation is set from the state that is left rather than grown per /// chunk. - fn resize(&mut self, bytes: usize) -> Result<()> { + pub(super) fn resize(&mut self, bytes: usize) -> Result<()> { self.reservation.try_resize(bytes)?; self.record(); Ok(()) @@ -2260,6 +2208,14 @@ struct NodeRuntime { /// `None` for every other node, and for a binary node that broadcasts a constant instead: that /// one holds no operand at all and is described by [`NodeRuntime::broadcast`]. binary: Option, + /// What an aggregate node holds: the fold over what it was handed. + /// + /// Unlike a binary node's operands, nothing of the input survives a chunk: the fold keeps its + /// states and the chunk it read is dropped, so this is the one node whose reservation does not + /// follow what it was handed but what it accumulated. That is also why this node cannot report + /// its input complete from what it holds -- there is nothing to hold -- and reads the driver's + /// [`NativeExecutor::inputs_finished`] and its child's [`NativeExecutionState::settled`] instead. + aggregate: Option, /// What a vector node holds. vector: Option, /// What a binary node holds when one operand is a constant: the operation and the layout of the @@ -2669,6 +2625,102 @@ impl NodeRuntime { self.buffer.resize(queued) } + /// Folds one chunk into this node's aggregate. + /// + /// The chunk is reserved for before it is folded, by what the node reports it is worth, and the + /// reservation settles afterwards on the fold that is left. Unlike [`Self::retain_binary`], what + /// stops being held here is the chunk itself: the fold keeps its states, so a chunk that closes + /// its series leaves nothing behind and the node is empty of input again. + fn retain_aggregate(&mut self, path: &NodePath, chunk: NativeChunk) -> Result<()> { + let NativeChunk::Evaluated(chunk) = chunk else { + return Err(DataFusionError::Execution(format!( + "the aggregate at {path} consumes evaluated chunks, not raw samples" + ))); + }; + // The figure is read before the chunk is folded, because reserving after the fact would be + // reserving for memory the fold already read. + let pending = self + .aggregate + .as_ref() + .ok_or_else(|| no_aggregate(path))? + .pending(&chunk); + let positions = chunk.series.time.len(); + self.buffer.grow(pending)?; + let metrics = Arc::clone(&self.buffer.metrics); + { + let _timer = metrics.timer(); + let (aggregate, buffer) = (&mut self.aggregate, &mut self.buffer); + aggregate + .as_mut() + .ok_or_else(|| no_aggregate(path))? + .accumulate(path, chunk, buffer, pending)?; + } + metrics.record_input(1, positions); + // A chunk that closed its series closed it for the node below: a group is not a series, so + // the fold completing nothing on it is why nothing is counted here. The series this node + // closes are the ones it hands out, counted when it queues them. + Ok(()) + } + + /// Finishes this node's fold and queues one chunk per group it produced. + /// + /// Each group is one series, decided in full by the fold, so the one chunk it is handed over as + /// closes it: a consumer above reads "nothing more here" from that chunk rather than from another + /// one arriving. + fn compute_aggregate(&mut self, path: &NodePath) -> Result<()> { + let metrics = Arc::clone(&self.buffer.metrics); + let _timer = metrics.timer(); + let groups = self + .aggregate + .as_ref() + .map_or(0, |aggregate| aggregate.groups()); + let queue_capacity = self.queued.capacity(); + let queue_target = self.queued.len().saturating_add(groups); + let queue_growth = queue_target + .saturating_sub(queue_capacity) + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + self.buffer.grow(queue_growth)?; + if groups > 0 { + self.queued.try_reserve_exact(groups).map_err(|error| { + DataFusionError::Execution(format!( + "cannot reserve aggregate output queue: {error}" + )) + })?; + } + let queue_bytes = queue_target + .max(queue_capacity) + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + let output = { + let aggregate = self.aggregate.as_mut().ok_or_else(|| no_aggregate(path))?; + aggregate.compute(path, &mut self.buffer, queue_bytes)? + }; + // The fold has returned its series; reserve their handover before appending them to the queue. + let Some(output) = output else { + return Ok(()); + }; + let mut closed = 0usize; + for series in output { + let bytes = series_bytes(&series); + let chunk = NativeChunk::Evaluated(EvaluatedChunk::try_new(series, true)?); + self.queued.push_back((chunk, bytes)); + closed += 1; + } + // A group is one series and this node closes each of them exactly once, on the one chunk it + // hands it over as. + self.buffer.metrics.record_completed_series(closed); + // The output Vec has been consumed; the node now holds only queue backing and chunks. + let queued_bytes = self + .queued + .iter() + .fold(0usize, |total, (_, bytes)| total.saturating_add(*bytes)); + let queue_backing = self + .queued + .capacity() + .saturating_mul(std::mem::size_of::<(NativeChunk, usize)>()); + self.buffer + .resize(queued_bytes.saturating_add(queue_backing)) + } + /// Feeds one chunk to this node's instant selection. /// /// `None` means the series this chunk belongs to has not ended: the selection reads one whole @@ -2871,6 +2923,8 @@ impl NodeRuntime { self.window = None; self.instant = None; self.binary = None; + // The fold is what an aggregate node holds, and dropping it is what gives its states back. + self.aggregate = None; self.vector = None; self.broadcast = None; self.source_schema = None; @@ -2963,15 +3017,14 @@ impl NativeExecutionState { // two vectors, or a constant broadcast over one -- is decided where the operands' // layouts are read. NativeExpr::Binary { .. } => NodeTag::Binary, + // An aggregate is built here and narrowed below, for the same reason: `validate` + // admits every kind and every stage as a legal tree, and the runtime computes one + // kind at one stage. + NativeExpr::Aggregate { .. } => NodeTag::Aggregate, NativeExpr::Vector { .. } => NodeTag::Vector, // A scalar is a value node: it produces the constant it was bound to, hands over no // chunk and holds no state, so it is built like any other node and pulled like none. NativeExpr::Scalar(_) => NodeTag::Scalar, - // An aggregate is legal in a bound tree, and `validate` stays the authority on that; - // this runtime simply does not execute it yet. Rejected before a single chunk is - // read: a tree the runtime cannot execute must not start, because a partial answer - // would be one the caller never asked for. - NativeExpr::Aggregate { .. } => return Err(not_implemented("Aggregate", &path)), }; let buffer = NodeBuffer::new(&path, handles, &pool); let mut runtime = NodeRuntime { @@ -2982,6 +3035,7 @@ impl NativeExecutionState { window: None, map: None, binary: None, + aggregate: None, vector: None, broadcast: None, source_schema: None, @@ -3081,6 +3135,29 @@ impl NativeExecutionState { // out. A broadcast that keeps the labels it reads is not where two series // become one either, so it keeps none; the one that drops them does, above. } + NativeExpr::Aggregate { input, op } => { + // A `Complete` aggregate neither adds nor drops a field and evaluates on its + // input's grid, so the node's own output is the layout it folds. What it does + // not execute -- another kind, or another stage -- is refused here, before the + // pass reaches the input that would feed it, and naming what it found is that + // rejection's job: a `Partial` aggregate's output is state and names no samples + // at all, so the input's layout stands in for it and the stage decides. + let schema = node + .output + .samples() + .or_else(|| input.output.samples()) + .cloned() + .ok_or_else(|| { + DataFusionError::Execution(format!( + "the aggregate at {path} does not produce samples" + )) + })?; + runtime.aggregate = Some(AggregateRuntime::new( + op.clone(), + schema, + &mut runtime.buffer, + )?); + } NativeExpr::Vector { op, .. } => { let schema = node.output.samples().cloned().ok_or_else(|| { DataFusionError::Execution(format!( @@ -3103,10 +3180,9 @@ impl NativeExecutionState { // A vector operation keeps the labels of the series it reads, or answers under // label sets of its own, so it keeps no ledger of what it handed out. } - NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { + NativeExpr::Scalar(_) => { // A scalar holds nothing: what its consumer reads is the constant the node was - // bound to, which is on the tree rather than in this state. The aggregate is - // rejected above, where the pass has the position to name. + // bound to, which is on the tree rather than in this state. } } // The inputs this node reads, in the order the tree declares them: a unary node's one @@ -3114,12 +3190,11 @@ impl NativeExecutionState { // pushed back to front, so the node the spec names is the one the pass reaches next and // nothing but its own subtree lies between the two. let reads_input = match &node.expr { - NativeExpr::Input(_) | NativeExpr::Scalar(_) | NativeExpr::Aggregate { .. } => { - false - } + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, NativeExpr::Instant { input, .. } | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } + | NativeExpr::Aggregate { input, .. } | NativeExpr::Vector { input, .. } => { pending.push((input.as_ref(), path.input())); true @@ -3189,6 +3264,14 @@ impl NativeExecutionState { holds_shared[index] = input_holds_shared; } NodeTag::Map => holds_shared[index] = true, + NodeTag::Aggregate => { + // An aggregate hands over one series per group, each under the label set its + // group was keyed by and each closed on its own chunk: two groups are two + // label sets, so nothing above has to join pieces of a series. Its own input + // is read whole -- the fold holds what it accumulated, not the chunks -- so + // nothing about its input's flag reaches it either. + holds_shared[index] = false; + } NodeTag::Scalar => { // A scalar produces one value and no labels, so it is not a stream of series // and nothing above it can read two of them under one label set. The flag @@ -3742,9 +3825,24 @@ impl NativeExecutor { // A scalar produces its value without reading anything, so it is settled from the // start: there is no chunk it owes its consumer. NodeTag::Scalar => true, + // These report as they are handed chunks, so a child with nothing left is all it takes + // for them to have nothing left either. NodeTag::Map | NodeTag::Window | NodeTag::Instant => self.state.nodes[index] .input .is_none_or(|child| self.settled(child)), + // An aggregate answers one series per group and only once its input ended, so its own + // fold has to have run as well: a child that has handed everything up is not enough, + // because every one of those chunks was folded and none of them answered yet. This is + // what keeps it out of the line above, which asks only about the child. + NodeTag::Aggregate => { + runtime + .aggregate + .as_ref() + .is_some_and(|aggregate| aggregate.computed) + && self.state.nodes[index] + .input + .is_none_or(|child| self.settled(child)) + } NodeTag::Vector => { runtime .vector @@ -3942,6 +4040,7 @@ impl NativeExecutor { // does, so its values are counted rather than read as its positions. runtime.pop_queued(evaluated_values) } + NodeTag::Aggregate => self.aggregate_node(index), NodeTag::Vector => { let child = input.ok_or_else(|| missing_input(tag))?; // Which samples are kept at a point is not decided until the whole input has @@ -3961,6 +4060,35 @@ impl NativeExecutor { } } + /// Folds everything the aggregate at `index` can read and hands over one group's series. + /// + /// What it reads is taken to the end of what its child has, because nothing the fold is handed + /// is retained: a chunk is folded as it arrives and dropped, so reading the child dry costs the + /// fold's states and not the input's bytes. Nothing is answered until the input has ended -- + /// before that, a group's value at a point is still missing the members that have not arrived. + /// + /// That is the difference from a binary node. A binary node retains both operands and so can ask + /// whether they hold anything still open; an aggregate keeps none of them and asks the driver + /// instead: [`Self::inputs_finished`] says the caller will hand over nothing more, and + /// [`NativeExecutionState::settled`] says the child has none left to produce. + fn aggregate_node(&mut self, index: usize) -> Result> { + let child = self.state.nodes[index] + .input + .ok_or_else(|| missing_input(NodeTag::Aggregate))?; + while let Some(chunk) = self.pull(child)? { + let path = &self.state.paths[index]; + self.state.runtimes[index].retain_aggregate(path, chunk)?; + } + if !self.inputs_finished || !self.settled(child) { + return Ok(None); + } + let path = &self.state.paths[index]; + self.state.runtimes[index].compute_aggregate(path)?; + // A group's series holds one value per field and point it decided, so its values are counted + // rather than read as its positions, as a binary result's are. + self.state.runtimes[index].pop_queued(evaluated_values) + } + /// Pulls the vector side of the broadcast at `index` until it hands over a chunk, and returns /// that chunk transformed. /// @@ -4188,11 +4316,7 @@ fn valid_values(chunk: &EvaluatedChunk) -> usize { .sum() } -/// The values a chunk a binary node hands over holds. -/// -/// A binary result is an evaluated chunk, so this is [`valid_values`]: one value per field sample -/// that holds one. A raw chunk cannot reach a binary node's queue, and is answered as its positions -/// so an impossible queue cannot under-count. +/// The values an evaluated chunk hands over, or the timestamps a raw source hands over. fn evaluated_values(chunk: &NativeChunk) -> usize { match chunk { NativeChunk::Evaluated(chunk) => valid_values(chunk), @@ -4200,18 +4324,6 @@ fn evaluated_values(chunk: &NativeChunk) -> usize { } } -/// The rejection of a node kind this runtime does not execute yet. -/// -/// It names the kind and the node's position: the position is what the metrics and `EXPLAIN` call -/// the same node, so a reader can find it in the tree. An aggregate is the one kind left here: a -/// scalar executes as a value node, and a binary node executes both as a vector match and as a -/// constant broadcast. -fn not_implemented(kind: &str, path: &NodePath) -> DataFusionError { - DataFusionError::NotImplemented(format!( - "the native runtime does not execute {kind} nodes yet: {kind} at {path}" - )) -} - /// The error for a window whose evaluator this execution already closed. fn closed(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the window at {path} is already closed")) @@ -4293,11 +4405,27 @@ fn truncated_input(path: &NodePath) -> DataFusionError { )) } +/// The rejection of a node kind this runtime does not execute yet. +/// +/// It names the kind and the node's position: the position is what the metrics and `EXPLAIN` call +/// the same node, so a reader can find it in the tree. +fn not_implemented(kind: &str, path: &NodePath) -> DataFusionError { + DataFusionError::NotImplemented(format!( + "the native runtime does not execute {kind} nodes yet: {kind} at {path}" + )) +} + /// The error for a node asked to compute without the runtime its kind holds. fn no_binary(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no binary operation")) } +/// The error for an aggregate node whose fold is missing. +fn no_aggregate(path: &NodePath) -> DataFusionError { + DataFusionError::Execution(format!("the node at {path} holds no aggregate")) +} + +/// The error for a node asked to run a vector operation without one. fn no_vector(path: &NodePath) -> DataFusionError { DataFusionError::Execution(format!("the node at {path} holds no vector operation")) } diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index 8adc04d3b38..cc6617a5fcc 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -45,8 +45,9 @@ use crate::instant_eval::select_instants; use crate::map::MathOp; use crate::native::metrics::{NodeMetric, NodeMetrics, NodePath}; use crate::native::plan::{ - AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundMap, BoundScalar, - BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, + AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, BoundInstant, + BoundMap, BoundScalar, BoundVectorOp, BoundWindow, Grouping, NativeExpr, NativeNode, + NativeTree, NativeType, }; use crate::range_eval::{RangeEvaluator, RangeFunction, RangeWindow}; use crate::select::SortOrder; @@ -485,6 +486,7 @@ fn source_runtime(pool: &Arc) -> NodeRuntime { window: None, map: None, binary: None, + aggregate: None, broadcast: None, vector: None, source_schema: None, @@ -507,6 +509,7 @@ fn binary_runtime(pool: &Arc, schema: &Arc) -> Nod Arc::new(binary_matching(BinaryOp::Add)), [Arc::clone(schema), Arc::clone(schema)], )), + aggregate: None, broadcast: None, vector: None, source_schema: None, @@ -1083,30 +1086,17 @@ fn a_downstream_window_emits_before_its_input_series_ends() { } } -/// Every node kind the runtime does not execute is rejected, naming the kind and its position. +/// Every binary shape the runtime does not execute is rejected, naming the kind and its position. /// /// A scalar is not one of them any more: it is a value node, and a binary node over a constant is -/// executed as a broadcast. The binary shapes that stay refused are the ones no kernel here takes, -/// and the aggregate is still a kind this runtime does not run at all. +/// executed as a broadcast. The binary shapes that stay refused are the ones no kernel here takes; +/// the aggregate kinds and stages this runtime does not execute are refused by +/// `an_aggregate_kind_or_stage_the_runtime_does_not_compute_is_named`. #[test] fn an_unimplemented_node_is_reported_not_panicked() { // One grid shared by the nodes that validate by allocation identity. let statement = || grid(0, 120_000, 2); - let aggregate_tree = { - let out = statement(); - NativeNode { - expr: NativeExpr::Aggregate { - input: Box::new(window(input(), Arc::clone(&out))), - op: BoundAggregate { - kind: AggregateKind::Sum, - group: Grouping::By(Arc::new(["job".into()])), - stage: AggregateStage::Complete, - }, - }, - output: NativeType::Samples(evaluated_schema(out)), - } - }; // Two constants make a scalar, which is a legal node and a shape no binary here broadcasts // over: there is no vector to pair either of them with. let two_scalars_tree = { @@ -1138,12 +1128,6 @@ fn an_unimplemented_node_is_reported_not_panicked() { }; let cases = [ - ( - "Aggregate", - "root", - "does not execute Aggregate", - aggregate_tree, - ), ("Binary", "root", "two constants", two_scalars_tree), ("Binary", "root", "matching modifiers", matched_scalar_tree), ]; @@ -6558,6 +6542,158 @@ fn a_constant_on_the_right_multiplies_every_value() { assert_eq!(execution.state().reserved_bytes(), 0); } +/// A `Complete` aggregate over `inner`, carrying `inner`'s fields and evaluated on its grid. +/// +/// A `Complete` aggregate neither adds nor drops a field and lands on its input's grid, so the node's +/// own output is the layout its input produces. +fn aggregate(inner: NativeNode, op: BoundAggregate) -> NativeNode { + let fields = inner + .output + .samples() + .expect("an aggregate reads samples") + .fields + .clone(); + let grid = inner + .output + .grid() + .expect("a complete aggregate reads an evaluated vector") + .clone(); + let output = NativeType::Samples(Arc::new(SeriesSchema { + fields, + eval_grid: Some(grid), + })); + NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(inner), + op, + }, + output, + } +} + +/// A `Complete` `sum` grouped as `group`. +fn sum_grouping(group: Grouping) -> BoundAggregate { + BoundAggregate { + kind: AggregateKind::Sum, + group, + stage: AggregateStage::Complete, + } +} + +/// `name{job=...,instance=...}`, the two labels these tests group by. +fn job_instance_metric(name: &str, job: &str, instance: &str) -> MetricName { + MetricName::try_new( + Some(name.into()), + vec![ + Label { + name: "job".into(), + value: job.into(), + }, + Label { + name: "instance".into(), + value: instance.into(), + }, + ], + ) + .unwrap() +} + +/// What one aggregate result reports: its name, its tags and the value at each point it decided. +type GroupResult = ( + Option, + Vec<(String, String)>, + Vec<(usize, Option)>, +); + +/// The groups a run produced, each with the absolute positions it decided, sorted by identity. +/// +/// Sorting is the test's, not the runtime's: a group's place in the answer is the order its first +/// input appeared, which the tests that check it assert on the stream rather than here. +fn groups_of(chunks: &[EvaluatedChunk]) -> Vec { + let mut results: Vec = chunks + .iter() + .map(|chunk| { + let mut tags: Vec<(String, String)> = chunk + .series + .metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + tags.sort(); + ( + chunk.series.metric.name().map(str::to_string), + tags, + decided(std::slice::from_ref(chunk)), + ) + }) + .collect(); + results.sort_by(|left, right| (&left.0, &left.1).cmp(&(&right.0, &right.1))); + results +} + +/// One group's answer under the name and tags `tags` give it, ordered by label name. +fn grouped(tags: &[(&str, &str)], points: Vec<(usize, Option)>) -> GroupResult { + let mut tags: Vec<(String, String)> = tags + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + tags.sort(); + (None, tags, points) +} + +/// One group's answer when every label of its input was dropped. +fn unlabelled(points: Vec<(usize, Option)>) -> GroupResult { + grouped(&[], points) +} + +/// `sum` with no grouping drops every label, name and tags alike, and adds the members per point. +#[test] +fn an_aggregate_drops_every_label_and_adds_its_group_per_point() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let set = ExecutionPlanMetricsSet::new(); + let mut execution = executor(&tree, pool(), &set); + for (metric, values) in [ + ( + job_instance_metric("requests", "api", "a"), + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + ), + ( + job_instance_metric("requests", "db", "b"), + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + ), + ] { + execution + .push_evaluated(evaluated_chunk_of(&schema, metric, 0, values, true)) + .unwrap(); + } + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ])], + "the two series were not folded into one unlabelled group" + ); + // One group is one series, so it closes exactly once and its chunk says so. + assert!(output.len() == 1 && output[0].is_last); + assert_eq!( + count(&set, &NodePath::root(), NodeMetric::COMPLETED_SERIES), + 1 + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + /// The side a constant is on decides which operand the vector's value is. /// /// Division is not commutative, so `x / 100` and `100 / x` answering different values is what says @@ -7031,6 +7167,208 @@ fn a_window_above_a_broadcast_merges_the_series_it_collapses() { assert_eq!(execution.state().reserved_bytes(), 0); } +/// Every supported `Complete` kind reaches its matching float kernel through the native runtime. +#[test] +fn complete_aggregate_kinds_fold_two_values() { + let cases = [ + (AggregateKind::Sum, 4.0), + (AggregateKind::Count, 2.0), + (AggregateKind::Avg, 2.0), + (AggregateKind::Min, 1.0), + (AggregateKind::Max, 3.0), + (AggregateKind::Group, 1.0), + (AggregateKind::Stddev, 1.0), + (AggregateKind::Stdvar, 1.0), + ]; + + for (kind, expected) in cases { + let grid = grid(0, 60_000, 1); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + BoundAggregate { + kind, + group: Grouping::By(Arc::new([])), + stage: AggregateStage::Complete, + }, + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + for value in [1.0, 3.0] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(value)], + true, + )) + .unwrap(); + } + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![(0, Some(expected))])], + "{kind:?} did not compute its native complete result" + ); + } +} + +/// The same series, delivered in two chunks, fold to what the whole series folds to. +#[test] +fn an_aggregate_folds_a_series_split_across_chunks_like_a_whole_one() { + let bound = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&bound)); + // A fresh tree per execution: the node's state is the execution's, not the tree's. + let tree = || { + NativeTree::try_new(aggregate( + evaluated_input(grid(0, 60_000, 4)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap() + }; + let (first, second) = ( + job_instance_metric("requests", "api", "a"), + job_instance_metric("requests", "api", "b"), + ); + + // One chunk per series. + let mut whole = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + for (metric, values) in [ + ( + first.clone(), + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + ), + ( + second.clone(), + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + ), + ] { + whole + .push_evaluated(evaluated_chunk_of(&schema, metric, 0, values, true)) + .unwrap(); + } + let whole = whole.finish().unwrap(); + + // The first series split in two, the second whole: the group is the same group, folded from + // positions the chunks name absolutely. + let mut split = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + split + .push_evaluated(evaluated_chunk_of( + &schema, + first.clone(), + 0, + vec![Some(1.0), Some(2.0)], + false, + )) + .unwrap(); + assert!( + split.next_chunk().unwrap().is_none(), + "a chunk that does not close its series answered for the group" + ); + split + .push_evaluated(evaluated_chunk_of( + &schema, + first, + 2, + vec![Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + + split + .push_evaluated(evaluated_chunk_of( + &schema, + second, + 0, + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + true, + )) + .unwrap(); + + let split = split.finish().unwrap(); + assert_eq!(groups_of(&split), groups_of(&whole)); + assert_eq!( + groups_of(&split), + vec![grouped( + &[("job", "api")], + vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ], + )], + "the split series was folded as two groups or at the wrong positions" + ); +} + +/// A point no member of a group has a value at is not a point of the group's result. +/// +/// Sparse input is what says the fold indexes its states by absolute position: a chunk holding +/// positions 0 and 1 and another holding 1 and 2 leave position 3 to nothing, and the hole between +/// the rows is never filled with a value. +#[test] +fn an_aggregate_reports_only_the_points_a_group_holds_a_value_at() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + // Positions 0 and 1 of the first series, 1 and 2 of the second, and nothing at 3. + execution + .push_evaluated(sparse_chunk( + &schema, + job_instance_metric("requests", "api", "a"), + &[(0, 2)], + vec![1.0, 2.0], + true, + )) + .unwrap(); + execution + .push_evaluated(sparse_chunk( + &schema, + job_instance_metric("requests", "db", "b"), + &[(1, 2)], + vec![20.0, 300.0], + true, + )) + .unwrap(); + + let output = execution.finish().unwrap(); + assert_eq!( + groups_of(&output), + vec![ + // Position 1 is the only one both members hold, and it is their sum. + unlabelled(vec![(0, Some(1.0)), (1, Some(22.0)), (2, Some(300.0))]), + ], + "the group was answered at a point no member holds, or summed the wrong members" + ); +} + +/// An aggregate handed no series produces nothing at all: no group, so no empty series either. +#[test] +fn an_aggregate_over_no_series_produces_nothing() { + let grid = grid(0, 60_000, 4); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + let output = execution.finish().unwrap(); + assert!( + output.is_empty(), + "an aggregate with no input answered {} series", + output.len() + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + /// A broadcast hands one chunk over for every chunk it reads, and closes the series once. /// /// The metrics are what says the chunk-per-chunk behaviour held: one input batch per chunk read, one @@ -7246,6 +7584,411 @@ fn a_broadcast_refuses_two_series_that_share_a_label_set_and_releases_the_reserv ); } +/// A series that ended is one member of a group, not the group: nothing is answered until the input +/// ends. +/// +/// The observable state between the two is what this pins: the fold has read the closing chunk of +/// the first series and answered nothing, and `computed` still says so. +#[test] +fn an_aggregate_answers_nothing_until_its_whole_input_ends() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + let computed = |execution: &NativeExecutor| { + execution.state().runtimes[0] + .aggregate + .as_ref() + .expect("the root holds an aggregate") + .computed + }; + + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", "api", "a"), + 0, + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().unwrap().is_none()); + assert!(!computed(&execution), "the fold finished on one series"); + assert!( + execution.next_chunk().unwrap().is_none(), + "the closing chunk of one series answered for its group" + ); + // The second series closed too, and the input still has not ended. + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", "db", "b"), + 0, + vec![Some(10.0), Some(20.0), Some(30.0), Some(40.0)], + true, + )) + .unwrap(); + assert!( + !computed(&execution), + "the fold finished before the input ended" + ); + assert!( + execution.next_chunk().unwrap().is_none(), + "every series closed, and the group was answered before the input ended" + ); + + execution.end_of_input().unwrap(); + assert!( + !computed(&execution), + "ending the input finished the fold before anything pulled for it" + ); + let mut output = Vec::new(); + while let Some(chunk) = execution.next_chunk().unwrap() { + output.push(chunk); + } + assert_eq!( + groups_of(&output), + vec![unlabelled(vec![ + (0, Some(11.0)), + (1, Some(22.0)), + (2, Some(33.0)), + (3, Some(44.0)), + ])], + "the group was not answered once the input ended" + ); +} + +/// A chunk of another layout is refused, naming the aggregate and its position, and the execution +/// stops with nothing reserved. +/// +/// The fold indexes its states by absolute position and by field, so a chunk of another field count +/// would be folded into states that mean something else. It is refused rather than folded, and what +/// the refusal leaves behind is a stopped execution whose reservation went back to the pool. +#[test] +fn an_aggregate_refuses_a_chunk_of_another_layout_and_returns_its_reservation() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + // Two fields where the tree bound one: the source node accepts it, the fold refuses it. + let wider = Arc::new(SeriesSchema { + fields: field_count(2), + eval_grid: Some(Arc::clone(&grid)), + }); + execution + .push_evaluated(evaluated_columns( + &wider, + metric(), + 0, + vec![vec![Some(1.0); 4], vec![Some(2.0); 4]], + true, + )) + .unwrap(); + + let error = execution + .next_chunk() + .expect_err("a chunk of another field count was folded"); + assert!( + matches!(error, DataFusionError::Execution(_)), + "the refusal is not an execution error: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("aggregate at root"), + "the refusal does not name the aggregate: {message}" + ); + assert!( + message.contains("field count"), + "the refusal does not say what disagreed: {message}" + ); + assert_eq!(execution.state().reserved_bytes(), 0); + // A stopped execution is not resumed: the fold must not be asked to continue after this. + assert!(execution.next_chunk().is_err()); +} + +/// A pool rejection while the first aggregate block is admitted releases the source and fold. +#[test] +fn aggregate_budget_rejects_first_block_and_releases_everything() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(2_000); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(2.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().is_err()); + assert_eq!(pool.reserved(), 0); + assert_eq!(execution.state().reserved_bytes(), 0); +} + +/// The aggregate runtime charges its output queue, releases each output as it is pulled, and keeps +/// queue backing until the runtime itself is released. +#[test] +fn aggregate_runtime_queue_keeps_only_backing_after_pops_until_release() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new(["job".into()]))), + )) + .unwrap(); + let pool = bounded_pool(1 << 20); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + for (job, value) in [("api", 1.0), ("db", 2.0)] { + execution + .push_evaluated(evaluated_chunk_of( + &schema, + job_instance_metric("requests", job, "a"), + 0, + vec![Some(value)], + true, + )) + .unwrap(); + } + execution.end_of_input().unwrap(); + let first = execution.next_chunk().unwrap().expect("first group output"); + assert_eq!(first.series.time.positions().count(), 1); + let runtime = &execution.state().runtimes[0]; + let entry = std::mem::size_of::<(NativeChunk, usize)>(); + assert_eq!(runtime.queued.len(), 1); + let remaining = runtime.queued[0].1; + let expected = runtime.queued.capacity() * entry + remaining; + assert_eq!(runtime.buffer.reservation.size(), expected); + + let second = execution + .next_chunk() + .unwrap() + .expect("second group output"); + assert_eq!(second.series.time.positions().count(), 1); + let runtime = &execution.state().runtimes[0]; + assert!(runtime.queued.is_empty()); + let backing = runtime.queued.capacity() * entry; + assert!(backing > 0); + assert_eq!(runtime.buffer.reservation.size(), backing); + assert!(execution.next_chunk().unwrap().is_none()); + // Pulling the terminal None releases the queue backing; only explicit close checks the + // runtime's final cleanup below. + assert_eq!(execution.state().runtimes[0].buffer.reservation.size(), 0); + + execution.state.close().unwrap(); + assert_eq!(execution.state().runtimes[0].queued.capacity(), 0); + assert_eq!(execution.state().runtimes[0].buffer.reservation.size(), 0); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate whose sparse output has a few points on a large grid can finish in a small pool. +#[test] +fn aggregate_sparse_output_budget_uses_rows_not_grid_length() { + let grid = grid(0, 1, 100_000); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(64 * 1024); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 10, + vec![Some(2.0)], + true, + )) + .unwrap(); + let output = execution.finish().unwrap(); + assert_eq!(groups_of(&output), vec![unlabelled(vec![(10, Some(2.0))])]); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate retains its fold before finish and release returns its reservation. +#[test] +fn aggregate_fold_state_is_released_after_input() { + let grid = grid(0, 1, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = NativeTree::try_new(aggregate( + evaluated_input(Arc::clone(&grid)), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap(); + let pool = bounded_pool(1 << 20); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + metric(), + 0, + vec![Some(2.0)], + true, + )) + .unwrap(); + assert!(execution.next_chunk().unwrap().is_none()); + let runtime = &execution.state.runtimes[0]; + let aggregate = runtime.aggregate.as_ref().unwrap(); + assert_eq!(aggregate.groups(), 1); + assert!(!aggregate.computed); + assert!(pool.reserved() > 0); + execution.state.close().unwrap(); + assert_eq!(pool.reserved(), 0); +} + +/// An aggregate kind or stage this runtime does not compute is refused, naming what was found. +#[test] +fn an_aggregate_kind_or_stage_the_runtime_does_not_compute_is_named() { + let grid = || grid(0, 60_000, 4); + let kind_tree = NativeTree::try_new(aggregate( + evaluated_input(grid()), + BoundAggregate { + kind: AggregateKind::Quantile { quantile: 0.5 }, + group: Grouping::By(Arc::new([])), + stage: AggregateStage::Complete, + }, + )) + .unwrap(); + let partial_tree = { + let out = grid(); + let group = Grouping::By(Arc::new(["job".into()])); + let state = Arc::new(AggregateStateSchema { + group: group.clone(), + accumulators: fields(), + grid: Arc::clone(&out), + }); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(evaluated_input(Arc::clone(&out))), + op: BoundAggregate { + kind: AggregateKind::Sum, + group, + stage: AggregateStage::Partial, + }, + }, + output: NativeType::AggregateState(state), + }) + .unwrap() + }; + + for (found, tree) in [("Quantile", kind_tree), ("Partial", partial_tree)] { + let error = match NativeExecutionState::try_new( + &tree, + pool(), + &ExecutionPlanMetricsSet::new(), + 0, + ) { + Ok(_) => panic!("the runtime accepted a {found} aggregate"), + Err(error) => error, + }; + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "{found}: {error}" + ); + let message = error.to_string(); + assert!(message.contains(found), "{found}: {message}"); + // A `Partial` rejection has to say why it is not computed: the stages it would need. + if found == "Partial" { + assert!(message.contains("Final"), "{message}"); + } + } +} + +/// An aggregate above a window that collects and merges its input waits for that merge. +/// +/// The multiplication drops the metric name, so the two input series land on one label set and the +/// `last_over_time` above them has to collect and merge before it evaluates. The aggregate above the +/// window must wait for that merge: its guard is `inputs_finished && settled(child)`, and with only +/// the first half it finalizes on the pull that follows the caller's last chunk, before the window +/// has answered. +#[test] +fn an_aggregate_above_a_collecting_window_waits_for_the_merge() { + let grid = grid(0, 60_000, 4); + let schema = evaluated_schema(Arc::clone(&grid)); + let tree = || { + NativeTree::try_new(aggregate( + window_with( + broadcast_right( + evaluated_input(Arc::clone(&grid)), + 2.0, + BinaryOp::Mul, + false, + ), + RangeFunction::LastOverTime, + 60_000, + Some(Arc::clone(&grid)), + Arc::clone(&grid), + ), + sum_grouping(Grouping::By(Arc::new([]))), + )) + .unwrap() + }; + + let mut execution = executor(&tree(), pool(), &ExecutionPlanMetricsSet::new()); + // Two series that only differ by name, each holding half of the grid the other leaves empty: + // after the name is dropped they are one series only once merged. + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("requests", "x"), + 0, + vec![Some(1.0), Some(2.0)], + true, + )) + .unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "the aggregate answered before the window above it had merged" + ); + execution + .push_evaluated(evaluated_chunk_of( + &schema, + named_metric("requests_total", "x"), + 2, + vec![Some(3.0), Some(4.0)], + true, + )) + .unwrap(); + + let out = run(&mut execution); + assert_eq!( + decided(&out), + vec![ + (0, Some(2.0)), + (1, Some(4.0)), + (2, Some(6.0)), + (3, Some(8.0)), + ], + "the aggregate did not fold the merged window result" + ); + assert_eq!(execution.state().reserved_bytes(), 0); +} + fn vector_op(inner: NativeNode, op: BoundVectorOp) -> NativeNode { let output = inner.output.clone(); NativeNode { diff --git a/src/promql/src/native/mod.rs b/src/promql/src/native/mod.rs index c9545a9fe69..5375c92f1c2 100644 --- a/src/promql/src/native/mod.rs +++ b/src/promql/src/native/mod.rs @@ -27,6 +27,7 @@ //! - A series closes exactly once, on the chunk that says so, so a consumer can tell "nothing more //! here" from "no more input". +mod aggregate; pub mod batch; mod display; pub mod execute; diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index d131af69ba9..7e50b2c644d 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -20,10 +20,10 @@ use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datatypes::arrow::array::Float64Array; use datatypes::arrow::buffer::{Buffer, ScalarBuffer}; -use datatypes::series::MetricName; use datatypes::series::SeriesOrdering; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; use datatypes::series::samples::{FieldSpec, SampleRef, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::{Label, MetricName}; use promql::evaluated::EvaluatedChunk; use promql::map::MathOp; use promql::native::execute::{NativeExecutionState, NativeExecutor}; @@ -3723,3 +3723,526 @@ fn whole_vector_functions_bind_a_vector_operation() { assert_eq!(root, format!(" root Vector {op}"), "`{query}`"); } } + +/// One raw chunk of `metric`, keyed by `key`, closing the series. +/// +/// `timestamps` are milliseconds and `values` one per timestamp, which is what a scan of the series +/// hands the tree. +fn agg_raw_chunk( + schema: &SeriesSchema, + key: &[u8], + metric: MetricName, + timestamps: &[i64], + values: &[f64], +) -> RawSeriesChunk { + agg_raw_chunk_with_last(schema, key, metric, timestamps, values, true) +} + +/// One raw chunk of `metric`, keyed by `key`, with an explicit series-ending flag. +fn agg_raw_chunk_with_last( + schema: &SeriesSchema, + key: &[u8], + metric: MetricName, + timestamps: &[i64], + values: &[f64], + is_last: bool, +) -> RawSeriesChunk { + RawSeriesChunk::try_new( + schema, + Buffer::from_slice_ref(key), + Some(metric), + RawTimestamps::try_new(ScalarBuffer::from(timestamps.to_vec())).unwrap(), + SeriesFields::one(Samples::from_float(Float64Array::from(values.to_vec()))), + is_last, + ) + .unwrap() +} + +/// `name` with the labels in `tags`, in the order given. +fn agg_metric(name: &str, tags: &[(&str, &str)]) -> MetricName { + MetricName::try_new( + Some(name.into()), + tags.iter() + .map(|(name, value)| Label { + name: (*name).into(), + value: (*value).into(), + }) + .collect(), + ) + .unwrap() +} + +/// What one result series of a native execution reports. +type AggSeries = ( + Option, + Vec<(String, String)>, + Vec<(usize, Option)>, +); + +/// Reads the outputs a native execution produced. +fn aggregate_results(output: &[EvaluatedChunk]) -> Vec { + let mut results: Vec = output + .iter() + .map(|chunk| { + let mut tags: Vec<(String, String)> = chunk + .series + .metric + .tags() + .iter() + .map(|tag| (tag.name.to_string(), tag.value.to_string())) + .collect(); + tags.sort(); + let points = chunk + .series + .time + .positions() + .map(|(position, row)| { + let sample = chunk.series.fields.as_slice()[0] + .iter() + .nth(row) + .expect("the axis and the field are aligned"); + ( + position, + match sample { + SampleRef::Float(value) => Some(value), + _ => None, + }, + ) + }) + .collect(); + (chunk.series.metric.name().map(str::to_string), tags, points) + }) + .collect(); + results.sort_by(|left, right| (&left.0, &left.1).cmp(&(&right.0, &right.1))); + results +} + +/// Runs `tree` over `chunks`, which enter at the tree's only input, and reads the answer. +/// +/// The answer is read the way a caller reads it: the whole input is pushed, the input is ended, and +/// what the root reports is collected. Each series is reported by its identity and the value it +/// decided at each absolute grid position. +fn run_aggregate_tree(tree: &NativeTree, chunks: Vec) -> Vec { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution + .push_raw(chunk) + .unwrap_or_else(|error| panic!("a source chunk was refused: {error}")); + } + let output = execution + .finish() + .unwrap_or_else(|error| panic!("the tree did not answer: {error}")); + aggregate_results(&output) +} + +/// The 21 grid points of a statement over 100 seconds in 5-second steps. +fn agg_points_0_to_100() -> Vec { + (0..21).collect() +} + +/// The timestamps of the same 21 points, in milliseconds. +fn agg_timestamps_0_to_100() -> Vec { + (0..21).map(|index| index * 5_000).collect() +} + +/// `sum(x)` over two series: one unlabelled group, and every label of both inputs dropped. +#[test] +fn a_compiled_sum_drops_every_label_and_adds_its_group() { + let tree = compile( + &eval_stmt("sum(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + ×tamps, + &(1..=21).map(f64::from).collect::>(), + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "db"), ("instance", "b")]), + ×tamps, + &(1..=21).map(|value| (2 * value) as f64).collect::>(), + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + // `3` at the first point, three more at each of the twenty after it: the two members summed. + let values: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(3.0 * (position as f64 + 1.0)))) + .collect(); + assert_eq!(output, vec![(None, Vec::new(), values)]); +} + +#[test] +fn a_compiled_sum_aggregates_series_chunks_after_input_ends() { + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let values: Vec = (1..=21).map(f64::from).collect(); + let cases = [ + ("sum(x)", None, Vec::new()), + ( + "sum by(job)(x)", + None, + vec![("job".to_string(), "api".to_string())], + ), + ( + "sum without(instance)(x)", + None, + vec![("job".to_string(), "api".to_string())], + ), + ]; + + for (query, name, tags) in cases { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + let metric = agg_metric("x", &[("job", "api"), ("instance", "a")]); + let chunks = vec![ + agg_raw_chunk_with_last( + &schema, + b"a", + metric.clone(), + ×tamps[..10], + &values[..10], + false, + ), + agg_raw_chunk_with_last( + &schema, + b"a", + metric, + ×tamps[10..], + &values[10..], + true, + ), + ]; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(&tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution.push_raw(chunk).unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "`{query}` produced output before the whole input ended" + ); + } + let output = execution.finish().unwrap(); + let results = aggregate_results(&output); + let expected_points = values + .iter() + .enumerate() + .map(|(position, value)| (position, Some(*value))) + .collect(); + assert_eq!( + results, + vec![(name.map(str::to_string), tags, expected_points)], + "`{query}`" + ); + } +} + +/// Empty source input is not a member of an aggregate's implicit group. +#[test] +fn a_compiled_aggregate_over_empty_input_produces_no_groups() { + for query in ["sum(x)", "sum by(job)(x)", "sum without(instance)(x)"] { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + assert_eq!( + run_aggregate_tree(&tree, Vec::new()), + Vec::new(), + "`{query}` invented a group for empty input" + ); + } +} + +/// `sum by(job)(x)` over five series: two groups, each the sum of its own members. +/// +/// The members cover different points, so each group is a sum over a subset that changes from point +/// to point: one series holds the whole grid, one holds a single sample between two grid points, and +/// one stops a whole lookback before the grid does. No last sample sits exactly one lookback behind +/// a grid point, so which points each member covers does not depend on how the boundary is read. +#[test] +fn a_compiled_sum_by_groups_its_members_and_sums_each_group() { + let tree = compile( + &eval_stmt("sum by(job)(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let all = agg_timestamps_0_to_100(); + let chunks = vec![ + // The whole grid, holding one. + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + &all, + &[1.0; 21], + ), + // One sample halfway between grid points 0 and 1, so only grid point 1 selects it. + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b")]), + &[2_500], + &[2.0], + ), + // Four samples between grid points, covering grid points 1 to 4. + agg_raw_chunk( + &schema, + b"e", + agg_metric("x", &[("job", "api"), ("instance", "e")]), + &[2_400, 7_400, 12_400, 17_400], + &[5.0; 4], + ), + // The whole grid, holding ten. + agg_raw_chunk( + &schema, + b"c", + agg_metric("x", &[("job", "db"), ("instance", "c")]), + &all, + &[10.0; 21], + ), + // One sample halfway between grid points 0 and 1, holding twenty. + agg_raw_chunk( + &schema, + b"d", + agg_metric("x", &[("job", "db"), ("instance", "d")]), + &[2_500], + &[20.0], + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + // `api` is `a` alone at points 0 and 5 to 20, and every member at the points they cover: + // point 1 is `1 + 2 + 5`, points 2 to 4 are `1 + 5`. + let api: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| { + let value = match position { + 1 => 8.0, + 2..=4 => 6.0, + _ => 1.0, + }; + (position, Some(value)) + }) + .collect(); + // `db` is the same shape over its own members: point 1 is `10 + 20`, the rest is `c` alone. + let db: Vec<(usize, Option)> = agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(if position == 1 { 30.0 } else { 10.0 }))) + .collect(); + assert_eq!( + output, + vec![ + (None, vec![("job".to_string(), "api".to_string())], api), + (None, vec![("job".to_string(), "db".to_string())], db), + ] + ); +} + +/// `sum without(instance)(x)`: `instance` goes, `job` and `zone` stay, and so does `__name__` going. +#[test] +fn a_compiled_sum_without_drops_the_named_label_and_the_name() { + let tree = compile( + &eval_stmt("sum without(instance)(x)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let schema = raw_schema(); + let timestamps = agg_timestamps_0_to_100(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a"), ("zone", "west")]), + ×tamps, + &[1.0; 21], + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b"), ("zone", "east")]), + ×tamps, + &[2.0; 21], + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + let values_of = |value: f64| -> Vec<(usize, Option)> { + agg_points_0_to_100() + .into_iter() + .map(|position| (position, Some(value))) + .collect() + }; + // Two groups, because `zone` is not the dropped label: the two series differ by more than the + // label the grouping removes. + assert_eq!( + output, + vec![ + ( + None, + vec![ + ("job".to_string(), "api".to_string()), + ("zone".to_string(), "east".to_string()), + ], + values_of(2.0), + ), + ( + None, + vec![ + ("job".to_string(), "api".to_string()), + ("zone".to_string(), "west".to_string()), + ], + values_of(1.0), + ), + ] + ); +} + +/// `sum by(job)(rate(x[5m]))`: an aggregate over a window, mid-chain. +/// +/// Every series is a counter rising by one per five seconds, and every grid point's window ends on a +/// sample and begins five seconds before one, so the extrapolated rate is the same at every point: +/// `59 / 295 = 0.2` per second, which the double arithmetic lands on the literals below. +#[test] +fn a_compiled_sum_over_a_rate_groups_the_window_it_reads() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + // The window reads 300 seconds, so the samples start a full window before the first grid point. + let schema = raw_schema(); + let timestamps: Vec = (0..=80).map(|index| index * 5_000).collect(); + let values: Vec = (0..=80).map(|index| f64::from(index) + 1.0).collect(); + let chunks = vec![ + agg_raw_chunk( + &schema, + b"a", + agg_metric("x", &[("job", "api"), ("instance", "a")]), + ×tamps, + &values, + ), + agg_raw_chunk( + &schema, + b"b", + agg_metric("x", &[("job", "api"), ("instance", "b")]), + ×tamps, + &values, + ), + agg_raw_chunk( + &schema, + b"c", + agg_metric("x", &[("job", "db"), ("instance", "c")]), + ×tamps, + &values, + ), + ]; + + let output = run_aggregate_tree(&tree, chunks); + let at_every_point = |value: f64| -> Vec<(usize, Option)> { + (0..21).map(|position| (position, Some(value))).collect() + }; + // One `rate` is `0.19999999999999998`, and the pair of them add to `0.39999999999999997`. + assert_eq!( + output, + vec![ + ( + None, + vec![("job".to_string(), "api".to_string())], + at_every_point(0.39999999999999997), + ), + ( + None, + vec![("job".to_string(), "db".to_string())], + at_every_point(0.19999999999999998), + ), + ] + ); +} + +/// Splitting one counter's samples does not change the rate aggregate or make it answer early. +#[test] +fn a_compiled_sum_over_a_rate_aggregates_series_chunks_after_input_ends() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + let schema = raw_schema(); + let timestamps: Vec = (0..=80).map(|index| index * 5_000).collect(); + let values: Vec = (0..=80).map(|index| f64::from(index) + 1.0).collect(); + let metric = agg_metric("x", &[("job", "api"), ("instance", "a")]); + let chunks = [ + agg_raw_chunk_with_last( + &schema, + b"a", + metric.clone(), + ×tamps[..41], + &values[..41], + false, + ), + agg_raw_chunk_with_last( + &schema, + b"a", + metric, + ×tamps[41..], + &values[41..], + true, + ), + ]; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 26)); + let state = + NativeExecutionState::try_new(&tree, pool, &ExecutionPlanMetricsSet::new(), 0).unwrap(); + let mut execution = NativeExecutor::new(state); + for chunk in chunks { + execution.push_raw(chunk).unwrap(); + assert!( + execution.next_chunk().unwrap().is_none(), + "the rate aggregate produced output before the whole input ended" + ); + } + let output = execution.finish().unwrap(); + assert_eq!( + aggregate_results(&output), + vec![( + None, + vec![("job".to_string(), "api".to_string())], + (0..21) + .map(|position| (position, Some(0.19999999999999998))) + .collect(), + )] + ); +} + +/// Empty source input is not a member of this aggregate's implicit group, with or without a window. +#[test] +fn a_compiled_sum_over_a_rate_with_empty_input_produces_no_groups() { + let statement = EvalStmt { + expr: promql_parser::parser::parse("sum by(job)(rate(x[5m]))").unwrap(), + start: UNIX_EPOCH + Duration::from_secs(300), + end: UNIX_EPOCH + Duration::from_secs(400), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(5), + }; + let tree = compile(&statement, raw_schema(), SeriesOrdering::KeyThenTime).unwrap(); + assert_eq!(run_aggregate_tree(&tree, Vec::new()), Vec::new()); +} From a8737d099610d5716109a6fe438e93b740e190c0 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 01:02:31 +0800 Subject: [PATCH 2/4] test(promql): skip an endpoint the platform cannot represent (#9894) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/native_compile/tests.rs | 25 ++++++++++++++++---- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 7e50b2c644d..d42ae4fa7b4 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::sync::Arc; -use std::time::{Duration, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use datafusion::error::DataFusionError; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; @@ -3138,11 +3138,25 @@ fn large_subquery_child_grids_compile_deferred() { assert_eq!(grid.2, range_ms); } +/// The largest millisecond endpoint the platform can represent as a `SystemTime`. +/// +/// The tests below need an endpoint close enough to `i64::MAX` milliseconds that a one-second +/// rounding or a one-millisecond offset crosses the representable range. Unix-like platforms can +/// hold that instant; Windows represents `SystemTime` as a 64-bit count of 100-nanosecond +/// intervals since 1601, so the same instant does not fit and `UNIX_EPOCH + duration` panics. The +/// behaviour under test is about arithmetic on the millisecond value, not about the platform's +/// clock range, so a platform that cannot represent the instant skips instead of failing. +fn extreme_millisecond_endpoint() -> Option { + UNIX_EPOCH.checked_add(Duration::from_millis(i64::MAX as u64)) +} + /// A sampled child endpoint can fit even when the unrounded read-window endpoint does not. #[test] fn a_subquery_accepts_an_in_range_sampled_endpoint() { + let Some(endpoint) = extreme_millisecond_endpoint() else { + return; + }; let mut stmt = eval_stmt("rate(x[1s:1s] offset -1ms)"); - let endpoint = UNIX_EPOCH + Duration::from_millis(i64::MAX as u64); stmt.start = endpoint; stmt.end = endpoint; stmt.interval = Duration::from_millis(1); @@ -3161,6 +3175,9 @@ fn a_subquery_accepts_an_in_range_sampled_endpoint() { /// An anchored fold offset that cannot be represented as i64 is a plan error. #[test] fn a_subquery_with_an_overflowing_fold_offset_is_refused() { + let Some(endpoint) = extreme_millisecond_endpoint() else { + return; + }; let mut stmt = eval_stmt("rate(x[5m:1m] @ 0)"); let PromExpr::Call(call) = &mut stmt.expr else { unreachable!("query parses to a call"); @@ -3168,9 +3185,7 @@ fn a_subquery_with_an_overflowing_fold_offset_is_refused() { let PromExpr::Subquery(subquery) = &mut *call.args.args[0] else { unreachable!("call argument parses to a subquery"); }; - subquery.at = Some(AtModifier::At( - UNIX_EPOCH + Duration::from_millis(i64::MAX as u64), - )); + subquery.at = Some(AtModifier::At(endpoint)); subquery.offset = Some(Offset::Neg(Duration::from_millis(1))); let error = compile(&stmt, raw_schema(), SeriesOrdering::KeyThenTime) From 3b09dd75dcad10e1ea7b5b9739402a79350ccb70 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 01:03:18 +0800 Subject: [PATCH 3/4] feat(promql): run a bound tree over raw series in one fused node (#9820) * feat(promql): run a bound tree over raw series in one fused node `FusedSeriesExec` evaluates a compiled `NativeTree` directly over the raw series its child produces, instead of materialising a row batch through the per-operator chain. It binds the time column, the fields and the tags once, checks the input against that binding when the plan is built, and converts the chunks the tree settles back into the statement's output columns. The node's partition count is the child's, one tree per partition. The child must be a raw-series scan that assigns every series whole to one partition: `PerSeries` distribution is what makes a per-partition evaluation equivalent to evaluating the whole series. A zero-partition input is refused; anything else is the child's business. A label set can recur across partitions once `rate` and the map kernels drop `__name__`, so the occupancy map is plan state rather than per-stream. Two series reaching one label set is an error only when their occupied grid points overlap, which is Prometheus' range-query rule; disjoint points are legal and are merged into one occupancy record. The reservation for that map is attached to the first execution's pool and released when the plan is dropped, because a later partition still has to collide with what an earlier one emitted. `occupied_runs` reads a chunk's absolute positions rather than assuming one contiguous run: a sparse evaluated chunk may hold several, and treating the first as the whole range would count the holes as rows. A poll does a bounded amount of work and wakes itself when the budget runs out, so a long run of chunks that settle nothing cannot monopolise the executor thread. Signed-off-by: discord9 * fix(promql): charge for the label sets the plan keeps, not for one stream The occupancy map is the plan's: entries are inserted by whichever partition sees a label set first and no path removes them. The reservation paying for them was the stream's, so the bytes went back when that stream ended while the entries stayed. Measured on a plan with three partitions producing three label sets: each grew 42 bytes, each stream's end returned them, and the pool then read 0 while every entry was still there and still colliding with later partitions. The plan's own drop produced no shrink at all. So a label set's bytes were charged for part of the time and invisible for the rest, and the invisibility grew with the number of label sets a query touches. The charge now belongs to the map it pays for: `LabelSets` holds the entries and the reservation together, `attach` registers it once on the first execution's pool, and it is released when the plan is dropped. `attach` is idempotent because every partition calls it. Re-registering would move a charge the other partitions keep growing, and would give it back when one stream ends while the entries stay - which is the shape of the bug. A regression test states the invariant: after the first partition's stream ends, the pool is still charged for its label set; a second partition with the same set on disjoint points adds no further charge; dropping the plan returns it to zero. Signed-off-by: discord9 * fix(promql): take a raw series frame, not a row batch, and keep its series whole Review found five things; this is all of them. The node read ordinary Arrow rows and converted them itself, so a real `SeriesScanExec` could never feed it: a frame's columns are `read_key`, `is_last`, and the `time` and `fields` structs, and none of them is the time column, a tag or a field the binding names, so construction refused it. It now takes a raw series frame the way `NativeInstantManipulateExec` does - the frame is validated when the plan is built and unpacked per partition when it runs - and the row-to-series mapping is gone, because turning rows into series is a different node's job. The series schema comes from the tree's own input node, so there is one statement of it rather than two that could disagree. A raw frame is a per-series stream, and this node holds one window per series and closes it on the chunk that marks the end. Declaring that as `KeyPartitioned(read_key)` keeps the optimizer from splitting a series across partitions, which it would otherwise be free to do; `read_key` is the series' primary key as a binary column, so hashing it groups exactly the rows of one series. A round-robin split would have let each partition rate half a series and answer half the truth; with the key declared, a series that somehow did arrive in two pieces is refused by the frame validator instead. Two label-set rules were wrong. Occupancy was tracked per position for any field, so two series that each hold a different field at the same point looked like a collision, while the rest of the stack merges exactly that as complementary fields - it is now tracked per field, and only one field holding two values at one point is refused. And the key encoded every tag including empty ones, while `grid_series::label_set` drops them, so `{host="h", zone=""}` and `{host="h"}` were two keys here and one everywhere else; the key is now that same normalized label set. Finally, a label set already in the map grew its run arrays on every merge without telling the pool. The charge now follows the capacity the merge adds, so a long-lived plan that keeps meeting one label set pays for what it holds. The test that asserted the opposite was wrong and is corrected. Signed-off-by: discord9 * fix(table): declare the distribution a raw series frame already has `SeriesScan` guarantees every series arrives whole in one partition - its own docs say so, and both modes enforce it in different ways: the two-phase collector places a series by its tsid, and the legacy divider only emits a batch once it holds all of a series' rows. But the leaf built for a raw frame declared `UnknownPartitioning`, so nothing on the plan said it. That mattered to a consumer holding one window per series: with the distribution unspecified an optimizer is free to round-robin the partition, and two halves of one series would each look final to their own stream. So the leaf now declares `Hash(read_key)`. `read_key` is the primary key as a binary column, so hashing it groups exactly the rows of one series and splits nothing; `FusedSeriesExec` asks for the same key and is satisfied without a shuffle. The three tests that pinned `UnknownPartitioning` now pin the key instead. Signed-off-by: discord9 * fix(promql): keep the capacity a label set already paid for `merge_runs` built a fresh vector sized to `stored + new` and replaced the old one. When adjacent runs joined - or when a claim added nothing to a field, which is what two complementary label sets do to each other - the entry came back smaller than it was. The charge in `claim` is the growth over the capacity, so the next merge into that field looked like growth again and was charged for it, every time, while the entry itself never held more than a couple of runs. A long query meeting one label set repeatedly would have paid for the same few bytes over and over and could exhaust a pool it never used. The merge now keeps whichever capacity is larger, and the charge follows that: a merge that adds nothing to a field costs nothing. Signed-off-by: discord9 * fix(promql): run the fused node over one series stream `FusedSeriesExec` held one window per series and closed it on the chunk carrying `is_last`, so it can answer a series only if nothing splits that run and nothing puts part of it in another tree. Those are two different promises, and only the first is a property of a stream: the frame decoder refuses a chunk interleaved with another series', but nothing it can see says that half a series is in a partition it never reads. The node used to ask for `KeyPartitioned(read_key)` instead, expecting the optimizer to keep whole series per partition. That does not hold: DataFusion considers two hash layouts with the same count co-partitioned, and above `subset_repartition_threshold` children it adds a hash exchange even when the requirement is met (`enforce_distribution.rs`, the `!partitioning_satisfied || target > child` branch). A hash exchange does not preserve order, so series chunks interleave and the second half of a series is answered as a series of its own. So the count is now fixed at one, in `try_new` and in the reported properties, and `benefits_from_input_partitioning` is false so no optimizer splits it for parallelism. This is the phase-1 contract: reading a series spread over partitions needs an exchange above this node, which does not exist yet. Merging #9786 also moved `unpack_raw_series` to two arguments, with the frame declaring its own label layout and ordering; the calls follow it, and the plan's binding is checked against that declaration with `ensure_frame_matches_label_source`. Signed-off-by: discord9 * Revert "fix(table): declare the distribution a raw series frame already has" The declaration was false. `Partitioning::Hash(read_key, n)` claims DataFusion placed rows by its own hash of `read_key` mod `n`, and DataFusion treats any two equal-count `Hash` layouts as co-partitioned, so a future co-partitioned consumer would pair partitions that hold unrelated series. `SeriesScan` places a series by tsid range (two-phase) or round-robin (legacy), neither of which is that hash. It also did not buy what it was added for. `FusedSeriesExec` now takes one input partition and reports `benefits_from_input_partitioning = false`, so nothing repartitions between the scan and it, and the ordering guarantee the declaration was meant to carry is the frame's own. This reverts commit c9d160f6a1b339980cf1c03290347472eafed045. Signed-off-by: discord9 * fix(promql): keep the input's partitions and forbid an exchange The node asked for `KeyPartitioned(read_key)` so the optimizer would keep whole series per partition. It does not: that is a hash requirement, and `hash_necessary = is_hash && multi_partitions` makes a hash exchange necessary as soon as the child has more than one partition - even where the child already satisfies the requirement. `with_preserve_order` then leaves `preserve_order` false, because the child declares no ordering, so the exchange reorders and the frame decoder closes a series it has not finished. Asking for nothing removes both ways an exchange can appear between the scan and this node: `is_hash` is false for `UnspecifiedDistribution`, so no hash exchange is necessary, and `benefits_from_input_partitioning` is false so no round-robin is added to raise parallelism. The child's layout then arrives as it left the scan, which is what this node needs: `SeriesScan` keeps a series in one partition, so one window per series is the right amount of state and the partitions run in parallel. So the node runs as many partitions as its input has, and only a zero-partition input is refused. The comments around the label-set map say the same thing again: the map is shared across partitions because two series that a range function and its maps reduce to one label set can land in different ones, and the runtime's per-node `Emitted` table never sees both. Signed-off-by: discord9 * test(promql): prove no exchange is inserted above the fused node The multi-partition contract rests on the optimizer leaving the input's layout alone, and the previous `KeyPartitioned` requirement failed exactly there: the rule inserted a hash `RepartitionExec` over `read_key` even though the child satisfied the requirement, because a hash requirement is necessary for any child with more than one partition. Reasoning about the flags is what missed it, so the claim is now checked against the rule. The new test runs `EnsureRequirements` over the node with a two-partition child and `target_partitions = 8`, and fails if any `RepartitionExec` sits between the input and the node. Ablating the requirement back to `KeyPartitioned(read_key)` makes it fail, naming the hash exchange the rule inserted at 8 partitions, so the test has the case it was written for; with `UnspecifiedDistribution` and `benefits_from_input_partitioning = false` it passes. The node also refuses, when it has more than one partition, a tree where a window or a map rewrites the series identity and a window above reads a whole series. The runtime does merge fragments that a rewrite made indistinguishable, but that merge belongs to one execution and this node runs one per partition, so each would merge its own half and the consumer would answer an empty result. The compiler cannot build that shape today - an outer window needs a subquery and `compile` refuses those - so the check costs no executable query. Signed-off-by: discord9 * test(promql): assert the optimizer leaves the partition count alone No exchange is not the same as no change. The ablation showed the old `KeyPartitioned` requirement not only inserted a hash exchange but raised the node from two partitions to eight, which is what `target_partitions` is set above the child for. A rule that changed the count some other way would split a series just as effectively, and the previous assertions only looked for a `RepartitionExec`. The test now also pins the node's partition count and its input's, both to the child's own two. Signed-off-by: discord9 * fix(promql): record no label set for a chunk holding no value A selection over a series that staleness or the range left empty still closes that series and hands over a closing chunk with no rows. Every chunk reached `LabelSets::claim`, which wrote an entry for it and charged the plan's reservation, and the plan holds that charge for its whole life: a query whose answer is empty grew the table and its reservation until a finite pool ran out, without producing a row. A chunk that occupies no position has no collision to remember and no entry to write, so it is skipped before the key is built. A later chunk under the same label set claims the positions it really holds, which a test now pins, together with the table staying empty over many empty series. Signed-off-by: discord9 * fix(promql): refuse a nested window over more than one partition `reject_a_split_logical_series` took a window to be a possible rewriter only when the function drops the metric name. That is not enough. `last_over_time` and `first_over_time` keep `__name__`, but a tag holding an empty value and an absent tag are one canonical label set and two raw keys, so their output can still hold two series under one label set. The runtime already knows this: its `holds_shared` pass computes `drops_metric_name || !input_holds_shared` for a window, so a window that did not merge its input leaves fragments for the window above it. That merge belongs to one execution, and this node runs one execution per partition. A bare tree says nothing about whether an input was merged, so every window is now taken to be a possible rewriter, and the existing shape check narrows it to the case that costs correctness: a window with a window above it. The refusal still applies only over more than one partition; one partition returns before the walk and keeps working. The compiler requires a range function to apply directly to a selector, so no nested-window tree is executable today and this costs no executable query. The regression test pins the hole: with the old predicate restored it fails on the two-partition tree, and the same tree over one partition is still accepted. Signed-off-by: discord9 * fix(promql): share the run bookkeeping so a sparse claim stays linear Every output chunk went through `LabelSets::claim`, which scanned the stored runs from index 0 and re-merged all of them into a fresh vector. A legal sparse stream - one grid point per chunk with a hole after it - therefore cost the square of the chunk count, under a mutex every partition waits on. Measured on the two functions alone: 4,000 chunks took about 58 ms, 16,000 about 962 ms. The native runtime had already solved this for its own `Emitted` table, so the two implementations are now one. `Run`, `absorb`, `append_runs`, `merge_runs` and the test-only rebuild counter move to `native::runs`, and both layers call them: `LabelSets` stores `Vec>`, looks a position up with `partition_point` instead of scanning, and appends through `absorb`, which only rebuilds when a chunk reaches back into what is stored. Keeping one copy is the point - the Fused layer had fallen behind the runtime's version, and two copies can drift again. The charge follows the vectors' capacity across the absorb, so the append path's real growth is what gets charged and nothing already paid for is charged twice. Two tests pin it: 64 forward chunks rebuild nothing, and a chunk that reaches back still merges (and still refuses a genuine collision). Signed-off-by: discord9 * fix(promql): pull the fused node's answers one at a time after its input ends `close()` called `finish()`, which ends the input and then drains what the tree settled into one `Vec` before anything is handed out. A tree that answers many chunks after its input ends therefore held all of them while the consumer waited for the first, and charged the pool for them at once - the opposite of what the runtime's `end_of_input` / `next_chunk` pair is for. The stream now ends the input once and pulls one chunk per poll, through the execution's own protocol: `end_of_input()` when the input ends, then `next_chunk()` per answer, and `None` - which is also when the execution closes and returns its reservations - is what ends the stream. The input is dropped as soon as it ends so a poll goes straight to pulling. Signed-off-by: discord9 * fix(promql): preserve fused claims on memory refusal Signed-off-by: discord9 * test(promql): verify retained fused claim allocations Signed-off-by: discord9 * test(promql): keep the pool blocker immutable Signed-off-by: discord9 * test(promql): cover a claim that shrinks its run vector Signed-off-by: discord9 * fix(promql): refuse an instant split under a window Signed-off-by: discord9 * fix(promql): read grid timestamps through their fallible accessor Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/promql/src/extension_plan.rs | 2 + src/promql/src/extension_plan/fused_series.rs | 1555 +++++++++ .../src/extension_plan/fused_series/tests.rs | 2854 +++++++++++++++++ src/promql/src/native/execute.rs | 120 +- src/promql/src/native/mod.rs | 1 + src/promql/src/native/runs.rs | 146 + 6 files changed, 4561 insertions(+), 117 deletions(-) create mode 100644 src/promql/src/extension_plan/fused_series.rs create mode 100644 src/promql/src/extension_plan/fused_series/tests.rs create mode 100644 src/promql/src/native/runs.rs diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 750ee41b5e8..baf57a275c8 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -14,6 +14,7 @@ mod absent; mod empty_metric; +mod fused_series; mod histogram_fold; mod instant_manipulate; mod normalize; @@ -37,6 +38,7 @@ use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; +pub use fused_series::{FusedSeries, FusedSeriesBinding, FusedSeriesExec, FusedSeriesStream}; pub use histogram_fold::{ HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream, }; diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs new file mode 100644 index 00000000000..d957851902a --- /dev/null +++ b/src/promql/src/extension_plan/fused_series.rs @@ -0,0 +1,1555 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! One node that runs a compiled native tree over raw series. +//! +//! [`FusedSeries`] is the logical half: it holds the PromQL text the tree was compiled from, the +//! tree itself and the binding its input was compiled against. [`FusedSeriesExec`] is the physical +//! half: it owns one execution of that tree per partition and decodes its input frames into raw +//! series chunks and converts its evaluated chunks back into batches. +//! +//! The two halves show different things in `EXPLAIN`, which is the point of having both. The +//! logical node prints the query a reader wrote ([`FusedSeries`]'s `fmt_for_explain`), and the +//! physical node prints the tree it will run ([`FusedSeriesExec`]'s `fmt_as`, which embeds +//! [`NativeTree`]'s own display rather than a second renderer). A reader comparing the two plans +//! can therefore see whether a query was compiled the way they expected without decoding node +//! names. +//! +//! Scope of this layer: the input is a stream of native raw series frames +//! ([`raw_series_frame_schema`](datatypes::series::frame::raw_series_frame_schema)) — one row per +//! chunk, carrying the series' encoded key, its finality, its timestamps and its fields, plus the +//! complete labels when the frame is explicit — with each series' rows contiguous and its +//! timestamps ascending, which is what the caller's [`SeriesOrdering`] names. The output is one +//! batch per evaluated chunk, with the statement's grid timestamps in the declared time column. +//! Everything the tree cannot do is rejected when the execution plan is built, not when the first +//! batch arrives: a tree the native runtime does not execute, a frame whose schema the tree's input +//! was not compiled against, or an output schema the conversion cannot classify. + +use std::cmp::Ordering; +use std::collections::hash_map::Entry; +use std::collections::{HashMap, VecDeque}; +use std::fmt::{self, Formatter}; +use std::hash::{Hash, Hasher}; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use common_query::series_frame::{ensure_frame_matches_label_source, unpack_raw_series}; +use datafusion::arrow::array::{ + Array, ArrayRef, Float64Array, StringArray, TimestampMillisecondArray, +}; +use datafusion::arrow::buffer::ScalarBuffer; +use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::DFSchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, +}; +use datatypes::series::SeriesOrdering; +use datatypes::series::raw::RawSeriesChunk; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::time::EvalGrid; +use futures::Stream; +use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream}; + +use crate::evaluated::EvaluatedChunk; +use crate::grid_series::{LabelSet, label_set}; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::metrics::{self, NodePath}; +use crate::native::plan::{NativeExpr, NativeTree}; +use crate::native::runs::{Run, absorb}; +use crate::source::materialize_raw_labels; + +#[cfg(test)] +mod tests; + +/// The label a metric name is stored under, which is not a tag column of the input. +const METRIC_NAME: &str = "__name__"; + +/// How many input steps one poll of [`FusedSeriesStream`] runs before it wakes itself. +/// +/// A step polls the input once and hands what it settled to the executor. Settling is not +/// guaranteed: raw samples before the first evaluation point settle nothing, so an input of many +/// such chunks would otherwise be read to its end inside one poll, holding the task's thread for +/// an unbounded time and delaying both other tasks on it and the cancellation the consumer asked +/// for. The budget makes that a poll that yields instead, and the waker makes sure it is polled +/// again, so nothing is lost by stopping early. +const POLL_BUDGET: usize = 64; + +/// What the tree's input was compiled against, and how its output columns are named. +/// +/// The binding travels with the tree: a tree is legal only against the schema it was bound to, and +/// a reader of the output needs to know which grid a chunk's absolute positions are on. Keeping it +/// in one value is what lets the logical node be built from the same description, and the physical +/// node reads the same facts back off the tree rather than trusting a second copy of them. +/// +/// What the logical node checks is that the schema the tree's input node declares is the one this +/// binding names. What it cannot check is whether the caller meant *some other* schema with the +/// same field names, or whether [`time_index`](Self::time_index) is the column the query reports +/// positions in: from here both are indistinguishable from a correct binding, so picking them is +/// the caller's obligation. The physical node reads neither the metric name nor a time column from +/// the input — a raw frame carries the labels of each series and its own timestamps. +#[derive(Clone, Debug)] +pub struct FusedSeriesBinding { + /// The raw series schema the tree's input node was bound to, without an evaluation grid. + pub series_schema: Arc, + /// The statement's grid, which is the grid every evaluated chunk is addressed on. + pub grid: Arc, + /// The metric name every input series is bound to, when the input reads one named metric. + /// + /// The logical node requires it when [`tags`](Self::tags) declares [`METRIC_NAME`]: a reader of + /// the output has to be able to tell that a range function dropped the name, which needs a name + /// to drop. The physical node takes the names from the frames, which carry each series' own. + pub metric_name: Option>, + /// The output column holding the grid timestamps. + pub time_index: String, + /// The logical tag names the output carries, read from each series' labels. + /// + /// [`METRIC_NAME`] is special: it is resolved to the series' metric name instead of a column, + /// which is how a reader sees that a range function dropped it. + pub tags: Vec, +} + +impl FusedSeriesBinding { + /// Whether `other` binds the same fields under the same names. + fn same_fields(&self, other: &SeriesSchema) -> bool { + self.series_schema.fields.len() == other.fields.len() + && self + .series_schema + .fields + .iter() + .zip(other.fields.iter()) + .all(|(mine, theirs)| mine.name == theirs.name) + } +} + +impl PartialEq for FusedSeriesBinding { + fn eq(&self, other: &Self) -> bool { + // The schemas are compared by the field names they declare, and the grids by their spec: + // neither type has structural equality, and neither identity nor allocation is meant here. + self.metric_name == other.metric_name + && self.time_index == other.time_index + && self.tags == other.tags + && self.grid.spec() == other.grid.spec() + && self.same_fields(&other.series_schema) + } +} + +impl Eq for FusedSeriesBinding {} + +impl Hash for FusedSeriesBinding { + fn hash(&self, state: &mut H) { + self.metric_name.hash(state); + self.time_index.hash(state); + self.tags.hash(state); + let spec = self.grid.spec(); + (spec.start_ms, spec.step_ms, spec.len).hash(state); + for field in self.series_schema.fields.iter() { + field.name.hash(state); + } + } +} + +impl PartialOrd for FusedSeriesBinding { + fn partial_cmp(&self, other: &Self) -> Option { + // A total order over the same fields equality reads, so equal bindings compare equal. + match self + .metric_name + .as_deref() + .partial_cmp(&other.metric_name.as_deref()) + { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.time_index.partial_cmp(&other.time_index) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tags.partial_cmp(&other.tags) { + Some(Ordering::Equal) => {} + ord => return ord, + } + let (mine, theirs) = (self.grid.spec(), other.grid.spec()); + match (mine.start_ms, mine.step_ms, mine.len).partial_cmp(&( + theirs.start_ms, + theirs.step_ms, + theirs.len, + )) { + Some(Ordering::Equal) => {} + ord => return ord, + } + Some( + self.series_schema + .fields + .iter() + .map(|field| field.name.as_ref()) + .cmp( + other + .series_schema + .fields + .iter() + .map(|field| field.name.as_ref()), + ), + ) + } +} + +/// The logical plan of one node that runs a compiled native tree. +/// +/// The node adds nothing to the tree: it records the query the tree was compiled from, so +/// `EXPLAIN` shows the reader their own query rather than the shape it compiled to, and it carries +/// the binding the tree's input was compiled against so the execution plan can be built from it. +#[derive(Debug)] +pub struct FusedSeries { + /// The PromQL text this tree was compiled from. + promql: Arc, + /// The compiled tree to run. + tree: NativeTree, + /// What the tree's input was compiled against. + binding: FusedSeriesBinding, + /// The plan producing the input series. + input: LogicalPlan, + /// The columns this node produces. + output_schema: DFSchemaRef, +} + +impl FusedSeries { + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeries" + } + + /// Creates the node, checking that `tree` is legal against `binding`. + pub fn try_new( + promql: Arc, + tree: NativeTree, + binding: FusedSeriesBinding, + input: LogicalPlan, + output_schema: DFSchemaRef, + ) -> DataFusionResult { + validate_binding(&tree, &binding)?; + Ok(Self { + promql, + tree, + binding, + input, + output_schema, + }) + } + + /// The PromQL text this node was compiled from. + pub fn promql(&self) -> &str { + &self.promql + } + + /// The compiled tree this node runs. + pub fn tree(&self) -> &NativeTree { + &self.tree + } + + /// What the tree's input was compiled against. + pub fn binding(&self) -> &FusedSeriesBinding { + &self.binding + } + + /// The physical node that runs `exec_input` as this node's input. + /// + /// `label_source` and `ordering` are the binding of the frames `exec_input` produces: the + /// frame carries the encoded key and, in explicit mode, the complete labels, but no schema can + /// say which of the two a producer meant or how it ordered the series, so the caller that built + /// the frame passes both on. Building the node rejects a tree the runtime cannot execute, a + /// frame schema the tree's input was not compiled against, and a schema the conversion cannot + /// classify, so a plan that is built is a plan that runs. + pub fn to_execution_plan( + &self, + exec_input: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + ) -> DataFusionResult> { + let output_schema: SchemaRef = self.output_schema.inner().clone(); + Ok(Arc::new(FusedSeriesExec::try_new( + Arc::clone(&self.promql), + exec_input, + label_source, + ordering, + Arc::new(self.tree.clone()), + output_schema, + self.binding.time_index.clone(), + self.binding.tags.clone(), + )?)) + } +} + +impl UserDefinedLogicalNodeCore for FusedSeries { + fn name(&self) -> &str { + Self::name() + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.output_schema + } + + fn expressions(&self) -> Vec { + // The tree's operands are bound inside it; there is no expression for an optimizer to + // rewrite separately from the tree that owns it. + Vec::new() + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeries: {}", self.promql) + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> DataFusionResult { + if inputs.len() != 1 { + return Err(DataFusionError::Internal( + "FusedSeries should have at exact one input".to_string(), + )); + } + + let input = inputs.pop().unwrap(); + Ok(Self { + promql: Arc::clone(&self.promql), + tree: self.tree.clone(), + binding: self.binding.clone(), + input, + output_schema: Arc::clone(&self.output_schema), + }) + } +} + +impl PartialEq for FusedSeries { + fn eq(&self, other: &Self) -> bool { + // `NativeTree` carries no structural equality, so the rendered tree stands in for it. The + // renderer names every node with its bound operation and grids, which is what makes two + // nodes the same computation; comparing the text alone would call two different trees + // equal whenever their PromQL happens to read the same. + self.promql == other.promql + && self.tree.to_string() == other.tree.to_string() + && self.binding == other.binding + && self.input == other.input + && self.output_schema == other.output_schema + } +} + +impl Eq for FusedSeries {} + +impl Hash for FusedSeries { + fn hash(&self, state: &mut H) { + self.promql.hash(state); + self.binding.hash(state); + self.input.hash(state); + self.output_schema.hash(state); + } +} + +impl PartialOrd for FusedSeries { + fn partial_cmp(&self, other: &Self) -> Option { + match self.promql.partial_cmp(&other.promql) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.binding.partial_cmp(&other.binding) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.input.partial_cmp(&other.input) { + Some(Ordering::Equal) => {} + ord => return ord, + } + match self.tree.to_string().partial_cmp(&other.tree.to_string()) { + Some(Ordering::Equal) => {} + ord => return ord, + } + if self.output_schema == other.output_schema { + Some(Ordering::Equal) + } else { + None + } + } +} + +/// What the tree's input node reads, read off the tree itself. +/// +/// The tree is the one authority on the schema its input was compiled against: a caller's binding +/// can only agree or disagree with it, and the frame this node reads has to be decoded against the +/// same schema, so both halves read it here instead of each carrying a copy. +struct TreeInput { + /// The position the tree's raw chunks enter at: a position is how the runtime addresses a node, + /// so nothing on the node has to carry an identity that the tree itself does not. + path: NodePath, + /// The raw series schema the tree's input node was compiled against. + schema: Arc, +} + +/// Reads what the tree's input node was compiled against, rejecting a tree this node cannot feed. +/// +/// A node feeds one raw input, and its output is addressed on a grid; a tree that does not have +/// that shape is one this layer cannot hand chunks to or read positions from. +fn tree_input(tree: &NativeTree) -> DataFusionResult { + let mut sources = 0; + let mut input_path = None; + let mut schema = None; + metrics::walk(tree.root(), &mut |node, path| { + let NativeExpr::Input(input) = &node.expr else { + return; + }; + sources += 1; + input_path.get_or_insert_with(|| path.clone()); + // Raw, not evaluated: this node feeds the tree from a scan, so a tree whose input was + // compiled against an already-evaluated schema would read the wrong thing. + schema = input + .output + .samples() + .filter(|schema| schema.eval_grid.is_none()) + .map(Arc::clone); + }); + if sources != 1 { + return Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree with {sources} inputs, and it needs exactly one" + ))); + } + let (Some(path), Some(schema)) = (input_path, schema) else { + return Err(DataFusionError::Plan( + "FusedSeries runs a tree whose input is not raw series".to_string(), + )); + }; + let evaluated = tree + .root() + .output + .samples() + .is_some_and(|schema| schema.eval_grid.is_some()); + if !evaluated { + return Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on a grid".to_string(), + )); + } + Ok(TreeInput { path, schema }) +} + +/// Refuses a tree whose logical series can be split across the partitions this node runs. +/// +/// Keeping the metric name is not what keeps two raw series apart. An empty tag value and an absent +/// tag are the same canonical label set, so `{host="h", zone=""}` and `{host="h"}` are one series to +/// Prometheus while the scan encodes them as two raw keys. A window that evaluated its input without +/// merging it therefore passes such a pair through as two series under one label set, exactly as the +/// functions that drop `__name__` do by merging them, and a map drops `__name__` outright. +/// +/// A stateful node above one of those needs both halves together, and the runtime does provide that +/// within one execution: a window whose input can hold two series under one label set collects the +/// fragments, merges each label set once, and only then evaluates. That merge is per execution, and +/// this node runs one execution per partition. Two raw series that become one label set can land in +/// different partitions - the scan keeps a *series* in one partition, which says nothing about two +/// different series that later become indistinguishable - so each execution merges only the half it +/// was given. The consumer above then answers from one half and reports nothing for the other, which +/// is an empty result rather than an error. +/// +/// Whether a window merged its input is runtime state this check does not have: it runs over a bare +/// tree, where a window's input is only something that produces samples. A window below another +/// window is therefore treated as one that did not merge, whatever function it runs, which refuses +/// `rate` under `rate` (which drops the name) and `last_over_time` under `rate` (which keeps it) +/// alike. +/// +/// Reading such a tree needs an exchange above this node that merges a logical series across +/// partitions, which does not exist yet. Until it does the tree is refused where it can hurt: over +/// one partition there is nothing to split, so only a tree run over more than one is checked. +/// +/// The compiler cannot build a nested window today - [`compile`](crate::promql::native_compile::compile) +/// requires a range function to apply directly to a selector, `([])`, +/// and a selector is raw samples - so no tree it produces has a window above a window, and this +/// check costs no executable query. It is checked because [`FusedSeriesExec::try_new`] accepts any +/// well-formed tree, and a tree it accepts is a tree it must answer correctly. +fn reject_a_split_logical_series( + tree: &NativeTree, + partition_count: usize, +) -> DataFusionResult<()> { + if partition_count == 1 { + return Ok(()); + } + // The walk is outermost first, so it meets nodes in the reverse of the order data flows: a + // consumer met before a rewriter is one that the rewriter feeds. + let mut consumer: Option = None; + let mut rejection: Option<(String, String)> = None; + metrics::walk(tree.root(), &mut |node, path| { + if rejection.is_some() { + return; + } + // A node can leave two raw series on one canonical label set. A map always can: it drops + // `__name__` and every empty-valued tag. An instant selection preserves its input labels, + // rather than dropping them, but can still leave two raw series on one canonical label set: + // `{zone=""}` and the absent `zone` are one label set and two raw keys. A window can + // whenever it did not merge its input, which a bare tree says nothing about: a check on the + // function alone accepts `last_over_time`, whose output keeps the name and can still hold + // both raw series. Every map, instant selection, and window is therefore taken to be a + // possible rewriter, and the refusal below is what narrows that to the shape that costs + // correctness - a window with a rewriter below it, a window being the only consumer here. + let rewrites = matches!( + node.expr, + NativeExpr::Map { .. } | NativeExpr::Instant { .. } | NativeExpr::Window { .. } + ); + if rewrites && let Some(above) = &consumer { + rejection = Some((path.as_str().to_string(), above.clone())); + return; + } + let consumes = matches!(node.expr, NativeExpr::Window { .. }); + if consumes && consumer.is_none() { + consumer = Some(path.as_str().to_string()); + } + }); + let Some((rewriter, above)) = rejection else { + return Ok(()); + }; + Err(DataFusionError::Plan(format!( + "FusedSeries runs a tree where `{rewriter}` can leave two raw series on one canonical label \ + set and `{above}` above it reads a whole series, over {partition_count} partitions: an \ + empty tag value and an absent tag are the same label set, the runtime merges such \ + fragments once per execution, and this node runs one execution per partition, so `{above}` \ + would read one half and answer an empty result. Reading this tree needs an exchange between \ + `{rewriter}` and `{above}` that merges a logical series across partitions, which is not \ + implemented" + ))) +} + +/// Checks that `tree` reads the schema `binding` declares, on the grid it declares. +/// +/// The output grid is compared by [`EvalGrid::spec`], not by allocation: that is the rule +/// [`NativeTree`]'s own validation and the kernels apply, and a chunk's timestamps already come +/// from its absolute grid positions, so a separately allocated grid with the same spec is the same +/// grid to everything downstream. +/// +/// Answers with what the tree's input reads, which is what a caller names when it hands a raw chunk +/// over. +fn validate_binding( + tree: &NativeTree, + binding: &FusedSeriesBinding, +) -> DataFusionResult { + if binding.series_schema.eval_grid.is_some() { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema, which cannot have an evaluation grid" + .to_string(), + )); + } + if binding.tags.iter().any(|tag| tag == METRIC_NAME) && binding.metric_name.is_none() { + return Err(DataFusionError::Plan(format!( + "FusedSeries declares a {METRIC_NAME:?} tag column but binds no metric name" + ))); + } + + // The tree's own input node says which schema it reads, so the binding is checked against the + // tree rather than trusted: a mismatch would evaluate one schema and label the result as + // another. + let input = tree_input(tree)?; + if !binding.same_fields(&input.schema) { + return Err(DataFusionError::Plan( + "FusedSeries binds a raw series schema the tree's input was not compiled against" + .to_string(), + )); + } + match tree.root().output.samples() { + Some(schema) + if schema + .eval_grid + .as_ref() + .is_some_and(|grid| grid.spec() == binding.grid.spec()) => + { + Ok(input) + } + _ => Err(DataFusionError::Plan( + "FusedSeries runs a tree whose output is not evaluated on the statement's grid" + .to_string(), + )), + } +} + +/// Where one output column of [`FusedSeriesExec`] comes from. +#[derive(Clone, Copy, Debug)] +enum OutputColumn { + /// The evaluation grid timestamps of the current chunk. + Time, + /// Series field by index in the bound series schema. + Field(usize), + /// Declared logical tag by index in the caller-supplied tag list. + Tag(usize), +} + +/// Classifies every output column, rejecting anything the conversion cannot produce. +fn classify_output( + output_schema: &SchemaRef, + series_schema: &SeriesSchema, + tags: &[String], + time_column: &str, +) -> DataFusionResult> { + for (index, tag) in tags.iter().enumerate() { + if tags[..index].contains(tag) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {tag:?} in FusedSeriesExec output schema" + ))); + } + if tag == time_column { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec cannot be the time column" + ))); + } + if series_schema + .fields + .iter() + .any(|field| field.name.as_ref() == tag) + { + return Err(DataFusionError::Plan(format!( + "tag column {tag:?} of FusedSeriesExec clashes with a series field" + ))); + } + } + + let mut time_seen = false; + let mut field_seen = vec![false; series_schema.fields.len()]; + let mut tag_seen = vec![false; tags.len()]; + let mut classification = Vec::with_capacity(output_schema.fields().len()); + for field in output_schema.fields() { + let name = field.name(); + if name == time_column { + if std::mem::replace(&mut time_seen, true) { + return Err(DataFusionError::Plan(format!( + "duplicate time column {name:?} in FusedSeriesExec output schema" + ))); + } + match field.data_type() { + DataType::Timestamp(TimeUnit::Millisecond, None) => {} + other => { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must be Timestamp(Millisecond, \ + None), got {other:?}" + ))); + } + } + if field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "time column {name:?} of FusedSeriesExec must not be nullable" + ))); + } + classification.push(OutputColumn::Time); + continue; + } + if let Some(index) = series_schema + .fields + .iter() + .position(|field| field.name.as_ref() == name) + { + if std::mem::replace(&mut field_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate field column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Float64 { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be Float64, got {:?}", + field.data_type() + ))); + } + // A grid point inside the statement's interval can have no value in this field. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "field column {name:?} of FusedSeriesExec must be nullable, because a grid \ + point can be missing" + ))); + } + classification.push(OutputColumn::Field(index)); + continue; + } + if let Some(index) = tags.iter().position(|tag| tag == name) { + if std::mem::replace(&mut tag_seen[index], true) { + return Err(DataFusionError::Plan(format!( + "duplicate tag column {name:?} in FusedSeriesExec output schema" + ))); + } + if field.data_type() != &DataType::Utf8 { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be Utf8, got {:?}", + field.data_type() + ))); + } + // `rate` and the map kernels drop `__name__`, so a tag can be absent from a row. + if !field.is_nullable() { + return Err(DataFusionError::Plan(format!( + "tag column {name:?} of FusedSeriesExec must be nullable, because range \ + functions and map kernels drop the metric name" + ))); + } + classification.push(OutputColumn::Tag(index)); + continue; + } + return Err(DataFusionError::Plan(format!( + "unknown column {name:?} in FusedSeriesExec output schema" + ))); + } + + if !time_seen { + return Err(DataFusionError::Plan(format!( + "missing time column {time_column:?} in FusedSeriesExec output schema" + ))); + } + for (index, seen) in field_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing field column {:?} in FusedSeriesExec output schema", + series_schema.fields[index].name + ))); + } + } + for (index, seen) in tag_seen.iter().enumerate() { + if !seen { + return Err(DataFusionError::Plan(format!( + "missing tag column {:?} in FusedSeriesExec output schema", + tags[index] + ))); + } + } + Ok(classification) +} + +/// The physical node that runs a compiled native tree over its input's raw series. +pub struct FusedSeriesExec { + /// The PromQL text the tree was compiled from, which is what the plan line shows. + promql: Arc, + /// The tree this node runs, shared by every partition's execution. + tree: Arc, + /// The raw series schema the tree's input was compiled against, which every frame is decoded + /// against. + series_schema: Arc, + /// How the input frames name their series: labels encoded in the key, or carried along. + label_source: RawLabelSource, + /// The ordering the input frames guarantee. + ordering: SeriesOrdering, + /// The position of the tree's input node, where every raw chunk a partition reads enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The output column holding the grid timestamps, kept so this node can be rebuilt over a new + /// child without losing which column its time came from. + time_column: String, + /// The tag column names, indexed by [`OutputColumn::Tag`]. + tags: Arc<[String]>, + /// The plan producing the raw series frames. + input: Arc, + /// The columns this node produces. + output_schema: SchemaRef, + /// The grid positions every label set in this plan's output already occupies, and what they + /// keep charged on the memory pool. + /// + /// Plan-level rather than per-partition: see [`LabelSets`]. + labels: SharedLabelSets, + /// Per-node metrics, registered by each execution of the tree. + metric: ExecutionPlanMetricsSet, + /// This node's partitioning, emission and boundedness. + properties: Arc, +} + +impl FusedSeriesExec { + /// Creates the node, rejecting whatever this node cannot execute or convert. + /// + /// `input` produces native raw series frames: the series' encoded key, its finality, its + /// timestamps and its fields, plus the complete labels when `label_source` is + /// [`RawLabelSource::Explicit`]. `time_column` and `tags` name the columns of the output, which + /// is a row table rather than a frame. + /// + /// The parameters stay separate because each is a fact its own caller has: the query, the plan + /// below, the frame's label binding and ordering, the tree, and the three things the output + /// schema is classified by. A struct over them would be the same list under another name. + #[allow(clippy::too_many_arguments)] + pub fn try_new( + promql: Arc, + input: Arc, + label_source: RawLabelSource, + ordering: SeriesOrdering, + tree: Arc, + output_schema: SchemaRef, + time_column: String, + tags: Vec, + ) -> DataFusionResult { + let tree_input = tree_input(&tree)?; + // Building a state is what proves the runtime executes this tree: a node kind it does not + // run yet is refused here, so a plan that is built is a plan that starts. The throwaway + // state is dropped at once, returning its (unbounded) reservations. + let _ = NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + )?; + let classification = + classify_output(&output_schema, &tree_input.schema, &tags, &time_column)?; + // Read here rather than below, because whether the split below can happen decides whether + // the tree shape is one this node may run at all. + let input_properties = input.properties(); + let partition_count = input_properties.partitioning.partition_count(); + if partition_count == 0 { + return Err(DataFusionError::Plan( + "FusedSeriesExec input has no partitions".into(), + )); + } + reject_a_split_logical_series(&tree, partition_count)?; + // The frame states its own label layout and series ordering; this plan's binding is an + // expectation checked against that declaration, not an input to decoding. + ensure_frame_matches_label_source(input.schema().as_ref(), &label_source)?; + // The frame is checked against what the tree's input reads, so a frame this node cannot + // decode is reported when the plan is built. The shared decoder is the authority on that + // layout, so it is asked rather than copied: handing it an empty stream of this schema makes + // it validate the declared frame with no batch to decode, which is the check + // `unpack_raw_series` makes before its first poll too. + unpack_raw_series( + Box::pin(RecordBatchStreamAdapter::new( + input.schema(), + futures::stream::empty(), + )), + Some(Arc::clone(&tree_input.schema)), + )?; + // One evaluator per input partition, and the input's own partitioning kept as it is. + // + // A window holds the state of the series it is reading and closes it on the chunk carrying + // `is_last`, so it can answer a series only if the whole series arrives in the partition it + // reads. The input guarantees that: a raw series frame is produced by `SeriesScan`, whose + // contract is that the same primary key always reaches the same partition, and the frame + // decoder holds each series' chunks contiguous within that partition. + // + // What is *not* asked of the optimizer is a distribution. A `KeyPartitioned` or + // `HashPartitioned` requirement is classed as a hash requirement, and a hash requirement + // makes a hash exchange necessary whenever the child has more than one partition - even if + // the child already satisfies it, and even when the partition counts match. A hash exchange + // does not preserve order, so it would interleave the chunks of one series and the decoder + // would close a series it has not finished. `UnspecifiedDistribution` leaves the child's + // layout alone, and `benefits_from_input_partitioning` is false so no round-robin is added + // to raise parallelism either. Between them no exchange can appear between the scan and this + // node, and the scan's guarantee arrives intact. + // As many partitions as the input hands over: the node runs one evaluator per partition and + // answers each as it settles. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(partition_count), + input_properties.emission_type, + input_properties.boundedness, + )); + // Shared by every partition's stream: a label set colliding is a property of the whole + // output vector, not of one partition. Building it needs no pool; the owner's reservation is + // attached in `execute`, which is the first point that has one. + let labels: SharedLabelSets = Arc::new(Mutex::new(LabelSets::default())); + Ok(Self { + promql, + tree, + series_schema: tree_input.schema, + label_source, + ordering, + input_path: tree_input.path, + classification: classification.into(), + time_column, + tags: tags.into(), + input, + output_schema, + labels, + metric: ExecutionPlanMetricsSet::new(), + properties, + }) + } + + /// The name this node is shown under. + pub const fn name() -> &'static str { + "FusedSeriesExec" + } +} + +impl std::fmt::Debug for FusedSeriesExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("FusedSeriesExec").finish() + } +} + +impl DisplayAs for FusedSeriesExec { + /// The query, then the tree it compiles to: the tree is rendered by [`NativeTree`]'s own + /// `Display`, so `EXPLAIN` and the node's metrics name the same nodes. + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "FusedSeriesExec: {}\n{}", self.promql, self.tree) + } +} + +impl ExecutionPlan for FusedSeriesExec { + fn name(&self) -> &str { + Self::name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + /// Nothing is asked of the input's distribution, and nothing may be inserted above it. + /// + /// A hash requirement would make a hash exchange necessary as soon as the child has more than + /// one partition - even where the child already satisfies the requirement - and a hash exchange + /// does not preserve order. The input's own layout is already the one this node needs, because + /// the scan it reads from guarantees a series stays in one partition, so the requirement says + /// only that the layout is left alone. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(vec![Distribution::UnspecifiedDistribution]) + } + + /// Splitting the input further would not help: a second partition is a second set of series + /// answered by a second tree, which this node already runs, and raising the count is what an + /// optimizer does by adding a round-robin exchange that reorders nothing it can rely on. + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + // The tree's operands are bound inside it; this node owns no expression of its own. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + let [child] = <[Arc; 1]>::try_from(children).map_err(|children| { + DataFusionError::Plan(format!( + "FusedSeriesExec requires exactly one child, got {}", + children.len() + )) + })?; + // Rebuilt rather than patched: the new child's schema is what the frame is validated + // against, so a child that changed shape is refused instead of decoded with a stale layout. + Ok(Arc::new(Self::try_new( + Arc::clone(&self.promql), + child, + self.label_source.clone(), + self.ordering, + Arc::clone(&self.tree), + self.output_schema.clone(), + self.time_column.clone(), + self.tags.to_vec(), + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + // The label-set state is the plan's, so the reservation that pays for it is attached exactly + // once, to the pool of the first execution. Every partition calls this and only the first + // takes effect: re-registering per stream would move a charge the other partitions keep + // growing, and would give it back when that stream ends while the entries stay. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .attach(&context.runtime_env().memory_pool); + let input = self.input.execute(partition, Arc::clone(&context))?; + // Batches become raw series chunks through the shared frame decoder, which is the inverse of + // what the producer packed: one chunk per frame row, each series closed exactly once by the + // row carrying `is_last`. + let input = unpack_raw_series(input, Some(Arc::clone(&self.series_schema)))?; + // A source bound in key mode carries its identity encoded and no decoded labels, and a chunk + // handed to the tree has to name the series it holds. + let input = + materialize_raw_labels(input).map_err(|e| DataFusionError::External(Box::new(e)))?; + let state = NativeExecutionState::try_new( + &self.tree, + Arc::clone(&context.runtime_env().memory_pool), + &self.metric, + partition, + )?; + // What the label sets hold is charged to the reservation attached above, so nothing is + // reserved per stream: the entries outlive every stream and are released with the plan. + Ok(Box::pin(FusedSeriesStream { + input: Some(input), + executor: Some(NativeExecutor::new(state)), + input_path: self.input_path.clone(), + classification: Arc::clone(&self.classification), + tags: Arc::clone(&self.tags), + output_schema: self.output_schema.clone(), + labels: Arc::clone(&self.labels), + ready: VecDeque::new(), + done: false, + })) + } + + fn metrics(&self) -> Option { + Some(self.metric.clone_inner()) + } +} + +/// Converts one evaluated chunk into one output batch. +/// +/// The time column is the absolute grid positions the chunk carries, missing samples are nulls and +/// ordinary NaN stays a value. +fn series_to_record_batch( + classification: &[OutputColumn], + tags: &[String], + output_schema: &SchemaRef, + chunk: &EvaluatedChunk, +) -> DataFusionResult { + let series = &chunk.series; + // A chunk that forms one run is read from its range, which lends the grid's buffers and so + // allocates nothing. A chunk whose rows form several runs has no single range to read + // ([`EvaluatedChunk::range`] answers `None` for it), and the positions between its runs are + // holes rather than rows: the axis is walked for its carried rows instead, which is the one + // traversal that visits exactly those. + let timestamps = match chunk.range() { + Some(range) => range.timestamps(), + None => series + .time + .timestamps() + .map(|timestamps| ScalarBuffer::from(timestamps.into_owned())), + } + .map_err(|error| DataFusionError::Execution(error.to_string()))?; + let rows = timestamps.len(); + let mut columns: Vec = Vec::with_capacity(classification.len()); + for (position, column) in classification.iter().enumerate() { + match column { + OutputColumn::Time => columns.push(Arc::new(TimestampMillisecondArray::new( + timestamps.clone(), + None, + ))), + OutputColumn::Field(index) => { + let samples = series.fields.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "evaluated series has no field at index {index}" + )) + })?; + if samples.len() != rows { + return Err(DataFusionError::Execution(format!( + "evaluated series field {index} has {} rows, time axis has {rows}", + samples.len() + ))); + } + columns.push(field_column(samples)?); + } + OutputColumn::Tag(index) => { + // `__name__` resolves to the metric name, every other tag through the label set. + let name = tags.get(*index).ok_or_else(|| { + DataFusionError::Execution(format!( + "output column {position} declares unavailable tag {index}" + )) + })?; + let value = series.metric.get(name); + columns.push(Arc::new(StringArray::from(vec![value; rows]))); + } + } + } + Ok(RecordBatch::try_new(output_schema.clone(), columns)?) +} + +/// Converts one evaluated series field into one `Float64` column. +/// +/// Dense float samples are reused as they are: the array already carries exactly the right nulls +/// and keeps NaN a value, so copying it through `Samples::iter()` would only cost an allocation +/// and a re-encode. Chunked samples go through the iterator, and histograms have no float column. +fn field_column(samples: &Samples) -> DataFusionResult { + match samples.dense() { + Some(DenseSamplesRef::Float(values)) => Ok(Arc::new(values.clone())), + Some(DenseSamplesRef::Histogram(_)) | Some(DenseSamplesRef::Mixed(_)) => { + Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )) + } + None => { + let values = samples + .iter() + .map(|sample| match sample { + SampleRef::Float(value) => Ok(Some(value)), + SampleRef::Missing => Ok(None), + SampleRef::Histogram(_) => Err(DataFusionError::NotImplemented( + "FusedSeriesExec cannot output histogram samples".into(), + )), + }) + .collect::>>()?; + Ok(Arc::new(Float64Array::from(values))) + } + } +} + +/// The bytes one label set holds, for the memory reservation. +/// +/// It counts each label's name and value with the pair that refers to them, which is the same +/// convention [`crate::native::execute`]'s own label-set table keeps. +fn labels_bytes(labels: &LabelSet) -> usize { + labels + .iter() + .map(|(name, value)| size_of::<(Arc, Arc)>() + name.len() + value.len()) + .sum() +} + +/// Renders the label set Prometheus compares the way it prints one, for the duplicate-label-set +/// error. +fn render_labels(labels: &LabelSet) -> String { + let parts = labels + .iter() + .map(|(name, value)| format!("{name}={value:?}")) + .collect::>(); + format!("{{{}}}", parts.join(", ")) +} + +/// The grid positions each final label set already occupies, per field and as merged, sorted runs. +/// +/// Shared by every partition's stream, because a label-set collision is a property of the whole +/// output vector and not of one partition. Two series that carry the same labels after the range +/// function and its maps drop a label land in different partitions, and the runtime's own per-node +/// `Emitted` table never sees both: it belongs to one execution. Checking only within a partition +/// would miss exactly those collisions. +/// +/// The key is [`crate::grid_series::label_set`]'s normalization of the final labels rather than +/// [`datatypes::series::MetricName`] itself: a label with an empty value is no label to +/// Prometheus, so `{host="h", zone=""}` and `{host="h"}` are one label set and have to be compared +/// as one. Encoding the tags as they are would file them apart and let one point hold two values. +/// +/// Occupancy is per field and per position, not per chunk: two series whose label sets became equal +/// while holding values in *different* fields merge cleanly in Prometheus, so only the same field +/// at the same position collides. That is the rule +/// [`crate::grid_series::merge_same_labelset`] applies. +/// +/// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed +/// for the partitions still running, so the entries are only released with the plan itself. The +/// charge has the same lifetime, which is what makes the bytes the map holds - one entry per label +/// set of the output, kept until the plan dies - visible to the memory pool. Charging each stream +/// for what it inserted instead returned those bytes when that stream ended, while the entries the +/// later partitions are checked against stayed: everything the map held after the first partition +/// was unaccounted for. +#[derive(Default)] +struct LabelSets { + /// The grid positions each final label set already occupies, one run list per field of the + /// series schema. Each list is sorted, disjoint and has the runs that touch joined, which is + /// what lets a claim look a position up instead of walking the history. + seen: HashMap>>, + /// What `seen` keeps charged on the task's memory pool. `None` until the first execution, + /// since construction has no `TaskContext` and therefore no pool to register against. + reservation: Option, +} + +impl LabelSets { + /// Attaches the reservation to the pool of the first execution. + /// + /// Later executions keep the first reservation: it already pays for what `seen` holds, and + /// moving it to another pool would drop the bytes other partitions are still growing. + fn attach(&mut self, pool: &Arc) { + if self.reservation.is_none() { + self.reservation = + Some(MemoryConsumer::new("FusedSeriesExec label sets").register(pool)); + } + } + + /// Claims the occupied grid positions of one chunk of a final series for its label set. + /// + /// A collision is an error only when both series hold a sample of *one field* at the same + /// evaluation timestamp, which is Prometheus' range-query rule. Series with disjoint points are + /// legal, and so are series that fill the same points field by field: Prometheus merges both + /// into one output series, while this node emits them as separate batches and defers that merge + /// to matrix assembly, so every sample survives. + fn claim(&mut self, series: &Series) -> DataFusionResult<()> { + let runs = occupied_runs(series)?; + // A chunk holding no value occupies nothing, so there is no collision to remember and no + // entry to write. Such a chunk is not rare: a selection over a series that staleness or the + // range left empty still closes it, and hands over a closing chunk with no rows. Recording + // one entry per such series would grow the table - and its charge, which the plan holds for + // its whole life - for a query that never answers a value, until an empty result exhausted + // the pool. The executor's own `Emitted::occupy` takes the same early return. + if runs.iter().all(Vec::is_empty) { + return Ok(()); + } + let key = label_set(&series.metric); + let Self { seen, reservation } = self; + // Nothing reaches a claim without an `execute` first, and an execution that ran attached the + // reservation. Claiming without one would leave the map growing uncharged, so it is refused + // rather than silently skipped. + let Some(reservation) = reservation.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec label-set state is not attached to a memory pool".to_string(), + )); + }; + match seen.entry(key) { + Entry::Occupied(entry) => { + if let Some(field) = colliding_field(entry.get(), &runs) { + return Err(DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset: {} (field \ + {field})", + render_labels(entry.key()) + ))); + } + let stored = entry.into_mut(); + // An entry recorded from a chunk of fewer fields is padded, so every field of this + // chunk has a slot of its own: a slot is a field index of the schema, so both sides + // count the same field the same way. A chunk of *fewer* fields than the entry + // leaves the remaining slots to the chunks that fill them. + // Stage only the affected run lists, and reserve before committing. A forward + // append whose existing capacity suffices needs no copy of its stored history. + let mut growth = 0usize; + let mut shrink = 0usize; + let mut merged_runs = Vec::with_capacity(runs.len()); + for (index, incoming) in runs.iter().enumerate() { + let Some(old) = stored.get(index) else { + let copy = incoming.clone(); + growth += copy.capacity(); + merged_runs.push(Some(copy)); + continue; + }; + let append_fits = incoming.is_empty() + || (old.last().is_none_or(|last| incoming[0].start >= last.end) + && old.capacity() + >= old.len() + incoming.len() + - usize::from(old.last().is_some_and(|last| { + incoming + .first() + .is_some_and(|first| first.start == last.end) + }))); + if append_fits { + merged_runs.push(None); + } else { + let mut merged = old.clone(); + absorb(&mut merged, incoming); + if merged.capacity() >= old.capacity() { + growth += merged.capacity() - old.capacity(); + } else { + shrink += old.capacity() - merged.capacity(); + } + merged_runs.push(Some(merged)); + } + } + reservation.try_grow(growth * size_of::())?; + for (index, (incoming, merged)) in runs.iter().zip(merged_runs).enumerate() { + if let Some(merged) = merged { + if index < stored.len() { + stored[index] = merged; + } else { + stored.push(merged); + } + } else if index < stored.len() { + absorb(&mut stored[index], incoming); + } + } + reservation.try_shrink(shrink * size_of::())?; + } + Entry::Vacant(entry) => { + // Only the label-set state is this node's own cross-poll state that grows with + // the output, so it is the only thing reserved here. The runs are charged by the + // capacity they were gathered with rather than by their length, so that a later + // claim appending into that slack is charged nothing for it. + let bytes = labels_bytes(entry.key()) + + runs.iter().map(|runs| runs.capacity()).sum::() * size_of::(); + reservation.try_grow(bytes)?; + entry.insert(runs); + } + } + Ok(()) + } +} + +type SharedLabelSets = Arc>; + +/// The grid positions the rows of one chunk hold a value at, one run list per field. +/// +/// Each [`Run`] is a half-open range of absolute grid indices, and the runs of one field are sorted +/// and disjoint. Occupancy is kept per field because it is what a collision is: two series whose +/// label sets became equal merge field by field in Prometheus, so only two values of the same field +/// at one position collide. Marking the whole position as taken as soon as *any* field held a value +/// would refuse the complementary pair that merges cleanly. +fn occupied_runs(series: &Series) -> DataFusionResult>> { + // Rows are addressed on the statement's grid, so a row sits at the absolute position two + // series can collide on. An explicit axis has no such position, and a chunk never carries + // one: this is the refusal `EvaluatedChunk` already makes when the chunk is built. + if series.time.grid().is_none() { + return Err(DataFusionError::Execution( + "FusedSeriesExec requires evaluated series on a time grid".to_string(), + )); + } + let rows = series.time.len(); + let range = series.time.as_grid_range(); + let mut runs: Vec> = vec![Vec::new(); series.fields.len()]; + // A field that is one dense, null-free float run holds a value on every row the chunk carries, + // so it occupies exactly the range the chunk reports and no row of it has to be inspected. This + // is the shape the fused evaluator produces for a range function over ordinary float samples. + // Every other field is walked below, and a field filled here is skipped there. + let mut walked: Vec<(usize, &Samples)> = Vec::new(); + for (index, samples) in series.fields.iter().enumerate() { + match (range, samples.dense()) { + (Some(range), Some(DenseSamplesRef::Float(values))) + if rows > 0 && values.len() == rows && values.null_count() == 0 => + { + runs[index] = vec![Run { + start: range.start(), + end: range.start() + rows, + }]; + } + _ => walked.push((index, samples)), + } + } + if walked.is_empty() { + return Ok(runs); + } + // `positions` walks a chunk that forms one run and one whose rows form several the same way, + // in position order, so the holes between two runs are never read as rows. + let mut cursors: Vec<_> = walked.iter().map(|(_, samples)| samples.cursor()).collect(); + for (position, row) in series.time.positions() { + for ((index, _), cursor) in walked.iter().zip(&mut cursors) { + if !matches!(cursor.at(row), Some(sample) if !matches!(sample, SampleRef::Missing)) { + continue; + } + match runs[*index].last_mut() { + Some(run) if run.end == position => run.end += 1, + _ => runs[*index].push(Run { + start: position, + end: position + 1, + }), + } + } + } + Ok(runs) +} + +/// The field of `runs` that holds a value where `stored` already holds one, if any. +/// +/// A collision is two values of one field at one position, so each field is checked on its own. +/// Each incoming run is looked up rather than the stored history walked: the stored runs are sorted +/// and disjoint, so the only run that can share a position with `run` is the first one that does not +/// end before it, which is the look-up the runtime's own `Emitted::occupy` makes. Walking the +/// history instead costs a sparse series one pass over everything it stored per chunk - a query +/// whose chunks move forward one position at a time, a hole apart, is legal and would cost the +/// square of its chunk count under the lock every partition shares. +fn colliding_field(stored: &[Vec], runs: &[Vec]) -> Option { + for (field, (stored, new)) in stored.iter().zip(runs).enumerate() { + for run in new { + let at = stored.partition_point(|stored| stored.end <= run.start); + if stored.get(at).is_some_and(|stored| stored.overlaps(*run)) { + return Some(field); + } + } + } + None +} + +/// The batches one execution of [`FusedSeriesExec`] yields. +pub struct FusedSeriesStream { + /// The input frames, read as raw series chunks, or `None` once the stream ended in an error. + /// + /// The decoder behind it holds its current batch, the chunks it decoded from it and its own + /// source, so a terminated stream that kept it alive would hold upstream resources until the + /// consumer drops the stream. + input: Option, + /// The execution of this node's tree, or `None` once it has been closed. + executor: Option, + /// The position of the tree's input node, which is where every chunk read enters. + input_path: NodePath, + /// Where each output column comes from. + classification: Arc<[OutputColumn]>, + /// The tag column names, indexed by `OutputColumn::Tag`. + tags: Arc<[String]>, + /// The schema of every batch this stream yields. + output_schema: SchemaRef, + /// The label-set occupancy shared with the other partitions' streams. + /// + /// What it holds is charged to the plan's own reservation, attached by the first execution. + labels: SharedLabelSets, + /// The batches the executor settled since the last poll of the input. + ready: VecDeque, + /// Whether the input ended and everything it settled has been handed out. + done: bool, +} + +impl FusedSeriesStream { + /// Hands one raw chunk over and queues what it settled. + fn push(&mut self, chunk: RawSeriesChunk) -> DataFusionResult<()> { + let Some(executor) = self.executor.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec stream pushed a chunk after its execution closed".to_string(), + )); + }; + // The chunk enters at the tree's input node by position: a position is how the runtime + // addresses a node, and a tree that gains a second input names them the same way. + executor.push_raw_at(&self.input_path, chunk)?; + let mut settled = Vec::new(); + while let Some(chunk) = executor.next_chunk()? { + settled.push(chunk); + } + self.convert(settled) + } + + /// Ends the input, so the tree settles what it was holding. + /// + /// `finish()` would be the same thing plus a drain, but the drain collects: it holds the whole + /// answer in a `Vec` before anything is handed out, so a tree that answers many chunks after its + /// input ends would sit on all of them while the consumer waits for the first. Only the input is + /// ended here, and the answers are pulled one per call below, so each chunk leaves as the tree + /// settles it and what has not been pulled yet stays charged to the node holding it. + fn close_input(&mut self) -> DataFusionResult<()> { + match self.executor.as_mut() { + Some(executor) => executor.end_of_input(), + // Every answer was pulled and the execution closed itself. + None => Ok(()), + } + } + + /// Pulls the next answer the tree settles after its input ended, or `None` when it has none. + /// + /// One chunk per call is what keeps this a stream: collecting them first would hold the whole + /// answer while the consumer waits for the first, and would charge the pool for all of it at + /// once. `None` is also when the execution closes and returns its reservations, so it is + /// dropped here rather than staying alive until the stream is. + fn pull_one(&mut self) -> DataFusionResult> { + let Some(executor) = self.executor.as_mut() else { + return Ok(None); + }; + match executor.next_chunk() { + Ok(Some(chunk)) => Ok(Some(chunk)), + Ok(None) => { + self.executor = None; + Ok(None) + } + Err(error) => Err(error), + } + } + + /// Ends the stream in an error, releasing the execution and the input at once. + /// + /// An ordinary drop of the stream releases all of this as well, so this is not about the last + /// release: it is about a consumer that keeps the terminated stream, which would otherwise + /// hold the evaluator, the decoder's held batch and chunks, and the plan upstream of it. + fn fail(&mut self) { + self.done = true; + self.executor = None; + self.input = None; + // The label-set map and its charge are the plan's, so a stream stopping early gives back + // nothing: the entries another partition is still checking against have to stay charged, and + // are released when the plan itself is dropped. + // A batch queued before the error was settled by a run that has now stopped. + self.ready.clear(); + } + + /// Converts the chunks the executor settled into batches, in the order it produced them. + fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { + for chunk in &chunks { + // The label set is claimed before its batch is queued, so a collision surfaces at the + // chunk that reached it and nothing this plan would refuse is handed out after it. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .claim(&chunk.series)?; + let batch = series_to_record_batch( + &self.classification, + &self.tags, + &self.output_schema, + chunk, + )?; + self.ready.push_back(batch); + } + Ok(()) + } +} + +impl RecordBatchStream for FusedSeriesStream { + fn schema(&self) -> SchemaRef { + self.output_schema.clone() + } +} + +impl Stream for FusedSeriesStream { + type Item = DataFusionResult; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll>> { + let this = &mut *self; + // A batch that is already settled is handed over before anything is read, so a bounded + // poll never delays one that is ready behind more input. + for _ in 0..POLL_BUDGET { + if let Some(batch) = this.ready.pop_front() { + return Poll::Ready(Some(Ok(batch))); + } + if this.done { + return Poll::Ready(None); + } + // An error is the only thing that clears the input, and it ends the stream with it. + let Some(input) = this.input.as_mut() else { + // The input ended, so what the tree answers now comes from what it held. One chunk + // per iteration keeps it a stream, and `None` is both "nothing more" and the point + // where the execution releases its reservations. + match this.pull_one() { + Ok(Some(chunk)) => { + if let Err(error) = this.convert(vec![chunk]) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Ok(None) => { + this.done = true; + continue; + } + Err(error) => { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + continue; + }; + let polled = input.as_mut().poll_next(cx); + match polled { + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => { + // The input ended, so every series it opened has to be closed. The tree settles + // what it was holding, and the answers are pulled below in this same poll + // rather than collected first. + this.input = None; + if let Err(error) = this.close_input() { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + Poll::Ready(Some(Err(error))) => { + // An input error stops the execution: a stopped evaluator must not be reused. + this.fail(); + return Poll::Ready(Some(Err(DataFusionError::External(Box::new(error))))); + } + Poll::Ready(Some(Ok(chunk))) => { + if let Err(error) = this.push(chunk) { + this.fail(); + return Poll::Ready(Some(Err(error))); + } + } + } + } + // The budget ran out on input that settled nothing. Everything read so far is kept, so the + // next poll continues where this one stopped; waking first is what asks for that poll. + cx.waker().wake_by_ref(); + Poll::Pending + } +} diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs new file mode 100644 index 00000000000..20bdebfe2f0 --- /dev/null +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -0,0 +1,2854 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! What this node has to get right: the two `EXPLAIN` lines it shows, the batches it runs end to +//! end, and what it refuses to run. +//! +//! The trees here are built node by node with the same bindings the PromQL compiler produces for +//! `rate(x[5m])`: the compiler itself lives in the `query` crate, which depends on this one, so a +//! test in this crate cannot call it. +//! +//! The input is what a scan hands over: native raw series frames, one batch per chunk here, packed +//! by the same code the scan uses ([`pack_raw_series_chunk`]). A frame carries the series' encoded +//! key, the chunk's finality, its timestamps and its fields, plus the complete labels, so a test +//! says which series a chunk belongs to and whether it closes it instead of relying on where a row +//! boundary happens to fall. +//! +//! Kernel numerics are not the subject: `rate` is checked against a counter that advances one +//! second per second, so the value it has to produce is exactly 1 (or 2) at every grid point and +//! any position or finalization mistake shows up as a wrong, missing or duplicated row. + +use std::collections::VecDeque; +use std::fmt::{self, Display, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray}; +use datafusion::arrow::buffer::{Buffer, ScalarBuffer}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::ToDFSchema; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, collect, displayable, +}; +use datafusion::prelude::SessionContext; +use datatypes::series::frame::{RawLabelMode, pack_raw_series_chunk, raw_series_frame_schema}; +use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; +use datatypes::series::samples::{FieldSpec, Samples, Series, SeriesFields, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; +use datatypes::series::{Label, MetricName, SeriesOrdering}; +use futures::{Stream, StreamExt, TryStreamExt}; +use store_api::region_engine::RawLabelSource; + +use super::{ + FusedSeries, FusedSeriesBinding, LabelSets, OutputColumn, POLL_BUDGET, label_set, + occupied_runs, series_to_record_batch, +}; +use crate::evaluated::EvaluatedChunk; +use crate::native::execute::{NativeExecutionState, NativeExecutor}; +use crate::native::plan::{ + AggregateKind, AggregateStage, BoundAggregate, BoundInput, BoundInstant, BoundWindow, Grouping, + NativeExpr, NativeNode, NativeTree, NativeType, +}; +use crate::native::runs::absorb_rebuilds; +use crate::range_eval::{RangeFunction, RangeWindow}; + +/// The query every tree here is compiled from. +const PROMQL: &str = "rate(x[5m])"; +/// The range the statement's window selects, which is the `5m` of [`PROMQL`]. +const RANGE_MS: i64 = 300_000; +/// The statement's grid: `rate(x[5m])` evaluated every minute from five minutes in. +const GRID_START: i64 = 300_000; +const GRID_STEP: i64 = 60_000; +const GRID_LEN: u64 = 5; +/// One raw sample per minute, from 0 to 9 minutes: enough to fill every window on the grid. +const SAMPLE_STEP: i64 = 60_000; +const SAMPLES: i64 = 10; +/// How many one-sample chunks the quiet series below arrives in. +/// +/// Two bounds hold at once: the series reaches the end of the last grid point's window (540s, so +/// indices `0..=540`), which is what lets it rate at 1 over the whole grid, and it arrives in more +/// batches than one poll's budget, which is what makes a poll that ignored the budget visible. +const QUIET_SAMPLES: i64 = 541; +/// One second between those samples, so the series rates at 1 once it is read to its end. +/// +/// The step is a second and not a minute on purpose: at a minute the series would arrive in ten +/// batches, fewer than [`POLL_BUDGET`], and the budget the test above exists to check would never +/// be reached. +const QUIET_STEP: i64 = 1_000; + +const TIME: &str = "time"; +const VALUE: &str = "value"; +const OTHER: &str = "other"; +const HOST: &str = "host"; +const ZONE: &str = "zone"; +const NAME: &str = "__name__"; +/// A metadata label, which the window drops along with the metric name. +const TYPE: &str = "__type__"; +const METRIC: &str = "http_requests_total"; + +fn grid() -> Arc { + Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ) +} + +/// The raw schema the tree's input is bound to: one float field called [`VALUE`]. +fn raw_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: VALUE.into() }]), + eval_grid: None, + }) +} + +/// A raw schema of two float fields, so one series can fill one of them and another the other. +fn two_field_schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([ + FieldSpec { name: VALUE.into() }, + FieldSpec { name: OTHER.into() }, + ]), + eval_grid: None, + }) +} + +/// The tree `rate(x[5m])` compiles to over `schema`: raw samples folded by one window onto the grid. +fn rate_tree_of(schema: Arc) -> NativeTree { + let grid = grid(); + // One schema shared by the window's binding and its output, which is what `NativeTree`'s + // validation compares by identity. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: schema.fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(schema); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + // The input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: grid, + }, + }, + output: evaluated, + }) + .unwrap() +} + +/// The tree `rate(x[5m])` compiles to over one field. +fn rate_tree() -> NativeTree { + rate_tree_of(raw_schema()) +} + +/// An instant selection feeding a range window. +fn instant_window_tree() -> NativeTree { + let instant = instant_tree().root().clone(); + let grid = instant + .output + .samples() + .expect("an instant selection produces samples") + .eval_grid + .clone() + .expect("an instant selection's output is evaluated on a grid"); + let evaluated = instant.output.clone(); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(instant), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: RANGE_MS, + offset_ms: 0, + anchor_ms: None, + }, + input_grid: Some(Arc::clone(&grid)), + output_grid: grid, + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// Two windows stacked, `rate(last_over_time(x[1s])[2s])`, which no compiler produces. +/// +/// A range function has to apply directly to a selector, and a selector is raw samples, so a tree +/// with a window above a window is only reachable by writing it down by hand - which is a tree +/// [`FusedSeriesExec`](super::FusedSeriesExec) accepts, and therefore one it has to answer correctly. +/// +/// The inner window is the interesting one: `last_over_time` keeps `__name__`, so the pair it can +/// split is not two series whose names it dropped but `{host="h", zone=""}` and `{host="h"}`, one +/// canonical label set and two raw keys. Its fragments are not merged unless the node above it +/// merges them, and that merge happens once per execution: with one execution per partition each +/// partition sees one fragment, so the outer window reads one half of the pair. +fn nested_window_tree() -> NativeTree { + let grid = grid(); + // The inner window's output, which is also the outer window's output: a window reports one value + // per input field on its own grid. + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Window { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundWindow { + function: RangeFunction::LastOverTime, + window: RangeWindow { + range_ms: 1_000, + offset_ms: 0, + anchor_ms: None, + }, + // The inner window's input is raw samples, so it has no grid of its own. + input_grid: None, + output_grid: Arc::clone(&grid), + }, + }, + output: evaluated.clone(), + }), + op: BoundWindow { + function: RangeFunction::Rate, + window: RangeWindow { + range_ms: 2_000, + offset_ms: 0, + anchor_ms: None, + }, + // The outer window reads the inner one's evaluated output, on that output's grid, + // and reports one value per field there like the inner one does. + input_grid: Some(Arc::clone(&grid)), + output_grid: Arc::clone(&grid), + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// The `rate(x[5m])` tree with a `sum by (job)` over it, which the runtime does not execute. +fn aggregate_tree() -> NativeTree { + let rate = rate_tree(); + let window = rate.root().clone(); + let grid = window + .output + .samples() + .expect("a window produces samples") + .eval_grid + .clone() + .expect("a window's output is evaluated on a grid"); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(window), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: Grouping::By(Arc::new(["job".into()])), + stage: AggregateStage::Complete, + }, + }, + output: NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(grid), + })), + }) + .expect("the bound tree is legal") +} + +/// The tree an instant selection compiles to: the last raw sample at or before each grid point. +/// +/// Unlike a window or a map, a selection rewrites no label: the series it answers carries the labels +/// its input did. +fn instant_tree() -> NativeTree { + let grid = grid(); + let evaluated = NativeType::Samples(Arc::new(SeriesSchema { + fields: raw_schema().fields.clone(), + eval_grid: Some(Arc::clone(&grid)), + })); + let raw = NativeType::Samples(raw_schema()); + NativeTree::try_new(NativeNode { + expr: NativeExpr::Instant { + input: Box::new(NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }), + op: BoundInstant { + lookback_ms: RANGE_MS, + offset_ms: 0, + output_grid: grid, + }, + }, + output: evaluated, + }) + .expect("the bound tree is legal") +} + +/// What the tree's input was compiled against, for the input schema and output tags given. +fn binding_of( + tree: &NativeTree, + series_schema: Arc, + tags: Vec, +) -> FusedSeriesBinding { + FusedSeriesBinding { + series_schema, + grid: tree + .root() + .output + .samples() + .expect("this tree produces samples") + .eval_grid + .clone() + .expect("this tree's output is evaluated on a grid"), + metric_name: Some(Arc::from(METRIC)), + time_index: TIME.to_string(), + tags, + } +} + +/// What the tree's input was compiled against, read off the tree itself. +fn binding(tree: &NativeTree) -> FusedSeriesBinding { + binding_of(tree, raw_schema(), vec![HOST.to_string(), NAME.to_string()]) +} + +/// The schema of a frame carrying series of `schema` with their labels explicit. +fn frame_schema_of(schema: &SeriesSchema) -> SchemaRef { + raw_series_frame_schema(schema, RawLabelMode::Explicit).expect("explicit labels are a frame") +} + +/// The schema of a frame carrying these series, which is what the node declares as its input. +fn frame_schema() -> SchemaRef { + frame_schema_of(&raw_schema()) +} + +/// The columns this node produces: the field, the tags and the grid timestamps. +fn output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +fn input_plan(schema: SchemaRef) -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.to_dfschema_ref().unwrap(), + }) +} + +/// A physical input producing `batches` in one partition, with `schema` as its schema. +fn memory_exec_with(schema: SchemaRef, batches: Vec) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[batches], schema, None).unwrap(), + ))) +} + +/// A physical input producing `batches` in one partition, on the frame schema of these tests. +fn memory_exec(batches: Vec) -> Arc { + memory_exec_with(frame_schema(), batches) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own, with +/// `schema` as its schema. +fn memory_exec_partitioned_with( + schema: SchemaRef, + partitions: Vec>, +) -> Arc { + Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&partitions, schema, None) + .expect("every partition holds batches of the given frame schema"), + ))) +} + +/// A physical input producing each of `partitions`' batches in a partition of its own. +fn memory_exec_partitioned(partitions: Vec>) -> Arc { + memory_exec_partitioned_with(frame_schema(), partitions) +} + +/// A task context whose memory pool is `pool`, so what the node charges can be read afterwards. +fn task_context(pool: &Arc) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(pool) as Arc) + .build() + .expect("the runtime environment builds"); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) +} + +/// A physical input whose one partition's stream `make` builds. +fn staged( + schema: SchemaRef, + make: impl Fn() -> SendableRecordBatchStream + Send + Sync + 'static, +) -> Arc { + Arc::new(StagedPlan { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + make: Box::new(make), + }) +} + +/// A physical input, so a test can stage what a memory source cannot: a pending poll, an error +/// after some rows, and a stream whose drop it can see. +struct StagedPlan { + make: Box SendableRecordBatchStream + Send + Sync>, + properties: Arc, +} + +impl std::fmt::Debug for StagedPlan { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("StagedPlan").finish() + } +} + +impl DisplayAs for StagedPlan { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { + Ok(()) + } +} + +impl ExecutionPlan for StagedPlan { + fn name(&self) -> &str { + "StagedPlan" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + // A scan answers the partitions it declares and no others, which is what makes an index + // past this plan's count an error the node above has to pass on. + if partition >= self.properties.partitioning.partition_count() { + return Err(DataFusionError::Execution(format!( + "StagedPlan has no partition {partition}" + ))); + } + Ok((self.make)()) + } +} + +/// The stream a [`StagedPlan`] hands over: `batches`, then an optional error, with a pending poll +/// and a drop a test can observe. +struct StagedStream { + schema: SchemaRef, + batches: VecDeque, + handed: usize, + /// Hand this many batches, then answer [`Poll::Pending`] once. + pending_after: Option, + /// The error to report once the batches are exhausted. + error: Option<&'static str>, + /// Set when this stream is dropped, so a test can see the input being released. + dropped: Option>, + /// Counts the batches this stream has handed over, so a test can see how far one poll read. + handed_out: Option>, +} + +impl StagedStream { + fn new(schema: SchemaRef, batches: Vec) -> Self { + Self { + schema, + batches: batches.into(), + handed: 0, + pending_after: None, + error: None, + dropped: None, + handed_out: None, + } + } + + fn pending_after(mut self, batches: usize) -> Self { + self.pending_after = Some(batches); + self + } + + fn error(mut self, message: &'static str) -> Self { + self.error = Some(message); + self + } + + fn dropped(mut self, flag: &Arc) -> Self { + self.dropped = Some(Arc::clone(flag)); + self + } + + fn counted(mut self, counter: &Arc) -> Self { + self.handed_out = Some(Arc::clone(counter)); + self + } +} + +impl Stream for StagedStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if this.pending_after == Some(this.handed) { + this.pending_after = None; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + match this.batches.pop_front() { + Some(batch) => { + this.handed += 1; + if let Some(counter) = &this.handed_out { + counter.fetch_add(1, Ordering::SeqCst); + } + Poll::Ready(Some(Ok(batch))) + } + None => match this.error.take() { + Some(message) => { + Poll::Ready(Some(Err(DataFusionError::Execution(message.to_string())))) + } + None => Poll::Ready(None), + }, + } + } +} + +impl RecordBatchStream for StagedStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Drop for StagedStream { + fn drop(&mut self) { + if let Some(dropped) = &self.dropped { + dropped.store(true, Ordering::SeqCst); + } + } +} + +fn fused(tree: NativeTree) -> FusedSeries { + let binding = binding(&tree); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The physical node of `node` over `input`, reading the complete labels every frame carries. +fn plan_of( + node: FusedSeries, + input: Arc, +) -> DataFusionResult> { + node.to_execution_plan(input, RawLabelSource::Explicit, SeriesOrdering::Contiguous) +} + +/// The physical node of `tree` over `input`. +fn plan( + tree: NativeTree, + input: Arc, +) -> DataFusionResult> { + plan_of(fused(tree), input) +} + +/// The timestamps of `rows`. +fn times_of(rows: &[(i64, f64)]) -> Vec { + rows.iter().map(|(time, _)| *time).collect() +} + +/// The values of `rows`, one sample per timestamp. +fn values_of(rows: &[(i64, f64)]) -> Vec> { + rows.iter().map(|(_, value)| Some(*value)).collect() +} + +/// One missing sample per timestamp of `rows`, which is a field the series holds no value in. +fn no_values(rows: &[(i64, f64)]) -> Vec> { + rows.iter().map(|_| None).collect() +} + +/// One chunk of `metric`, keyed `key`, holding `times` and one sample per field of `schema`. +fn chunk_of( + schema: &SeriesSchema, + key: &str, + metric: MetricName, + times: &[i64], + values: Vec>>, + last: bool, +) -> RawSeriesChunk { + assert_eq!( + schema.fields.len(), + values.len(), + "one value list per field of the schema" + ); + assert!( + values.iter().all(|values| values.len() == times.len()), + "one sample per timestamp per field" + ); + let fields = values + .into_iter() + .map(|values| Samples::from_float(Float64Array::from(values))) + .collect::>(); + RawSeriesChunk::try_new( + schema, + Buffer::from_slice_ref(key.as_bytes()), + Some(metric), + RawTimestamps::try_new(ScalarBuffer::from(times.to_vec())).expect("ascending timestamps"), + SeriesFields::try_from_vec(fields).expect("one sample per timestamp per field"), + last, + ) + .expect("one series' samples are a legal raw chunk") +} + +/// One chunk of the series `{__name__=METRIC, host=}`: `rows` as its samples, and the chunk +/// that closes the series when `last`. The host names the series, so it doubles as the key. +fn series_chunk(host: &str, rows: &[(i64, f64)], last: bool) -> RawSeriesChunk { + chunk_of( + &raw_schema(), + host, + metric(&[(HOST, host)], Some(METRIC)), + ×_of(rows), + vec![values_of(rows)], + last, + ) +} + +/// The metric these tests read, with `labels` on it and the name given. +fn metric(labels: &[(&str, &str)], name: Option<&str>) -> MetricName { + MetricName::try_new( + name.map(Arc::from), + labels + .iter() + .map(|(name, value)| Label { + name: Arc::from(*name), + value: Arc::from(*value), + }) + .collect(), + ) + .expect("the labels of these tests are legal") +} + +/// One frame batch carrying `chunk` of `schema`, packed by the scan's own encoder. +fn frame_of(schema: &SeriesSchema, chunk: &RawSeriesChunk) -> RecordBatch { + let mut batches = pack_raw_series_chunk(schema, RawLabelMode::Explicit, chunk) + .expect("a chunk of this schema packs into a frame"); + assert_eq!(1, batches.len(), "one small chunk is one frame batch"); + batches.remove(0) +} + +/// One frame batch carrying `chunk` of the one-field schema of these tests. +fn frame(chunk: &RawSeriesChunk) -> RecordBatch { + frame_of(&raw_schema(), chunk) +} + +/// The frames carrying `chunks` of `schema`, one batch each, so a test can say how much input one +/// poll reads. +fn frames_of(schema: &SeriesSchema, chunks: &[RawSeriesChunk]) -> Vec { + chunks.iter().map(|chunk| frame_of(schema, chunk)).collect() +} + +/// The frames carrying `chunks` of the one-field schema of these tests, one batch each. +fn frames(chunks: &[RawSeriesChunk]) -> Vec { + frames_of(&raw_schema(), chunks) +} + +/// One series' samples: one per minute from 0, `scale` seconds per second. +fn series_samples(scale: f64) -> Vec<(i64, f64)> { + (0..SAMPLES) + .map(|index| { + let time = index * SAMPLE_STEP; + (time, scale * time as f64 / 1000.0) + }) + .collect() +} + +/// Two series, each split over two chunks. A series' chunks stay adjacent, and the second of each +/// pair closes it, which is the input contract: a series' chunks arrive together. +fn two_series_frames() -> Vec { + let first = series_samples(1.0); + let second = series_samples(2.0); + frames(&[ + series_chunk("web-1", &first[..3], false), + series_chunk("web-1", &first[3..], true), + series_chunk("web-2", &second[..3], false), + series_chunk("web-2", &second[3..], true), + ]) +} + +/// One series split over three chunks, so its chunks have to be stitched. +fn one_series_frames() -> Vec { + let rows = series_samples(1.0); + frames(&[ + series_chunk("web-1", &rows[..4], false), + series_chunk("web-1", &rows[4..7], false), + series_chunk("web-1", &rows[7..], true), + ]) +} + +/// The rows of `batches`: timestamp, value, host and metric name. +type Row = (i64, Option, Option, Option); + +fn rows(batches: &[RecordBatch]) -> Vec { + let mut rows = Vec::new(); + for batch in batches { + let time = batch + .column_by_name(TIME) + .expect("every batch has a time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let value = batch + .column_by_name(VALUE) + .expect("every batch has the field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + let host = batch + .column_by_name(HOST) + .expect("every batch has the tag column") + .as_any() + .downcast_ref::() + .expect("the tag column is Utf8"); + let name = batch + .column_by_name(NAME) + .expect("every batch has the metric name column") + .as_any() + .downcast_ref::() + .expect("the metric name column is Utf8"); + for row in 0..batch.num_rows() { + rows.push(( + time.value(row), + value.is_valid(row).then(|| value.value(row)), + host.is_valid(row).then(|| host.value(row).to_string()), + name.is_valid(row).then(|| name.value(row).to_string()), + )); + } + } + rows +} + +/// The output of [`zone_node`]: the field, `host`, `zone`, the dropped name and the grid times. +fn zone_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(ZONE, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// `rate(x[5m])` over frames that carry `zone` as well as `host`, with both in the output. +fn zone_node() -> FusedSeries { + let tree = rate_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The instant-selection node over frames that carry `zone` as well as `host`. +/// +/// Its output labels are the input's, so an empty-valued label reaches this node's own state as the +/// frame wrote it. +fn instant_node() -> FusedSeries { + let tree = instant_tree(); + let binding = binding_of( + &tree, + raw_schema(), + vec![HOST.to_string(), ZONE.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + zone_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The two chunks of the collision: one series once with an empty `zone` and once with no `zone`. +/// +/// They are one identity to Prometheus — an empty label value is no label — but two different raw +/// series: the scan encodes the empty value and the absent one differently, so they reach the +/// runtime as two series that close on the same label set. +fn zone_chunks() -> Vec { + let rows = series_samples(1.0); + vec![ + chunk_of( + &raw_schema(), + "zone-empty", + metric(&[(HOST, "web-1"), (ZONE, "")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + chunk_of( + &raw_schema(), + "zone-absent", + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&rows), + vec![values_of(&rows)], + true, + ), + ] +} + +/// The canonical label set these tests collide on, as the runtime renders it in its refusal. +const DUPLICATE_LABEL_SET: &str = + "the node at root produces two series under the label set [(\"host\", \"web-1\")]"; + +/// Two raw series that are already one Prometheus identity are refused by the runtime that would +/// have to answer both. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one identity to Prometheus: an empty label +/// value is no label, so the two chunks are the same series and `rate` creates no collision here. +/// What the two *raw* keys differ in is the empty `zone` against the absent one, which the scan +/// encodes differently and the decoder therefore opens two input series for. `rate` drops the +/// metric name and the empty value, so the second raw series arrives on the canonical label set the +/// first one already occupies; the refusal names that label set rather than answering it, because +/// answering would report one point twice. +#[tokio::test] +async fn differing_raw_keys_that_close_on_one_label_set_are_refused() { + let exec = plan_of(zone_node(), memory_exec(frames(&zone_chunks()))) + .expect("the runtime executes a rate tree"); + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second raw key cannot be answered under the first one's label set"); + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// A runtime error after buffered rows is reported, and the terminated stream releases its input. +/// +/// The decoder holds the batch it is reading, the chunks it decoded from it and its own source, so a +/// stream that kept its input after the error would hold all three until the consumer drops it. The +/// collision above is the error: it happens on the second series, after the first one's rows have +/// been handed out. +#[tokio::test] +async fn a_runtime_error_after_rows_releases_the_input() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = frame_schema(); + let batches = frames(&zone_chunks()); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).dropped(&dropped)) + }); + let mut stream = plan_of(zone_node(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!( + "the collision was answered instead of refused: {:?}", + rows(&answered) + ), + } + }; + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); + // The first series' rows were handed out before the error, and the input is gone although the + // terminated stream is still alive here. + assert!( + !answered.is_empty(), + "the first series is answered before the error" + ); + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// Two series with *different* names collapse once `rate` drops the name. +/// +/// This is the collision the canonical label set exists for, at the layer that sees it: +/// `http_requests_total{host="web-1"}` and `http_errors_total{host="web-1"}` reach the window as two +/// identities and leave it as `{host="web-1"}`. [`FusedSeries`] reads one metric per scan, so this +/// collision cannot be staged through it; the runtime is driven directly, with the names on the raw +/// chunks. +#[test] +fn two_differently_named_series_collide_once_rate_drops_the_name() { + let tree = rate_tree(); + let mut execution = NativeExecutor::new( + NativeExecutionState::try_new( + &tree, + Arc::new(UnboundedMemoryPool::default()), + &ExecutionPlanMetricsSet::new(), + 0, + ) + .expect("the runtime executes a rate tree"), + ); + let samples = series_samples(1.0); + execution + .push_raw(chunk_of( + &raw_schema(), + METRIC, + metric(&[(HOST, "web-1")], Some(METRIC)), + ×_of(&samples), + vec![values_of(&samples)], + true, + )) + .expect("the first series is a legal raw chunk"); + assert!( + execution + .next_chunk() + .expect("the first series runs") + .is_some(), + "the first series is answered" + ); + execution + .push_raw(chunk_of( + &raw_schema(), + "http_errors_total", + metric(&[(HOST, "web-1")], Some("http_errors_total")), + ×_of(&samples), + vec![values_of(&samples)], + true, + )) + .expect("the second series is a legal raw chunk"); + let error = execution + .next_chunk() + .expect_err("the second series puts a value where the first one already holds one"); + assert!( + error.to_string().contains(DUPLICATE_LABEL_SET), + "refused for another reason: {error}" + ); +} + +/// The grid timestamps of the statement, in order. +fn grid_timestamps() -> Vec { + (0..GRID_LEN) + .map(|index| GRID_START + index as i64 * GRID_STEP) + .collect() +} + +/// A waker a test can ask whether it was woken. +/// +/// A poll that stops on its own budget has to wake itself, and the default waker a `poll!` uses +/// swallows that, so the tests that check it poll with this one. +#[derive(Default)] +struct CountingWake(AtomicBool); + +impl futures::task::ArcWake for CountingWake { + fn wake_by_ref(woken: &Arc) { + woken.0.store(true, Ordering::SeqCst); + } +} + +impl CountingWake { + fn woken(&self) -> bool { + self.0.load(Ordering::SeqCst) + } +} + +/// One series as many chunks that settle nothing: one sample per chunk, all before the grid starts. +/// +/// The counter advances one second per second, so `rate` answers 1 at every grid point once the +/// series is read to its end; until a sample reaches the grid's first point, a push decides +/// nothing and settles no chunk. +fn quiet_series_frames() -> Vec { + frames( + &(0..QUIET_SAMPLES) + .map(|index| { + series_chunk( + "web-1", + &[(index * QUIET_STEP, index as f64)], + index == QUIET_SAMPLES - 1, + ) + }) + .collect::>(), + ) +} + +/// Asserts that `rows` are `rates`' series, each over the whole grid in order. +fn assert_series_rates(rows: &[Row], rates: &[(f64, &str)]) { + let expected = rates + .iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, *rate, *host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } +} + +fn count(metrics: &MetricsSet, name: &str) -> usize { + match metrics.sum_by_name(name) { + Some(MetricValue::Count { count, .. }) => count.value(), + other => panic!("no count named {name}: {other:?}"), + } +} + +/// Prints `node` the way the logical `EXPLAIN` prints it. +struct Explain<'a>(&'a FusedSeries); + +impl Display for Explain<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.0.fmt_for_explain(f) + } +} + +/// The logical node prints the query a reader wrote, not the tree it compiled to. +#[test] +fn fmt_for_explain_prints_the_promql() { + assert_eq!( + "FusedSeries: rate(x[5m])", + Explain(&fused(rate_tree())).to_string() + ); +} + +/// The physical node prints the tree it was built from, rendered by the tree's own display. +#[test] +fn fmt_as_prints_the_native_tree() { + let exec = + plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes a rate tree"); + let text = displayable(exec.as_ref()).indent(false).to_string(); + + assert!( + text.starts_with("FusedSeriesExec: rate(x[5m])"), + "the node line names the query first: {text}" + ); + // A distinctive line of the compiled tree: the window's bound range, offset, anchor and grid. + assert!( + text.contains("root Window rate [300000ms] offset 0 anchor - on grid 300000/60000x5"), + "the node line carries the tree it runs: {text}" + ); + assert!( + text.contains("input Input raw samples"), + "the node line carries the tree's source: {text}" + ); +} + +/// `rate(x[5m])` runs end to end over two raw series: values, tags and the dropped metric name. +#[tokio::test] +async fn runs_rate_over_two_series() { + let exec = plan(rate_tree(), memory_exec(two_series_frames())) + .expect("the runtime executes a rate tree"); + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // A counter advancing one second per second rates at 1, and at 2 for the series that advances + // two: the value at each grid point says the right window was read for the right series. + let expected = [(1.0, "web-1"), (2.0, "web-2")] + .into_iter() + .flat_map(|(rate, host)| { + grid_timestamps() + .into_iter() + .map(move |time| (time, rate, host)) + }) + .collect::>(); + assert_eq!(expected.len(), rows.len(), "{rows:?}"); + for (row, (expected_time, expected_rate, expected_host)) in rows.iter().zip(expected) { + let (time, value, host, _) = row; + assert_eq!(&expected_time, time); + assert_eq!(Some(expected_host), host.as_deref()); + let value = value.expect("rate has a value at every grid point"); + assert!( + (expected_rate - value).abs() < 1e-9, + "{expected_host} at {expected_time}: {value} is not {expected_rate}" + ); + } + // `rate` drops `__name__`, so the tag column that reports it is empty for every row. + assert!( + rows.iter().all(|(_, _, _, name)| name.is_none()), + "rate drops the metric name: {rows:?}" + ); +} + +/// A series that spans several input chunks keeps its absolute grid positions and closes once. +#[tokio::test] +async fn a_series_spanning_chunks_keeps_absolute_positions_and_closes_once() { + let exec = plan(rate_tree(), memory_exec(one_series_frames())) + .expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("the node runs"); + let rows = rows(&batches); + + // Every grid position exactly once, in order: a chunk's range is never read as a row number. + let times = rows.iter().map(|row| row.0).collect::>(); + assert_eq!(grid_timestamps(), times); + assert!( + rows.iter().all(|(_, value, ..)| value.is_some()), + "every grid point of this series has a value: {rows:?}" + ); + + // The source closed the series on exactly one chunk, and the window saw exactly that one + // closing: one `is_last`, not one per input batch. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(1, count(&metrics, "native.root.input.completed_series")); + assert_eq!(1, count(&metrics, "native.root.completed_series")); +} + +/// A tree the runtime does not execute is refused when the execution plan is built. +#[tokio::test] +async fn a_tree_the_runtime_cannot_run_is_rejected_when_the_exec_is_built() { + let node = fused(aggregate_tree()); + let error = plan_of(node, memory_exec(Vec::new())) + .expect_err("the runtime does not execute aggregate nodes"); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "the runtime refuses the tree it cannot run: {error}" + ); + assert!(error.to_string().contains("Aggregate"), "{error}"); +} + +/// A tree evaluated on another grid than the node declares is refused before it runs. +#[test] +fn a_tree_bound_to_another_grid_is_rejected() { + let tree = rate_tree(); + let other_grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(0, GRID_STEP, GRID_LEN).unwrap()).unwrap()); + let binding = FusedSeriesBinding { + grid: other_grid, + ..binding(&tree) + }; + let error = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect_err("the tree is not evaluated on the declared grid"); + assert!(error.to_string().contains("grid"), "{error}"); +} + +/// A grid allocated separately from the tree's, with the same spec, is the tree's grid. +/// +/// The comparison is by spec, which is the rule [`NativeTree`]'s own validation and the kernels +/// apply, and a chunk's timestamps come from its absolute grid positions rather than from the +/// allocation, so nothing downstream needs the `Arc` to be shared. This pairs with +/// [`a_tree_bound_to_another_grid_is_rejected`]: a different spec is still refused. +#[test] +fn a_separately_allocated_grid_with_an_equal_spec_is_accepted() { + let tree = rate_tree(); + let equal = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ); + assert!( + !Arc::ptr_eq(&equal, &binding(&tree).grid), + "the two grids are separate allocations" + ); + let binding = FusedSeriesBinding { + grid: equal, + ..binding(&tree) + }; + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema()), + output_schema().to_dfschema_ref().unwrap(), + ) + .expect("an equal grid spec is the same grid"); +} + +/// An input that ends without a series still runs the tree, and yields no rows. +#[tokio::test] +async fn an_empty_input_runs_the_tree_and_yields_no_rows() { + let exec = + plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes a rate tree"); + let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx()) + .await + .expect("a tree with no series in it still closes"); + assert!(rows(&batches).is_empty(), "{:?}", rows(&batches)); + // The run closed on no series, so nothing was completed and nothing was decided. + let metrics = exec + .metrics() + .expect("the node reports its runtime's metrics"); + assert_eq!(0, count(&metrics, "native.root.completed_series")); + assert_eq!(0, count(&metrics, "native.root.input.completed_series")); +} + +/// An upstream error after some rows is reported, and the terminated stream releases that input. +#[tokio::test] +async fn an_upstream_error_after_rows_is_reported() { + let dropped = Arc::new(AtomicBool::new(false)); + let schema = frame_schema(); + let batches = one_series_frames(); + let source = staged(schema.clone(), { + let dropped = Arc::clone(&dropped); + move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + .error("the source failed") + .dropped(&dropped), + ) + } + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The error is reported rather than swallowed by the rows the input already handed over. + let mut answered = Vec::new(); + let error = loop { + match stream.next().await { + Some(Ok(batch)) => answered.push(batch), + Some(Err(error)) => break error, + None => panic!("the source failure was swallowed: {:?}", rows(&answered)), + } + }; + assert!( + error.to_string().contains("the source failed"), + "reported another error: {error}" + ); + // Still alive here: the release is the stream's, not the drop of a temporary. + assert!( + dropped.load(Ordering::SeqCst), + "the terminated stream still holds its input" + ); +} + +/// A pending poll loses nothing when the next-item future is dropped and the stream polled again. +#[tokio::test] +async fn a_pending_boundary_loses_no_rows() { + let expected = rows( + &collect( + plan(rate_tree(), memory_exec(one_series_frames())) + .expect("the runtime executes a rate tree"), + SessionContext::default().task_ctx(), + ) + .await + .expect("the node runs"), + ); + assert!(!expected.is_empty(), "the fixture decides some grid points"); + + let schema = frame_schema(); + let batches = one_series_frames(); + let source = staged(schema.clone(), move || { + Box::pin( + StagedStream::new(schema.clone(), batches.clone()) + // One batch is in the decoder's hands when the stream reports pending. + .pending_after(1), + ) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The future is dropped at the end of this block, as a consumer that gives up waiting does. + { + let next = stream.next(); + futures::pin_mut!(next); + assert!( + futures::poll!(next.as_mut()).is_pending(), + "the input's pending poll reaches the consumer" + ); + } + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the pending poll left it"); + assert_eq!(expected, rows(&batches)); +} + +/// A frame whose labels are encoded in its key is refused when the plan is built. +/// +/// The tags this node declares are read off each chunk's labels, and a key-mode frame carries none: +/// decoding it with a binding that promises complete labels would produce chunks that name no +/// series. The frame's own declared layout says which of the two it is, so the mismatch is caught +/// before the first batch, not when the missing labels reach the output. +#[test] +fn a_frame_whose_labels_are_encoded_is_rejected_when_the_plan_is_built() { + let keyed = + raw_series_frame_schema(&raw_schema(), RawLabelMode::Key).expect("key labels are a frame"); + let error = plan(rate_tree(), memory_exec_with(keyed, Vec::new())) + .expect_err("the frames carry no labels to read tags from"); + assert!( + error.to_string().contains("raw frame"), + "refused for another reason: {error}" + ); +} + +/// A frame that carries an evaluation grid is refused when the plan is built. +/// +/// A frame holds raw series: its rows are samples, not grid positions, and the grid every output +/// row is addressed on is the tree's own. A schema that declares a grid is one this node would +/// read positions out of, so it is refused rather than decoded. The source below is one partition, +/// which is the shape the node accepts, so what refuses the plan here is the frame's own +/// declaration and not the node's partition count: the refusal has to name the grid. +#[test] +fn a_frame_with_an_evaluation_grid_is_rejected_when_the_plan_is_built() { + let mut metadata = frame_schema().metadata().clone(); + metadata.insert("greptime.series.grid".to_string(), "v1".to_string()); + let gridded = Arc::new(Schema::new_with_metadata( + frame_schema().fields().clone(), + metadata, + )); + // One empty partition: every check but the frame's passes. + let source = memory_exec_with(gridded, Vec::new()); + assert_eq!( + 1, + source.properties().output_partitioning().partition_count(), + "the node only looks at the frame when its input has a partition" + ); + let error = plan(rate_tree(), source).expect_err("the frame declares a grid"); + let message = error_chain(&error); + assert!( + message.contains("raw series schema cannot have an evaluation grid"), + "refused for another reason: {message}" + ); + assert!( + !message.contains("no partitions"), + "the partition count is not what refuses this plan: {message}" + ); +} + +/// A frame of another series schema is refused when the plan is built. +/// +/// The tree's input node was compiled against one schema, and a frame for a different one names +/// fields the tree does not read and omits ones it does: decoding it would hand the tree a column +/// it never bound. Nothing about the two schemas says which one was meant, so the frame is refused +/// instead of half-read. +#[test] +fn a_frame_of_another_series_schema_is_rejected_when_the_plan_is_built() { + let other = frame_schema_of(&two_field_schema()); + let error = plan(rate_tree(), memory_exec_with(other, Vec::new())) + .expect_err("the tree's input reads one field"); + assert!( + error.to_string().contains("raw frame"), + "refused for another reason: {error}" + ); +} + +/// An input with no partitions is refused when the plan is built, and no other count is. +/// +/// The node runs one evaluator per input partition and adds none of its own, so an input with no +/// partition is one that could answer nothing: the plan is refused rather than built as a node that +/// would report an empty result for data it never read. Any non-zero count is the input's to +/// choose: the node takes it and answers each partition's own series. +#[test] +fn an_input_with_no_partitions_is_rejected_when_the_plan_is_built() { + let schema = frame_schema(); + let empty = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[], Arc::clone(&schema), None) + .expect("no partitions is a legal source"), + ))); + let error = plan(rate_tree(), empty).expect_err("there is no partition to run"); + assert!( + error.to_string().contains("input has no partitions"), + "refused for another reason: {error}" + ); + // The only count refused is zero: two partitions build and answer, which is what the + // multi-partition tests below run. + let two = plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("a two-partition input is not refused"); + assert_eq!( + 2, + two.properties().output_partitioning().partition_count(), + "the node takes the input's count" + ); +} + +/// Two windows stacked are refused over more than one partition, because the inner one's fragments +/// need not have been merged. +/// +/// The counterexample is legal PromQL once nesting is expressible: the inner `last_over_time` runs +/// over two raw series that are one canonical label set, `{host="h", zone=""}` holding `(1s, 1)` and +/// `{host="h"}` holding `(2s, 2)`, and the outer `rate` needs both in one window. In a single +/// partition the outer window merges the two fragments and answers `(2s, 1)`; in two partitions each +/// executor is handed one fragment and has one sample to rate, so the outer window answers nothing. +/// That is an empty result rather than an error, which is why no check further down can see it. +/// +/// `last_over_time` keeps `__name__`, so a check that looked only for a dropped metric name accepts +/// this tree. What makes it unsafe is that a window's output can hold two series under one label set +/// whenever it did not merge its input, which a bare tree cannot say. +/// +/// One partition stays accepted: there is nothing to split there, which is what makes the refusal +/// about the split rather than about the shape of the tree. +#[test] +fn a_nested_window_over_two_partitions_is_refused() { + let error = plan( + nested_window_tree(), + memory_exec_partitioned(vec![vec![], vec![]]), + ) + .expect_err("each of two partitions would merge only the fragment it was given"); + let message = error.to_string(); + assert!( + message.contains("`root.input`"), + "the refusal must name the inner window: {message}" + ); + assert!( + message.contains("`root`"), + "the refusal must name the window above it: {message}" + ); + assert!( + message.contains("2 partitions"), + "the refusal must name the partition count: {message}" + ); + + // The same tree over one partition is accepted, because one execution sees every fragment and + // merges them into one series before the window above it evaluates. + plan(nested_window_tree(), memory_exec(Vec::new())) + .expect("a single partition cannot split a logical series"); +} + +/// An instant selection can leave distinct raw labels on one canonical label set for a window. +#[test] +fn an_instant_then_window_over_two_partitions_is_refused() { + let error = plan( + instant_window_tree(), + memory_exec_partitioned(vec![vec![], vec![]]), + ) + .expect_err("each partition could hand the window only one fragment"); + let message = error.to_string(); + assert!( + message.contains("`root.input`"), + "the refusal must name the instant selection: {message}" + ); + assert!( + message.contains("`root`"), + "the refusal must name the consuming window: {message}" + ); + assert!( + message.contains("2 partitions"), + "the refusal must name the partition count: {message}" + ); +} + +/// A single partition cannot split a logical series, even when an instant feeds a window. +#[test] +fn an_instant_then_window_over_one_partition_is_accepted() { + plan(instant_window_tree(), memory_exec(Vec::new())).expect("one partition sees all fragments"); +} + +/// A window directly over raw input remains valid across partitions without a preceding rewriter. +#[test] +fn a_window_over_raw_input_on_two_partitions_is_accepted() { + plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("ordinary input-to-window trees do not split a logical series"); +} + +/// One poll reads at most [`POLL_BUDGET`] chunks of input, and stopping there loses nothing. +/// +/// `rate(x[5m])` decides nothing until a sample reaches the grid's first point, so a series can +/// arrive as many chunks that settle no batch at all: a poll that kept reading them would hold the +/// task's thread until the series ended, which is what stops other tasks on that thread and delays +/// the cancellation a consumer asked for. The first series here is [`QUIET_SAMPLES`] such chunks; +/// the second one covers the grid, so a poll that resumes properly still answers both. +#[tokio::test] +async fn one_poll_reads_at_most_the_budget_and_loses_nothing() { + let mut batches = quiet_series_frames(); + batches.push(frame(&series_chunk("web-2", &series_samples(2.0), true))); + let total = batches.len(); + assert!( + total > POLL_BUDGET, + "the fixture has to outlast one poll: {total} batches" + ); + + let pulled = Arc::new(AtomicUsize::new(0)); + let woken = Arc::new(CountingWake::default()); + let schema = frame_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + // The first poll: the quiet series settles nothing, so the budget is what ends this poll. + let waker = futures::task::waker(Arc::clone(&woken)); + let mut context = Context::from_waker(&waker); + assert!( + stream.as_mut().poll_next(&mut context).is_pending(), + "chunks that settle nothing cannot answer a poll" + ); + let read = pulled.load(Ordering::SeqCst); + assert!(read > 0, "the poll did read its input"); + // One step polls the decoder once, and the decoder reads at most one frame per poll: every + // frame here carries one chunk. The bound is therefore one frame per step, and the budget is + // what it is measured in. + assert!( + read <= POLL_BUDGET + 1, + "one poll read {read} of {total} batches, past its budget of {POLL_BUDGET} steps" + ); + assert!( + read < total, + "one poll read the whole input ({read} batches): the budget did not stop it" + ); + assert!( + woken.woken(), + "a poll that stopped on its budget has to wake itself" + ); + + // Nothing the budget stopped is lost: the stream resumes where the poll left it. + let batches = stream + .try_collect::>() + .await + .expect("the stream resumes where the bounded poll left it"); + assert_series_rates(&rows(&batches), &[(1.0, "web-1"), (2.0, "web-2")]); +} + +/// A batch that is ready is handed over before any batch behind it is read. +/// +/// Each series below is one chunk, and a chunk that carries `is_last` closes its series, so the +/// first answer needs only the first chunk. A poll that read on would answer the same batch only +/// after the whole input, which is exactly what the budget above exists to prevent. +#[tokio::test] +async fn a_ready_batch_is_returned_before_later_input_is_read() { + let batches = (0..20) + .map(|index| { + frame(&series_chunk( + &format!("web-{index}"), + &series_samples(1.0), + true, + )) + }) + .collect::>(); + let total = batches.len(); + let pulled = Arc::new(AtomicUsize::new(0)); + let schema = frame_schema(); + let source = staged(schema.clone(), { + let pulled = Arc::clone(&pulled); + move || Box::pin(StagedStream::new(schema.clone(), batches.clone()).counted(&pulled)) + }); + let mut stream = plan(rate_tree(), source) + .expect("the runtime executes a rate tree") + .execute(0, SessionContext::default().task_ctx()) + .expect("the node executes"); + + let first = stream + .next() + .await + .expect("the first series is answered") + .expect("the first series is answered without an error"); + let answered = rows(&[first]); + assert_series_rates(&answered, &[(1.0, "web-0")]); + // One batch: the series' own chunk closes it, so nothing behind it has to be read. + assert_eq!( + 1, + pulled.load(Ordering::SeqCst), + "the answer waited for input behind the series it reports" + ); + assert!(total > 1, "the fixture has to outlast the first answer"); +} + +/// A chunk whose rows form several runs is converted row by row from the rows it carries. +/// +/// A sparse result is one chunk holding two runs with a hole between them +/// ([`EvaluatedChunk::range`] answers `None` for it). Reading the hole as rows would invent grid +/// points this series holds no value at, and reading positions as compact rows would shift every +/// row after the hole onto the wrong timestamp: the batch has to hold exactly the carried rows, at +/// their absolute grid times, with nulls nowhere else. +#[test] +fn a_chunk_of_several_runs_converts_its_own_rows_only() { + let grid = grid(); + let axis = TimeAxis::from_grid_spans( + Arc::clone(&grid), + // Grid positions 0..2 and 4..5: position 2 and 3 hold no row at all. + vec![GridSpan::new(0, 2), GridSpan::new(4, 1)], + ) + .expect("the spans are increasing and inside the grid"); + let chunk = EvaluatedChunk::try_new( + Series { + metric: MetricName::try_new(Some(Arc::from(METRIC)), Vec::new()) + .expect("one metric name is a legal identity"), + time: axis, + fields: SeriesFields::try_from_vec(vec![Samples::from_float(Float64Array::from( + // One value per carried row, so the second run's value sits at grid position 4. + vec![7.0, 8.0, 11.0], + ))]) + .expect("one field is a legal field list"), + }, + true, + ) + .expect("a chunk on a grid"); + + // The columns of `output_schema`, in its order: the field, the two tags, then the time. + let batch = series_to_record_batch( + &[ + OutputColumn::Field(0), + OutputColumn::Tag(0), + OutputColumn::Tag(1), + OutputColumn::Time, + ], + &[HOST.to_string(), NAME.to_string()], + &output_schema(), + &chunk, + ) + .expect("the chunk converts"); + + let stamps = batch + .column_by_name(TIME) + .expect("the batch has its time column") + .as_any() + .downcast_ref::() + .expect("the time column is a millisecond timestamp"); + let values = batch + .column_by_name(VALUE) + .expect("the batch has its field column") + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + assert_eq!(3, batch.num_rows(), "only the carried rows are rows"); + let grid_stamps = grid.timestamps().expect("the grid timestamps build"); + assert_eq!( + vec![grid_stamps[0], grid_stamps[1], grid_stamps[4]], + stamps.values()[..3].to_vec(), + "the carried rows keep their absolute grid positions" + ); + assert_eq!(vec![7.0, 8.0, 11.0], values.values()[..3].to_vec()); +} + +/// Every message in `error`'s chain, so an error a layer above wrapped is still read. +fn error_chain(error: &DataFusionError) -> String { + let mut message = error.to_string(); + let mut source = std::error::Error::source(error); + while let Some(cause) = source { + message.push_str(": "); + message.push_str(&cause.to_string()); + source = cause.source(); + } + message +} + +/// Two partitions, one series each, and both series leave the window on the same label set. +/// +/// `rate` drops `__name__`, so `http_requests_total{host="web-1"}` reads as `{host="web-1"}` in +/// the output. The key column keeps each of the two whole inside its own partition, so a check that +/// lived inside one execution would see one series each and report nothing: only state shared by +/// the partitions sees that the output vector holds the label set twice. +#[tokio::test] +async fn a_duplicate_label_set_across_partitions_is_refused() { + let series = frame(&series_chunk("web-1", &series_samples(1.0), true)); + let exec = plan( + rate_tree(), + memory_exec_partitioned(vec![vec![series.clone()], vec![series]]), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("the second partition lands on the label set the first one claimed"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// The same label set on points that do not overlap is legal, and every point survives. +/// +/// A rate needs two samples in its window, so the first series only has a value at the first grid +/// point and the second only at the last: the two share a final label set but no evaluation +/// timestamp. Prometheus merges such series into one; this node emits both batches and leaves the +/// merge to matrix assembly, so the check has to allow what it would otherwise refuse. +#[tokio::test] +async fn a_duplicate_label_set_on_disjoint_grid_points_is_emitted() { + let first = vec![frame(&series_chunk( + "web-1", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + true, + ))]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![frame(&series_chunk( + "web-1", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + true, + ))]; + let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("disjoint grid points are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + let mut rows = rows(&batches); + rows.sort_by_key(|row| row.0); + assert_eq!( + vec![GRID_START, last], + rows.iter().map(|row| row.0).collect::>(), + "both series must survive with the one point each holds" + ); + for ((_, value, host, _), rate) in rows.iter().zip([1.0, 1.0]) { + assert_eq!(Some("web-1"), host.as_deref()); + let value = value.expect("the point each series holds has a value"); + assert!((rate - value).abs() < 1e-9, "{value} is not {rate}"); + } +} + +/// What a label set holds stays charged after the stream that first claimed it ends, and what a +/// later partition adds to it is charged as well. +/// +/// The map belongs to the plan: it outlives every partition and is released with the plan. Charging +/// each stream for what it inserted returned the bytes of a label set when *that* stream ended, +/// while the entry the next partition is checked against stayed in the map, so everything the map +/// held after the first partition was invisible to the pool. What is charged here is therefore read +/// between the partitions: it has to survive the first stream, and the runs the second partition +/// merges into the entry have to cost what they hold. +#[tokio::test] +async fn a_label_set_stays_charged_after_its_first_stream_ends() { + // The first series holds a value only at the first grid point and the second only at the last, + // so the two share a final label set without sharing an evaluation timestamp: the second + // partition is legal and has to be answered, and it has to find the first one's entry there. + let first = vec![frame(&series_chunk( + "web-0", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + true, + ))]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![frame(&series_chunk( + "web-0", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + true, + ))]; + let exec = plan(rate_tree(), memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + + // The first partition claims the label set, answers it and ends. + let mut stream = exec + .execute(0, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the first partition runs")); + } + assert_eq!(1, answered.len(), "one series over one grid point"); + drop(stream); + let charged = pool.reserved(); + assert!( + charged > 0, + "the first partition's label set is still in the plan's map, but nothing is charged for it" + ); + + // The second partition lands on that same label set on points the first one does not hold: it + // is answered, and the run it merges into the entry is what its charge grows by. + let mut stream = exec + .execute(1, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the second partition runs")); + } + assert_eq!(1, answered.len(), "disjoint points are legal"); + drop(stream); + assert!( + pool.reserved() > charged, + "the second partition's run grew the entry, but nothing was charged for it: {} is not \ + above {charged}", + pool.reserved() + ); + + // The entries live as long as the plan, so dropping it is what gives their bytes back. + drop(exec); + assert_eq!( + 0, + pool.reserved(), + "the plan's drop releases what its label sets held" + ); +} + +/// Each partition runs its own series and answers it, with no label set claimed twice. +#[tokio::test] +async fn every_partition_evaluates_its_own_series() { + let partitions = [1.0, 2.0, 3.0] + .into_iter() + .enumerate() + .map(|(index, scale)| { + vec![frame(&series_chunk( + &format!("web-{index}"), + &series_samples(scale), + true, + ))] + }) + .collect::>(); + let exec = plan(rate_tree(), memory_exec_partitioned(partitions)) + .expect("the runtime executes a rate tree"); + assert_eq!( + 3, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("each partition runs its own tree"); + assert_eq!(3, batches.len(), "one batch per partition"); + let rows = rows(&batches); + assert_eq!( + 3 * GRID_LEN as usize, + rows.len(), + "every partition reports its own series over the whole grid" + ); + // Batches arrive in whatever order the partitions settle, so the series are read apart. + for (index, scale) in [1.0, 2.0, 3.0].into_iter().enumerate() { + let host = format!("web-{index}"); + let mut series = rows + .iter() + .filter(|row| row.2.as_deref() == Some(host.as_str())) + .cloned() + .collect::>(); + series.sort_by_key(|row| row.0); + assert_eq!(GRID_LEN as usize, series.len(), "{host}: {series:?}"); + assert_series_rates(&series, &[(scale, host.as_str())]); + } +} + +/// The optimizer inserts no exchange between a two-partition scan and this node. +/// +/// This is the claim the whole arrangement rests on, so it is checked against the real rule rather +/// than against the flags that are supposed to satisfy it: a plan that reports the right +/// requirements can still be repartitioned, and that is exactly how the previous `KeyPartitioned` +/// requirement failed - a hash requirement is classed as necessary for any child with more than one +/// partition, so a hash exchange was inserted even though the child already satisfied it. +/// +/// The input here is a two-partition source, and `target_partitions` is raised above it because +/// that is when an optimizer is most willing to repartition. Both are the conditions under which the +/// old requirement broke, so a plan that survives them is one where the guarantee the scan makes +/// arrives at this node intact. +#[test] +fn the_optimizer_inserts_no_exchange_between_the_input_and_this_node() { + let exec = plan( + rate_tree(), + memory_exec_partitioned(vec![ + vec![frame(&series_chunk("web-0", &series_samples(1.0), true))], + vec![frame(&series_chunk("web-1", &series_samples(2.0), true))], + ]), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs a child with more than one partition" + ); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 8; + let optimized = EnsureRequirements::new() + .optimize(Arc::clone(&exec), &config) + .expect("the rule accepts the plan"); + + let exchanges = exchanges_between(&optimized, &exec); + assert!( + exchanges.is_empty(), + "the optimizer inserted {exchanges:?} between the input and the node, and an exchange \ + reorders a series' chunks: the node would close a series it has not finished" + ); + // No exchange is not the same as no change: a rule that raised the count some other way would + // still split a series. The count is the input's, which is what the node reports and what it + // runs one evaluator per. + assert_eq!( + 2, + optimized + .properties() + .output_partitioning() + .partition_count(), + "the optimizer kept no exchange but changed how many partitions the node runs" + ); + assert_eq!( + 2, + optimized.children()[0] + .properties() + .output_partitioning() + .partition_count(), + "the input's own partition count changed under the node" + ); +} + +/// The `RepartitionExec`s in `plan` that sit above `stop` and below no other `RepartitionExec`. +/// +/// Named by the kind the rule chose, so a failure says whether it was a hash or a round-robin +/// exchange: the two are inserted for different reasons and only the hash one is the trap. +fn exchanges_between(plan: &Arc, stop: &Arc) -> Vec { + let mut found = Vec::new(); + let mut stack = vec![Arc::clone(plan)]; + while let Some(node) = stack.pop() { + if Arc::ptr_eq(&node, stop) { + continue; + } + if let Some(exchange) = + node.downcast_ref::() + { + found.push(format!( + "{}({:?})", + exchange.name(), + exchange.partitioning() + )); + } + stack.extend(node.children().into_iter().cloned()); + } + found +} + +/// The node asks the optimizer for nothing, so no exchange can be inserted under it. +/// +/// A window holds the state of the series it is reading and closes it on the chunk carrying +/// `is_last`, so it can answer a series only if the whole series arrives in the partition it reads. +/// What makes that true is the input: a raw series frame comes from `SeriesScan`, whose contract is +/// that the same series (primary key) always reaches the same partition +/// (`src/mito2/src/read/series_scan.rs`), and the frame decoder keeps each series' chunks contiguous +/// within that partition. The node therefore needs the input's own layout kept, not a particular +/// layout of its own. +/// +/// The trap is the requirement writing that as `KeyPartitioned`, which reads like the honest +/// description of the guarantee. A `Hash|KeyPartitioned` requirement is classed as a hash +/// requirement, and a hash requirement forces a hash exchange whenever the child has more than one +/// partition - even where the child already satisfies it, and even when the partition counts match. +/// A hash exchange does not preserve order, so it would interleave one series' chunks across +/// partitions and the decoder would close a series it has not finished. Asking for +/// [`Distribution::UnspecifiedDistribution`] is what leaves the child alone, and declining extra +/// partitioning is the second half of the same thing: no round-robin is added to raise parallelism +/// either. The node is not pinned to one partition — it runs one tree per input partition and keeps +/// the input's count. +#[test] +fn the_input_is_left_alone_so_no_exchange_can_be_inserted() { + let exec = plan(rate_tree(), memory_exec_partitioned(vec![vec![], vec![]])) + .expect("the runtime executes a rate tree"); + let distributions = exec.input_distribution_requirements().into_per_child(); + // `Distribution` has no equality of its own, so the exact requirement is matched structurally: + // one child, and what it asks for is that the child's layout is left alone. + assert!( + matches!( + distributions.as_slice(), + [Distribution::UnspecifiedDistribution] + ), + "anything but leaving the child alone makes a hash exchange necessary: {distributions:?}" + ); + assert_eq!( + vec![false], + exec.benefits_from_input_partitioning(), + "raising the partition count would add a round-robin exchange" + ); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "the node's own count is its input's" + ); +} + +/// A partition index the input does not have is refused rather than answered. +/// +/// The node runs one evaluator per input partition and adds none of its own, so it forwards the +/// index to its input: the count it declares is the input's, and an index past that is the input's +/// refusal, not a series read from another partition. +#[tokio::test] +async fn a_partition_past_the_inputs_count_is_refused() { + let schema = frame_schema(); + let source = staged(schema.clone(), move || { + Box::pin(StagedStream::new(schema.clone(), Vec::new())) + }); + let exec = plan(rate_tree(), source).expect("the runtime executes a rate tree"); + assert_eq!( + 1, + exec.properties().output_partitioning().partition_count(), + "the node's partition count is its input's" + ); + exec.execute(0, SessionContext::default().task_ctx()) + .expect("the one partition the input has"); + let error = match exec.execute(1, SessionContext::default().task_ctx()) { + Ok(_) => panic!("a partition the input does not have is answered"), + Err(error) => error, + }; + assert!( + error.to_string().contains("partition"), + "refused for another reason: {error}" + ); +} + +/// The `rate(x[5m])` tree over two fields, with both fields and the tags in the output. +fn two_field_node() -> FusedSeries { + let tree = rate_tree_of(two_field_schema()); + let binding = binding_of( + &tree, + two_field_schema(), + vec![HOST.to_string(), NAME.to_string()], + ); + FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding, + input_plan(frame_schema_of(&two_field_schema())), + two_field_output_schema().to_dfschema_ref().unwrap(), + ) + .expect("the binding matches the tree") +} + +/// The output of [`two_field_node`]: both fields, the tags and the grid times. +fn two_field_output_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, true), + Field::new(OTHER, DataType::Float64, true), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])) +} + +/// One chunk of `{__name__=METRIC, host=, __type__=}` whose `field` holds `rows` and +/// whose other field is missing everywhere. +fn one_field_chunk(field: usize, host: &str, kind: &str, last: bool) -> RawSeriesChunk { + let rows = series_samples(1.0); + let (value, other) = if field == 0 { + (values_of(&rows), no_values(&rows)) + } else { + (no_values(&rows), values_of(&rows)) + }; + chunk_of( + &two_field_schema(), + &format!("{host}-{kind}"), + metric(&[(HOST, host), (TYPE, kind)], Some(METRIC)), + ×_of(&rows), + vec![value, other], + last, + ) +} + +/// Two series that close on one label set may still fill it, one field each. +/// +/// `rate` drops `__name__` and `__type__`, so the two series below leave the window as +/// `{host="web-1"}` and read as one label set. They hold values in *different* fields, which is the +/// pair [`crate::grid_series::merge_same_labelset`] merges: one value of one field at each position, +/// never two. Occupancy is therefore per field, and refusing this pair would refuse a series +/// Prometheus answers. +#[tokio::test] +async fn a_duplicate_label_set_of_complementary_fields_is_emitted() { + let exec = plan_of( + two_field_node(), + memory_exec_partitioned_with( + frame_schema_of(&two_field_schema()), + vec![ + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "counter", true)], + ), + frames_of( + &two_field_schema(), + &[one_field_chunk(1, "web-1", "gauge", true)], + ), + ], + ), + ) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let batches = collect(exec, SessionContext::default().task_ctx()) + .await + .expect("complementary fields are legal"); + assert_eq!(2, batches.len(), "each partition emits its own series"); + // Each batch answers one field of the shared label set and leaves the other missing: the two + // together are the merged series, with no value of one field holding a position twice. + let answered = batches + .iter() + .map(|batch| { + let host = batch + .column_by_name(HOST) + .and_then(|host| host.as_any().downcast_ref::()) + .expect("every batch has its tag column"); + assert!( + (0..host.len()).all(|row| host.value(row) == "web-1"), + "both batches answer the one label set" + ); + let present = [VALUE, OTHER] + .into_iter() + .filter(|field| { + let column = batch + .column_by_name(field) + .unwrap_or_else(|| panic!("every batch has a {field} column")) + .as_any() + .downcast_ref::() + .expect("the field column is Float64"); + column.null_count() < column.len() + }) + .collect::>(); + assert_eq!( + 1, + present.len(), + "each series fills exactly one field: {present:?}" + ); + present[0] + }) + .collect::>(); + let mut answered = answered; + answered.sort_unstable(); + assert_eq!(vec![OTHER, VALUE], answered, "the two fields are filled"); +} + +/// Two series that close on one label set may not hold the same field at the same position. +/// +/// This is the pair no merge resolves: at one grid position one field of the label set would carry +/// two values. It is refused even though the two series name different `__type__` labels, because +/// the window drops that label and what reaches the output is one label set twice. +#[tokio::test] +async fn a_duplicate_label_set_of_one_field_is_refused() { + let exec = plan_of( + two_field_node(), + memory_exec_partitioned_with( + frame_schema_of(&two_field_schema()), + vec![ + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "counter", true)], + ), + frames_of( + &two_field_schema(), + &[one_field_chunk(0, "web-1", "gauge", true)], + ), + ], + ), + ) + .expect("the runtime executes a rate tree"); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("both series hold a value of one field at one position"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// Two partitions whose label sets differ only in an empty value hold one label set. +/// +/// `{host="web-1", zone=""}` and `{host="web-1"}` are one series to Prometheus: an empty label value +/// is no label. An instant selection rewrites nothing, so both label sets reach this node's own +/// state as their frames wrote them, and a state that compared the tags as they arrived would file +/// them apart. Neither partition sees the other's series, so nothing else in the plan can refuse +/// the pair: this node's state is what compares them, and what it has to compare is the label set +/// Prometheus compares. +#[tokio::test] +async fn partitions_whose_label_sets_differ_only_in_an_empty_value_are_refused() { + let [empty, absent] = <[_; 2]>::try_from(zone_chunks()).expect("the fixture has two chunks"); + let exec = plan_of( + instant_node(), + memory_exec_partitioned(vec![frames(&[empty]), frames(&[absent])]), + ) + .expect("the runtime executes an instant selection"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + + let error = collect(exec, SessionContext::default().task_ctx()) + .await + .expect_err("both partitions hold one value of one field at the same position"); + let message = error_chain(&error); + assert!( + message.contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {message}" + ); + assert!( + message.contains("host=\"web-1\""), + "the label set must be readable: {message}" + ); +} + +/// A chunk holding no value records nothing, so a query that answers nothing holds nothing. +/// +/// A selection over a series that staleness or the range left empty still closes that series, and +/// hands over a closing chunk with no rows - [`crate::native::execute::no_row_chunk`] is where it is +/// built. Every chunk reaches [`LabelSets::claim`], so an implementation that wrote an entry per +/// chunk would record one per empty series and charge the plan's reservation for it, for a query +/// whose answer is empty: enough of them would exhaust a finite pool without a single row produced. +/// +/// The entries are checked rather than the charge, because a charge of zero is what a mistake here +/// would still show if the label bytes happened to be free: what must not grow is the table. +#[test] +fn a_chunk_holding_no_value_is_claimed_without_recording_it() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) + .expect("an empty run over the grid"); + let empty_series = |index: usize| Series { + metric: metric(&[(HOST, &format!("web-{index}"))], Some(METRIC)), + time: axis.clone(), + // Five positions were never filled, so no field occupies one. + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), + }; + + for index in 0..64 { + labels + .claim(&empty_series(index)) + .expect("an empty chunk cannot collide with anything"); + } + + assert_eq!( + 0, + labels.seen.len(), + "64 empty series were recorded, and the table outlives the query: {:?}", + labels.seen.keys().collect::>() + ); + assert_eq!( + 0, + pool.reserved(), + "nothing is held, so nothing may be charged" + ); +} + +/// An empty chunk does not stop a later chunk of the same label set from being recorded. +/// +/// Skipping the empty chunk is only safe while it also records nothing a later one needs: a series +/// that answers nothing holds no position, so a second series under one label set has to be free to +/// claim the positions it really holds. +#[test] +fn a_later_chunk_still_claims_the_label_set_an_empty_one_skipped() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let empty_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 0)]) + .expect("an empty run over the grid"); + let full_axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) + .expect("one run over the whole grid"); + + labels + .claim(&Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: empty_axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![f64::NAN; 0]))), + }) + .expect("an empty chunk cannot collide with anything"); + labels + .claim(&Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: full_axis, + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), + }) + .expect("the first chunk to hold a value claims the positions it holds"); + + assert_eq!(1, labels.seen.len(), "the positions are recorded once"); + assert!(pool.reserved() > 0, "the recorded positions are charged"); +} + +/// The claim compares the label set Prometheus compares, so an empty value is no label. +/// +/// Driven directly at the layer that keeps the state: a window and a map both drop an empty-valued +/// label before a chunk reaches this node, so no tree this crate can build produces the pair below. +/// The state has to key on the same normalization they apply all the same, or the equality it +/// enforces would be one Prometheus does not have. +#[test] +fn a_claim_compares_label_sets_without_empty_values() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let axis = TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 5)]) + .expect("one run over the whole grid"); + let series_of = |metric: MetricName| Series { + metric, + time: axis.clone(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; 5]))), + }; + + labels + .claim(&series_of(metric(&[(HOST, "h"), (ZONE, "")], Some(METRIC)))) + .expect("the first label set is claimed"); + let error = labels + .claim(&series_of(metric(&[(HOST, "h")], Some(METRIC)))) + .expect_err("an empty label value is no label"); + assert!( + error + .to_string() + .contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {error}" + ); + assert_eq!( + vec![(Arc::from(HOST), Arc::from("h"))], + label_set(&metric(&[(HOST, "h"), (ZONE, "")], None)), + "the empty value is not part of the label set" + ); +} + +/// A series whose chunks move forward, a hole apart, appends to its stored runs rather than +/// re-merging them. +/// +/// One grid point per chunk with a hole after it (`0, 2, 4, ...`) is a legal stream and the shape a +/// short `count_over_time` over sub-sampled, chunked input produces: the chunks of one label set +/// hold their values at disjoint positions, so they merge into one series and no position ever +/// carries two values. What it used to cost was quadratic - `claim` walked the whole stored history +/// to look for a collision and then rebuilt all of it with a merge, for each chunk - and all of it +/// ran under the label-set lock every partition shares, so a sparse series of `K` chunks made the +/// other partitions wait for `K(K+1)/2` run moves. +/// +/// The counter is the runtime's own, so this test measures the append path the native table takes +/// as well: one fix, not two that can drift apart again. +#[test] +fn a_forward_chunk_does_not_rebuild_the_stored_runs() { + // How many one-position chunks the series arrives in: enough that the quadratic cost it used to + // pay would be unmistakable, since that cost was the square of this figure. + const CHUNKS: usize = 64; + + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + // Two positions per chunk, so a hole follows every point the series holds. + let grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 2 * CHUNKS as u64).unwrap()) + .unwrap(), + ); + let chunk = |position: usize| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(position, 1)]) + .expect("one position of the grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + + let before = absorb_rebuilds(); + for index in 0..CHUNKS { + labels + .claim(&chunk(2 * index)) + .expect("the chunks of one label set hold values at disjoint positions"); + } + + assert_eq!( + 0, + absorb_rebuilds() - before, + "the {CHUNKS} forward chunks rebuilt the stored runs instead of appending to them" + ); + assert_eq!( + 1, + labels.seen.len(), + "every chunk belongs to one label set, so the history is one entry" + ); +} + +/// A backward disjoint claim after a spare-capacity run merge updates the exact memory charge. +#[test] +fn a_backward_disjoint_claim_after_spare_capacity_is_charged_correctly() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |start, len| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) + .unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), + }; + labels.claim(&chunk(10, 2)).unwrap(); + let runs = labels + .seen + .values_mut() + .next() + .unwrap() + .first_mut() + .unwrap(); + let old_capacity = runs.capacity(); + runs.reserve(16); + let new_capacity = runs.capacity(); + labels + .reservation + .as_mut() + .unwrap() + .try_grow((new_capacity - old_capacity) * std::mem::size_of::()) + .unwrap(); + let before = pool.reserved(); + + labels + .claim(&chunk(0, 1)) + .expect("the position is disjoint and earlier than stored history"); + let runs = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 1), (10, 12)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert_eq!( + before - (new_capacity - runs.capacity()) * std::mem::size_of::(), + pool.reserved(), + "the accounting follows the replacement vector's actual capacity" + ); +} + +/// A backward claim that joins two spare-capacity runs charges a replacement vector's shrinkage. +#[test] +fn a_backward_claim_merging_spare_capacity_runs_accounts_for_shrinkage() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 3).unwrap()).unwrap()); + let chunk = |values: Vec>| Series { + metric: metric(&[(HOST, "shrink")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 3)]) + .expect("the chunk covers the small grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(values))), + }; + + labels + .claim(&chunk(vec![Some(1.0), None, Some(1.0)])) + .expect("the first claim stores runs at positions 0 and 2"); + let runs = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 1), (2, 3)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + let before_capacity = runs.capacity(); + let before_reserved = pool.reserved(); + + labels + .claim(&chunk(vec![None, Some(1.0), None])) + .expect("the disjoint middle position joins both existing runs"); + + let runs = labels.seen.values().next().unwrap().first().unwrap(); + let after_capacity = runs.capacity(); + assert_eq!( + vec![(0, 3)], + runs.iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert!( + before_capacity > after_capacity, + "the regression requires replacement-vector shrinkage (before={before_capacity}, after={after_capacity})" + ); + let run_size = std::mem::size_of::(); + assert_eq!( + before_reserved - (before_capacity - after_capacity) * run_size, + pool.reserved(), + "the reservation tracks the replacement vector's actual capacity decrease" + ); +} + +/// A newly retained field is charged for the cloned run vector's capacity, not incoming slack. +#[test] +fn a_new_field_claim_is_charged_for_its_retained_capacity() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 4).unwrap()).unwrap()); + let first = Series { + metric: metric(&[(HOST, "new-field")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels.claim(&first).unwrap(); + + let incoming = Series { + metric: metric(&[(HOST, "new-field")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(0, 4)]).unwrap(), + fields: SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![None, None, None, None])), + Samples::from_float(Float64Array::from(vec![Some(1.0), None, Some(2.0), None])), + ]) + .unwrap(), + }; + let incoming_runs = occupied_runs(&incoming).unwrap(); + assert_eq!(2, incoming_runs[1].len()); + assert!( + incoming_runs[1].capacity() > incoming_runs[1].len(), + "the incoming run vector has spare capacity" + ); + let before = pool.reserved(); + labels.claim(&incoming).unwrap(); + + let retained = &labels.seen.values().next().unwrap()[1]; + assert_eq!(2, retained.len()); + assert_eq!( + retained.capacity() * std::mem::size_of::(), + pool.reserved() - before, + "the reservation charges the retained cloned vector capacity" + ); + assert!( + retained.capacity() < incoming_runs[1].capacity(), + "the incoming spare capacity was not retained or charged" + ); +} + +/// Chunked fields are traversed sequentially while occupancy is converted into grid runs. +#[test] +fn chunked_samples_occupancy_preserves_sparse_rows() { + let values = Samples::try_from_chunks(vec![ + Samples::from_float(Float64Array::from(vec![Some(1.0), None])), + Samples::from_float(Float64Array::from(vec![None, Some(2.0)])), + ]) + .unwrap(); + let series = Series { + metric: metric(&[(HOST, "chunked")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid()), vec![GridSpan::new(0, 4)]).unwrap(), + fields: SeriesFields::one(values), + }; + let runs = occupied_runs(&series).expect("chunked occupancy is read through a forward cursor"); + assert_eq!( + vec![(0, 1), (3, 4)], + runs[0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); +} + +/// Appending a run that touches stored history reuses its last slot instead of rebuilding it. +#[test] +fn a_touching_append_with_full_capacity_does_not_rebuild_history() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |start| Series { + metric: metric(&[(HOST, "touch")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels.claim(&chunk(0)).unwrap(); + let stored = labels + .seen + .values_mut() + .next() + .unwrap() + .first_mut() + .unwrap(); + stored.shrink_to_fit(); + let cap = stored.capacity(); + let allocation = stored.as_ptr(); + assert_eq!( + stored.len(), + cap, + "the test requires a full-capacity run list" + ); + let before = absorb_rebuilds(); + labels + .claim(&chunk(1)) + .expect("the touching position extends the same run"); + let stored = labels.seen.values().next().unwrap().first().unwrap(); + assert_eq!( + vec![(0, 2)], + stored + .iter() + .map(|run| (run.start, run.end)) + .collect::>() + ); + assert_eq!( + cap, + stored.capacity(), + "extending a run did not allocate or copy its history" + ); + assert_eq!( + allocation, + stored.as_ptr(), + "extending a run retained the stored vector allocation" + ); + assert_eq!( + before, + absorb_rebuilds(), + "touching append uses the append path" + ); +} + +/// A failed claim does not mutate the shared entry, while its reservation owner remains alive. +#[test] +fn a_refused_growth_keeps_shared_label_entry_unchanged() { + use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; + let limit = 4096; + let pool = Arc::new(GreedyMemoryPool::new(limit)); + let erased: Arc = pool.clone(); + let mut labels = LabelSets::default(); + labels.attach(&erased); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + let chunk = |host: &str, start| Series { + metric: metric(&[(HOST, host)], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, 1)]).unwrap(), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0]))), + }; + labels + .claim(&chunk("shared", 0)) + .expect("initial label fits"); + let key = label_set(&chunk("shared", 0).metric); + let before_runs = labels.seen[&key][0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>(); + let before_reserved = pool.reserved(); + let blocker: MemoryReservation = MemoryConsumer::new("test blocker").register(&erased); + blocker + .try_grow(limit - before_reserved) + .expect("the blocker occupies remaining capacity"); + let error = labels + .claim(&chunk("shared", 2)) + .expect_err("growth exceeds the pool"); + assert!(!error.to_string().is_empty(), "pool refusal is reported"); + assert_eq!( + before_runs, + labels.seen[&key][0] + .iter() + .map(|run| (run.start, run.end)) + .collect::>(), + "the persistent entry is unchanged" + ); + assert_eq!(limit, pool.reserved(), "both reservations remain live"); + drop(blocker); + assert_eq!( + before_reserved, + pool.reserved(), + "the plan-owned entry remains charged" + ); +} + +/// Binding ordering includes field names, while schema ordering refuses unequal DF schemas. +#[test] +fn unequal_binding_and_output_schema_are_not_ordered_equal() { + let tree = rate_tree(); + let left = binding(&tree); + let other_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: "different".into(), + }]), + eval_grid: None, + }); + let right = binding_of(&tree, other_schema, left.tags.clone()); + assert_ne!(left, right); + assert_ne!(Some(std::cmp::Ordering::Equal), left.partial_cmp(&right)); + + let left_node = fused(tree.clone()); + let nullability_changed = Arc::new(Schema::new(vec![ + Field::new(VALUE, DataType::Float64, false), + Field::new(HOST, DataType::Utf8, true), + Field::new(NAME, DataType::Utf8, true), + Field::new( + TIME, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let right_node = FusedSeries::try_new( + Arc::from(PROMQL), + tree, + binding(&rate_tree()), + input_plan(frame_schema()), + nullability_changed.to_dfschema_ref().unwrap(), + ) + .unwrap(); + assert_ne!(left_node, right_node); + assert_ne!( + Some(std::cmp::Ordering::Equal), + left_node.partial_cmp(&right_node) + ); +} + +/// A chunk reaching back before the last stored run still merges into the history. +/// +/// The forward path only takes a chunk that starts at or after the end of the last stored run; a +/// chunk that starts before it can interleave with what is stored, so it has to go through the full +/// merge, and the result still has to be sorted, disjoint and with the runs that touch joined - +/// which is what the look-up in [`LabelSets::claim`] relies on. +/// +/// A chunk reaching back over a stored run is the other case, and it stays refused: `[0, 11)` shares +/// position 10 with the stored `[10, 12)`, so it is a collision and the stored history is left +/// as it was. +#[test] +fn a_chunk_reaching_back_still_merges() { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let mut labels = LabelSets::default(); + labels.attach(&pool); + let grid = + Arc::new(EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 16).unwrap()).unwrap()); + // One chunk holding `len` positions from `start`, every one of them with a value. + let chunk = |start: usize, len: usize| Series { + metric: metric(&[(HOST, "h")], Some(METRIC)), + time: TimeAxis::from_grid_spans(Arc::clone(&grid), vec![GridSpan::new(start, len)]) + .expect("one run over the grid"), + fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![1.0; len]))), + }; + // The runs stored under the one label set, which one field of the output holds. + let stored = |labels: &LabelSets| -> Vec<(usize, usize)> { + labels + .seen + .values() + .next() + .expect("the chunk was recorded under one label set") + .first() + .expect("the series has one field") + .iter() + .map(|run| (run.start, run.end)) + .collect() + }; + + labels + .claim(&chunk(10, 2)) + .expect("the first chunk claims the positions it holds"); + assert_eq!( + vec![(10, 12)], + stored(&labels), + "the run `[10, 12)` is stored" + ); + + let error = labels + .claim(&chunk(0, 11)) + .expect_err("`[0, 11)` holds a value at position 10, where `[10, 12)` holds one"); + assert!( + error + .to_string() + .contains("vector cannot contain metrics with the same labelset"), + "refused for another reason: {error}" + ); + assert_eq!( + vec![(10, 12)], + stored(&labels), + "a refused claim left the stored history as it was" + ); + + let before = absorb_rebuilds(); + labels + .claim(&chunk(0, 10)) + .expect("`[0, 10)` reaches back before the last stored run without sharing a position"); + assert_eq!( + vec![(0, 12)], + stored(&labels), + "the run reaching back did not merge into the stored one and join it" + ); + assert_eq!( + 1, + absorb_rebuilds() - before, + "a chunk starting before the last stored run was appended rather than merged" + ); +} diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index 73fe2e94ee5..de52accb9b0 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -118,6 +118,9 @@ use crate::native::plan::{ BoundInstant, BoundScalar, BoundVectorOp, Grouping, NativeExpr, NativeNode, NativeTree, NativeType, }; +#[cfg(test)] +pub(crate) use crate::native::runs::absorb_rebuilds; +use crate::native::runs::{Run, absorb}; use crate::range_eval::RangeEvaluator; use crate::select::{self, SortOrder}; use crate::vector_match::{ @@ -576,20 +579,6 @@ impl NodeBuffer { } } -/// One run of grid positions, half-open, that hold a value in one field. -#[derive(Clone, Copy, Debug)] -struct Run { - start: usize, - end: usize, -} - -impl Run { - /// Whether two runs share a position. - fn overlaps(self, other: Self) -> bool { - self.start < other.end && other.start < self.end - } -} - /// The positions one node has handed out, per output label set and field. /// /// A node that rewrites labels — a `Map` dropping the metric name, a `Window` dropping the name @@ -824,109 +813,6 @@ fn misaligned_field( )) } -// How many times `absorb` rebuilt the stored runs instead of appending to them. -// -// Test-only, and per thread: the runs stored afterwards are the same either way, so a rebuild shows -// only in what it costs, and a test that reads its own thread's figure is not disturbed by another -// test's chunks. -#[cfg(test)] -thread_local! { - static ABSORB_REBUILDS: std::cell::Cell = const { std::cell::Cell::new(0) }; -} - -/// How many times this thread's [`absorb`] rebuilt the stored runs. -#[cfg(test)] -fn absorb_rebuilds() -> usize { - ABSORB_REBUILDS.with(std::cell::Cell::get) -} - -/// Merges `incoming` into `stored`, keeping the result in position order and joining the runs that -/// touch. -/// -/// Within a chunk one pass over both sides is what keeps this linear, since each side is sorted and -/// disjoint. Adding the runs one at a time and rebuilding the array after each would cost the square -/// of the runs a chunk yields: a chunk of one hundred thousand positions alternating between values -/// and nulls yields fifty thousand runs, and re-merging after each is over a billion element moves -/// for a single series. -/// -/// One pass per chunk is still a pass over everything the series has stored so far, so it is not -/// free once the *history* grows across chunks: a series split into `K` chunks of one position each, -/// a hole apart, stores one more run per chunk and moves `K(K+1)/2` of them, with no label collision -/// and no illegal input anywhere. A chunk that starts no earlier than the end of the last stored run -/// only ever adds to the end — the stored runs are sorted and disjoint, so every one of them ends -/// before anything this chunk holds — and it is appended instead, which is amortised constant per -/// run. That is the shape a scan hands over as its positions move forward chunk after chunk. A chunk -/// that starts before the last stored run can split a run or reach back past it, so it still costs -/// the full merge. -/// -/// Either way the result is sorted, disjoint, and has the runs that touch joined, which is what -/// [`Emitted::occupy`] looks a position up by. -fn absorb(stored: &mut Vec, incoming: &[Run]) { - if incoming.is_empty() { - return; - } - if stored - .last() - .is_none_or(|last| incoming[0].start >= last.end) - { - append_runs(stored, incoming); - } else { - merge_runs(stored, incoming); - } -} - -/// Appends `incoming` to `stored`, all of which lies past the last stored run. -/// -/// Nothing already stored has to be placed again: the two sides are each in position order and -/// nothing `incoming` holds can reach back into what is stored, so its runs are pushed on the end and -/// the vector's own growth is the only move — amortised over the runs appended. Where its first run -/// touches the end of the last stored run the two are one run, as they are after a merge, so it is -/// joined rather than left beside it. -fn append_runs(stored: &mut Vec, incoming: &[Run]) { - let mut incoming = incoming; - if let Some(last) = stored.last_mut() - && let Some(first) = incoming.first() - && first.start <= last.end - { - last.end = last.end.max(first.end); - incoming = &incoming[1..]; - } - stored.extend_from_slice(incoming); -} - -/// Merges `incoming` into `stored`, where the two can interleave. -/// -/// One pass is enough because each side is sorted and disjoint: whichever run starts first is -/// placed, and it joins the run before it where the two touch. -fn merge_runs(stored: &mut Vec, incoming: &[Run]) { - #[cfg(test)] - ABSORB_REBUILDS.with(|rebuilds| rebuilds.set(rebuilds.get() + 1)); - let mut merged: Vec = Vec::with_capacity(stored.len() + incoming.len()); - let (mut a, mut b) = (0, 0); - while a < stored.len() || b < incoming.len() { - let larger = match (stored.get(a), incoming.get(b)) { - (Some(left), Some(right)) if right.start < left.start => { - b += 1; - *right - } - (Some(left), _) => { - a += 1; - *left - } - (None, Some(right)) => { - b += 1; - *right - } - (None, None) => break, - }; - match merged.last_mut() { - Some(last) if larger.start <= last.end => last.end = last.end.max(larger.end), - _ => merged.push(larger), - } - } - *stored = merged; -} - /// What a map node applies. struct MapRuntime { ops: Arc<[MathOp]>, diff --git a/src/promql/src/native/mod.rs b/src/promql/src/native/mod.rs index 5375c92f1c2..fd04dcfda32 100644 --- a/src/promql/src/native/mod.rs +++ b/src/promql/src/native/mod.rs @@ -33,3 +33,4 @@ mod display; pub mod execute; pub mod metrics; pub mod plan; +pub(crate) mod runs; diff --git a/src/promql/src/native/runs.rs b/src/promql/src/native/runs.rs new file mode 100644 index 00000000000..8de325a5386 --- /dev/null +++ b/src/promql/src/native/runs.rs @@ -0,0 +1,146 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The grid positions a field occupies, as sorted, disjoint runs. +//! +//! A node that rewrites labels — a `Map` dropping the metric name, a `Window` dropping the name and +//! the empty-valued labels — can turn two input series into one label set, and every layer above +//! then reads the two chunks as one series. What decides whether that pair is legal is which +//! positions of each field already hold a value: two values of one field at one position cannot be +//! merged away, while a chunk that fills positions the set left empty merges cleanly, as +//! [`crate::grid_series::merge_same_labelset`] does. +//! +//! The state is kept per label set, per field, as runs, so a chunk of many positions costs one run +//! rather than one entry per position. Keeping it in one place is what keeps the two layers that +//! check it — the native runtime's own [`crate::native::execute::Emitted`] table and the +//! cross-partition [`crate::extension_plan::fused_series::LabelSets`] — from drifting apart: both +//! ask the same functions for the same answers, so the fast path one of them earned is the fast path +//! the other gets. + +/// One run of grid positions, half-open, that hold a value in one field. +#[derive(Clone, Copy, Debug)] +pub(crate) struct Run { + pub(crate) start: usize, + pub(crate) end: usize, +} + +impl Run { + /// Whether two runs share a position. + pub(crate) fn overlaps(self, other: Self) -> bool { + self.start < other.end && other.start < self.end + } +} + +// How many times `absorb` rebuilt the stored runs instead of appending to them. +// +// Test-only, and per thread: the runs stored afterwards are the same either way, so a rebuild shows +// only in what it costs, and a test that reads its own thread's figure is not disturbed by another +// test's chunks. +#[cfg(test)] +thread_local! { + static ABSORB_REBUILDS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// How many times this thread's [`absorb`] rebuilt the stored runs. +#[cfg(test)] +pub(crate) fn absorb_rebuilds() -> usize { + ABSORB_REBUILDS.with(std::cell::Cell::get) +} + +/// Merges `incoming` into `stored`, keeping the result in position order and joining the runs that +/// touch. +/// +/// Within a chunk one pass over both sides is what keeps this linear, since each side is sorted and +/// disjoint. Adding the runs one at a time and rebuilding the array after each would cost the square +/// of the runs a chunk yields: a chunk of one hundred thousand positions alternating between values +/// and nulls yields fifty thousand runs, and re-merging after each is over a billion element moves +/// for a single series. +/// +/// One pass per chunk is still a pass over everything the series has stored so far, so it is not +/// free once the *history* grows across chunks: a series split into `K` chunks of one position each, +/// a hole apart, stores one more run per chunk and moves `K(K+1)/2` of them, with no label collision +/// and no illegal input anywhere. A chunk that starts no earlier than the end of the last stored run +/// only ever adds to the end — the stored runs are sorted and disjoint, so every one of them ends +/// before anything this chunk holds — and it is appended instead, which is amortised constant per +/// run. That is the shape a scan hands over as its positions move forward chunk after chunk. A chunk +/// that starts before the last stored run can split a run or reach back past it, so it still costs +/// the full merge. +/// +/// Either way the result is sorted, disjoint, and has the runs that touch joined, which is what a +/// look-up by [`Run::overlaps`] on the runs stored here relies on. +pub(crate) fn absorb(stored: &mut Vec, incoming: &[Run]) { + if incoming.is_empty() { + return; + } + if stored + .last() + .is_none_or(|last| incoming[0].start >= last.end) + { + append_runs(stored, incoming); + } else { + merge_runs(stored, incoming); + } +} + +/// Appends `incoming` to `stored`, all of which lies past the last stored run. +/// +/// Nothing already stored has to be placed again: the two sides are each in position order and +/// nothing `incoming` holds can reach back into what is stored, so its runs are pushed on the end and +/// the vector's own growth is the only move — amortised over the runs appended. Where its first run +/// touches the end of the last stored run the two are one run, as they are after a merge, so it is +/// joined rather than left beside it. +fn append_runs(stored: &mut Vec, incoming: &[Run]) { + let mut incoming = incoming; + if let Some(last) = stored.last_mut() + && let Some(first) = incoming.first() + && first.start <= last.end + { + last.end = last.end.max(first.end); + incoming = &incoming[1..]; + } + stored.extend_from_slice(incoming); +} + +/// Merges `incoming` into `stored`, where the two can interleave. +/// +/// One pass is enough because each side is sorted and disjoint: whichever run starts first is +/// placed, and it joins the run before it where the two touch. +fn merge_runs(stored: &mut Vec, incoming: &[Run]) { + #[cfg(test)] + ABSORB_REBUILDS.with(|rebuilds| rebuilds.set(rebuilds.get() + 1)); + let mut merged: Vec = Vec::with_capacity(stored.len() + incoming.len()); + let (mut a, mut b) = (0, 0); + while a < stored.len() || b < incoming.len() { + let larger = match (stored.get(a), incoming.get(b)) { + (Some(left), Some(right)) if right.start < left.start => { + b += 1; + *right + } + (Some(left), _) => { + a += 1; + *left + } + (None, Some(right)) => { + b += 1; + *right + } + (None, None) => break, + }; + match merged.last_mut() { + Some(last) if larger.start <= last.end => last.end = last.end.max(larger.end), + _ => merged.push(larger), + } + } + *stored = merged; +} From 85c5dc0d69a06baf825e05a78b9fbb392c526a88 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 08:35:30 +0800 Subject: [PATCH 4/4] test(promql): pin the refusal to a kind the runtime still refuses (#9900) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/extension_plan/fused_series/tests.rs | 36 ++++++++++++++++--- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index 20bdebfe2f0..017f4ef6628 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -274,7 +274,30 @@ fn nested_window_tree() -> NativeTree { .expect("the bound tree is legal") } -/// The `rate(x[5m])` tree with a `sum by (job)` over it, which the runtime does not execute. +/// The same tree asking for `quantile`, whose state the runtime does not keep yet. +fn quantile_aggregate_tree() -> NativeTree { + let complete = aggregate_tree(); + let root = complete.root(); + let NativeExpr::Aggregate { input, op } = &root.expr else { + unreachable!("the fixture is an aggregate over a window"); + }; + NativeTree::try_new(NativeNode { + expr: NativeExpr::Aggregate { + input: input.clone(), + op: BoundAggregate { + kind: AggregateKind::Quantile { quantile: 0.5 }, + group: op.group.clone(), + stage: AggregateStage::Complete, + }, + }, + output: root.output.clone(), + }) + .expect("the bound tree is legal") +} + +/// The `rate(x[5m])` tree with a `sum by (job)` over it. +/// +/// A `Complete` aggregate: every contribution to each group lies inside one execution. fn aggregate_tree() -> NativeTree { let rate = rate_tree(); let window = rate.root().clone(); @@ -1194,16 +1217,21 @@ async fn a_series_spanning_chunks_keeps_absolute_positions_and_closes_once() { } /// A tree the runtime does not execute is refused when the execution plan is built. +/// +/// The runtime folds the supported `Complete` kinds now, so the tree below asks for one it still +/// refuses on purpose: `quantile` needs the parameter and the member buffer that the other kinds +/// do not. What the test pins is the refusal happening while the plan is built, not the particular +/// kind of node that provokes it. #[tokio::test] async fn a_tree_the_runtime_cannot_run_is_rejected_when_the_exec_is_built() { - let node = fused(aggregate_tree()); + let node = fused(quantile_aggregate_tree()); let error = plan_of(node, memory_exec(Vec::new())) - .expect_err("the runtime does not execute aggregate nodes"); + .expect_err("the runtime does not execute a quantile aggregate"); assert!( matches!(error, DataFusionError::NotImplemented(_)), "the runtime refuses the tree it cannot run: {error}" ); - assert!(error.to_string().contains("Aggregate"), "{error}"); + assert!(error.to_string().contains("Quantile"), "{error}"); } /// A tree evaluated on another grid than the node declares is refused before it runs.