rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 06:46:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(promql): add a series frame output to the fused native node (#9943)

Browse Source 
I hereby agree to the terms of the [GreptimeDB
CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md).

## Refer to a related PR or issue link (optional)

- **Stacked on #9942** (PR A, the unified SeriesFrame codec).
- GitHub cannot target a fork branch, so this PR's base reads `narthex`.
**Review #9942 first.**
- The commits unique to this PR are `b32573f9cbf`, `16b3e428951`,
`d5547c55bdf` and
  `f2041d8d9de` (test-only lint cleanup). The chain is
  `narthex (771347df5d) → #9942 (8bc53c4871) → this PR (CLEANUP_SHA)`.
- Merge order remains **#9942 → #9943**. Once #9942 lands, rebase this
PR onto the resulting
  `narthex`; do not merge the codec changes twice.
- Spec: Narthex Unified SeriesFrame refactor, PR B of two.

## What's changed and what's your intention?

The fused native node could only answer one way: a row table with the
grid timestamps, one column
per series field and one per declared tag. That is what a query returns,
but it is not what a
consumer of the native path needs when it has to tell series and
positions apart — the labels of a
series, and which grid positions it actually holds, are both flattened
away by the row shape.

This adds a second output format, **fixed when the plan is built**
rather than chosen at run time:

```rust
pub enum FusedOutputFormat { PlainRows, SeriesFrame }
```

| | `PlainRows` (default) | `SeriesFrame` |
| --- | --- | --- |
| output schema | the existing row table (time / fields / tags) | the
codec's declared evaluated frame |
| encoding | `series_to_record_batch` | straight from `EvaluatedChunk`
via `SeriesFrameCursor` |
| reachable from a query | yes | no |
| `EXPLAIN` | byte-identical | appends ` (output: series_frame)` |

The format is part of the plan because it decides the node's output
schema, so it is part of the
node's identity (`PartialEq`/`Hash`/`PartialOrd`), its `EXPLAIN` text,
and its rebuild over a new
child. A format error is refused while the plan is built, not when the
first poll happens.

The `SeriesFrame` schema is derived from the
**`NativeTree.root().output` evaluated `SeriesSchema`**,
never from the original scan schema or the input tags: after a label
rewrite or an aggregate, the
result's labels and fields are what count.

### Independent chains remain independent

`is_last` closes the current chain, not the entire label set. Separate
chains may therefore carry
identical complete labels and each close independently. This PR adds no
plan-owned global label
identity management; combining contributions across partitions remains
the consumer's responsibility.
Tests check exact positions, payload bits and separate finality for both
consecutive same-label
chains in one encoder and chains in separate partitions.

### Parallelism is unchanged

For the supported single-source Frame path, output partition count
equals the input scan's
partition count. No `CoalescePartitionsExec`, `RepartitionExec` or
round-robin is introduced in
either format.

`series_frame_preserves_four_input_partitions` checks the in-memory
boundary. The standalone
integration test additionally captures a real single-region
`SeriesScanExec` at `parallelism=4`
and executes every partition in each output format.

### Implementation notes worth knowing

1. **Stream and batch schemas agree.** The codec-derived output schema
includes the explicit
`contiguous` ordering declaration. Published batches are checked against
that declared schema.
2. **The end of the execution must flush.** The frame builder buffers up
to a batch across chunks,
   so missing the final flush loses the tail rows.
3. **`Poll::Pending` also flushes.** Otherwise a closing row stays in
the builder behind an input
that may never become ready — the same convention PR A's `PackStream`
already follows.

### Key code locations

- `src/promql/src/extension_plan.rs`: exports `FusedOutputFormat`.
- `src/promql/src/extension_plan/fused_series.rs`: default/explicit
format constructors, root-output
schema derivation, plan identity/rebuild/EXPLAIN, direct chunk encoding
and stream flushing.
- `src/promql/src/extension_plan/fused_series/tests.rs` and
`tests/frame_output.rs`: internal
  logical/physical execution and lifecycle/semantic regressions.
- `tests-integration/tests/http.rs`: real standalone scan replay in both
output formats.

### Review fixes and acceptance

- **Batch-level backpressure:** the stream retains the existing
`SeriesFrameCursor` and advances
only until one batch is ready, then returns control to the consumer. It
no longer encodes an
entire evaluated chunk into a `ready` queue before returning the first
batch. Pending native
output is drained before more upstream input; a pre-EOF stall does not
close the executor.
  PlainRows and partial-state conversion keep their original timing.
- **Real scan acceptance:** the standalone File integration fixture has
one physical region and
`parallelism=4`. The test captures its actual `SeriesScanExec`, replays
the internal fused plan
in both formats, executes all four nonempty partitions, and checks the
exact 24 expected rows,
eight Frame chain closings and per-partition native output metrics.
Frame is not exposed as a
  user query's final result to perform this test.
- Replaced the vacuous sparse fixture with exact nonempty positions
`0,2,5`, two fields with
partial NULLs and preserved ordinary-NaN bits, and a PlainRows
comparison. Added fixed-value
`abs(rate(...))` and label-rewrite comparisons, including a new label
outside the input tag
projection and two independent same-label chains in one encoder.
Cross-partition chains also
  check exact positions and value bits.
- Added retained-cursor/one-batch checks, buffered upstream-error
termination, and early-drop
ownership checks. Reintroducing eager encoding compiles but makes the
backpressure regression
fail at its queue assertion; restoring the reviewed source makes it
pass.
- Removed the redundant physical scalar/Frame guard; the shared
root-output validator remains.
- Independent correctness and anti-over-design review: **APPROVE** for
`8bc53c48714..d5547c55bdf`. Review approval is separate from new-head
CI.

### Limitations

- `SeriesFrame` is **not reachable from a user query** yet (production
defaults to `PlainRows`).
Per spec §8.7 the acceptance executes the internal physical plan and
decodes its output, instead
  of exposing the frame as a query's final result format.
- No sqlness result changes: the production planner still constructs
`PlainRows`. The existing
standalone integration case is extended to execute and decode the
internal Frame plan.
- Rebased onto `narthex` at `771347df5d9`, including #9936, #9882 and
#9886. The native zero-source
scalar, `timestamp()`/instant-`@`, multi-input and partial-state paths
are preserved. A Frame over
a scalar root is rejected because this format requires evaluated
samples.
- The multi-source case is refused for frames rather than implemented:
this node still answers only
partition 0 when it has more than one raw source, so a frame over
several sources would have to
  change that first.
- Consumer contract: a sparse chunk decodes into **several** `Series`
(one frame row per grid run),
so a consumer must concatenate by position and cannot assume "one chunk
→ one decoded result".
Nothing consumes frames in this PR, but Complete Aggregate's SeriesFrame
input must respect it.
- Not verified: Windows; the results below are Linux.

### Does this break API or data compatibility?

`PlainRows` remains the default and the only format a user query
reaches. This increment leaves
its conversion and plan display paths in place. The added tests compare
exact Frame semantics
without changing the user-facing result format. No persisted format is
touched.

### Tests

Coverage includes codec-derived schema, logical rebuild and EXPLAIN,
dense and genuinely sparse
output, multi-field NULL/NaN values, composed math and label rewrites,
four in-memory and four real
scan partitions, independent same-label chains, partial-state schema,
constructor refusals,
zero-sample closing, Pending/EOF flush, buffered error, cancellation and
incremental encoding.

<!-- CURRENT_VERIFICATION_START -->
Main verification on `d5547c55bdf` (Linux; default features):

- `make fmt-check`: passed.
- `cargo nextest run -p datatypes -p common-query -p table -p promql -p
query --retries 0`:
  **2698 passed, 2 skipped**.
- `cargo nextest run -p tests-integration --test main -E
'test(test_native_promql_route_plans_the_query) & test(file)' --retries
0 --success-output immediate`:
**1 passed, 771 skipped**. Both formats produced per-partition decoded
row counts
**[15, 3, 3, 3]**; exact 24-row results and eight Frame chain closings
were asserted.
- `cargo check --workspace --all-targets`: passed.
- The final strict Clippy pass found two test-only hygiene issues in the
new tests: an
unnecessary `f64` cast and two `println!` calls that the workspace's
`print_stdout` lint
rejects. The follow-up types the existing float literal explicitly,
removes the cast, and drops
the diagnostics, whose values were already asserted by the surrounding
exact checks. It changes
  no production code or assertions.

- Final test-only cleanup `f2041d8d9de` on top of the reviewed
`d5547c55bdf`:
`make fmt-check` passed; `cargo clippy -p promql -p query -p
tests-integration --all-targets -- -D warnings`
passed; the affected
`abs_rate_frame_matches_plain_rows_with_exact_values_and_labels`
(1 passed) and real File scan replay (1 passed, 771 skipped) passed
again. The full five-crate
suite above was run at `d5547c55bdf`; this last commit touches tests
only and was not followed by
  another full-suite run.

Other focused evidence: the eager-encoding mutation compiled and then
failed the queue-boundary
assertion; restored production source passed. License/TOML/diff checks
passed. These are local
Linux results, not a claim of Windows or new-head CI success; full
all-feature workspace Clippy
and the full workspace test suite were not run.
<!-- CURRENT_VERIFICATION_END -->

## PR Checklist
Please convert it to a draft if some of the following conditions are not
met.

- [x] I have written the necessary rustdoc comments.
- [x] I have added the necessary unit tests and integration tests.
- [ ] This PR requires documentation updates.
- [x] API changes are backward compatible.
- [x] Schema or data changes are backward compatible.
- [ ] This PR needs to be backported to release branches, and I have
added the `backport-<target>` labels (e.g. `backport-v1.3` targets
`release/v1.3`).
This commit is contained in:
discord9
2026-10-26 15:00:05 +08:00
committed by GitHub
parent eeca48bc43 f2041d8d9d
commit 1cab5f1624
7 changed files with 1721 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -15119,6 +15119,7 @@ dependencies = [
"paste",
"pipeline",
"plugins",
"promql",
"prost 0.14.1",
"query",
"rand 0.9.4",
src/promql/src/extension_plan.rs
+3 -1
View File
@@ -41,7 +41,9 @@ 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 fused_series::{
FusedOutputFormat, FusedSeries, FusedSeriesBinding, FusedSeriesExec, FusedSeriesStream,
};
pub use histogram_fold::{
HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream,
};
src/promql/src/extension_plan/fused_series.rs
+254 -17
View File
@@ -44,7 +44,9 @@ 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 common_query::series_frame::{
ensure_frame_matches_label_source, pack_evaluated_series, unpack_raw_series,
};
use datafusion::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray};
use datafusion::arrow::buffer::ScalarBuffer;
use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
@@ -64,6 +66,7 @@ use datafusion::physical_plan::{
Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
};
use datatypes::series::SeriesOrdering;
use datatypes::series::frame::{FrameLabelMode, SeriesFrameBuilder, SeriesFrameCursor};
use datatypes::series::raw::RawSeriesChunk;
use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesSchema};
use datatypes::series::time::EvalGrid;
@@ -164,6 +167,39 @@ impl PartialOrd for FusedSeriesBinding {
}
}
/// The output representation fixed when a fused plan is built.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum FusedOutputFormat {
/// The existing tabular query result.
PlainRows,
/// A self-describing evaluated SeriesFrame.
SeriesFrame,
}
fn evaluated_output(tree: &NativeTree) -> DataFusionResult<Arc<SeriesSchema>> {
match &tree.root().output {
NativeType::Samples(schema) if schema.eval_grid.is_some() => Ok(Arc::clone(schema)),
NativeType::Samples(_) => Err(DataFusionError::Plan(
"FusedSeries frame output requires evaluated samples with a grid".into(),
)),
NativeType::AggregateState(_) => Err(DataFusionError::Plan(
"FusedSeries frame output does not support an aggregate-state root".into(),
)),
NativeType::Scalar(_) => Err(DataFusionError::Plan(
"FusedSeries frame output does not support a scalar root".into(),
)),
}
}
fn series_frame_output_schema(schema: &SeriesSchema) -> DataFusionResult<SchemaRef> {
let stream =
pack_evaluated_series(Box::pin(futures::stream::empty()), Arc::new(schema.clone()))
.map_err(|error| {
DataFusionError::Plan(format!("FusedSeries cannot declare frame schema: {error}"))
})?;
Ok(stream.schema())
}
/// 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
@@ -186,6 +222,8 @@ pub struct FusedSeries {
/// This row-state bridge is temporary: native final aggregation will consume series frames
/// directly, removing state-row encoding and reconstruction.
partial_fold: Option<BoundAggregate>,
/// Fixed output representation; partial state export always remains PlainRows.
format: FusedOutputFormat,
}
impl FusedSeries {
@@ -201,8 +239,39 @@ impl FusedSeries {
binding: FusedSeriesBinding,
children: Vec<LogicalPlan>,
output_schema: DFSchemaRef,
) -> DataFusionResult<Self> {
Self::try_new_with_format(
promql,
tree,
binding,
children,
output_schema,
FusedOutputFormat::PlainRows,
)
}
/// Creates a node with an explicitly selected output format.
pub fn try_new_with_format(
promql: Arc<str>,
tree: NativeTree,
binding: FusedSeriesBinding,
children: Vec<LogicalPlan>,
output_schema: DFSchemaRef,
format: FusedOutputFormat,
) -> DataFusionResult<Self> {
validate_binding(&tree, &binding, children.len())?;
if format == FusedOutputFormat::SeriesFrame && children.len() > 1 {
return Err(DataFusionError::Plan(
"FusedSeries SeriesFrame output does not yet support multiple sources".into(),
));
}
let output_schema = if format == FusedOutputFormat::SeriesFrame {
Arc::new(datafusion::common::DFSchema::try_from(
series_frame_output_schema(evaluated_output(&tree)?.as_ref())?,
)?)
} else {
output_schema
};
Ok(Self {
promql,
tree,
@@ -210,6 +279,7 @@ impl FusedSeries {
children,
output_schema,
partial_fold: None,
format,
})
}
@@ -243,6 +313,7 @@ impl FusedSeries {
children,
output_schema,
partial_fold: Some(op),
format: FusedOutputFormat::PlainRows,
})
}
@@ -296,6 +367,7 @@ impl FusedSeries {
self.binding.time_index.clone(),
self.binding.tags.clone(),
self.partial_fold.clone(),
self.format,
)?))
}
}
@@ -327,6 +399,9 @@ impl UserDefinedLogicalNodeCore for FusedSeries {
self.promql,
render_partial_fold(op)
),
None if self.format == FusedOutputFormat::SeriesFrame => {
write!(f, "FusedSeries: {} (output: series_frame)", self.promql)
}
None => write!(f, "FusedSeries: {}", self.promql),
}
}
@@ -344,6 +419,7 @@ impl UserDefinedLogicalNodeCore for FusedSeries {
children: inputs,
output_schema: Arc::clone(&self.output_schema),
partial_fold: self.partial_fold.clone(),
format: self.format,
})
}
}
@@ -361,6 +437,7 @@ impl PartialEq for FusedSeries {
&& self.children == other.children
&& self.output_schema == other.output_schema
&& partial_key(self.partial_fold.as_ref()) == partial_key(other.partial_fold.as_ref())
&& self.format == other.format
}
}
@@ -373,6 +450,7 @@ impl Hash for FusedSeries {
self.children.hash(state);
self.output_schema.hash(state);
partial_key(self.partial_fold.as_ref()).hash(state);
self.format.hash(state);
}
}
@@ -400,6 +478,10 @@ impl PartialOrd for FusedSeries {
Some(Ordering::Equal) => {}
ord => return ord,
}
match self.format.cmp(&other.format) {
Ordering::Equal => {}
ord => return Some(ord),
}
match tree_input_fields(&self.tree).cmp(&tree_input_fields(&other.tree)) {
Ordering::Equal => {}
ord => return Some(ord),
@@ -783,6 +865,7 @@ pub struct FusedSeriesExec {
metric: ExecutionPlanMetricsSet,
/// Aggregate state exported by each partition, when this is the partial mode.
partial_fold: Option<BoundAggregate>,
format: FusedOutputFormat,
/// This node's partitioning, emission and boundedness.
properties: Arc<PlanProperties>,
}
@@ -807,9 +890,19 @@ impl FusedSeriesExec {
time_column: String,
tags: Vec<String>,
) -> DataFusionResult<Self> {
Self::build(promql, inputs, tree, output_schema, time_column, tags, None)
Self::build(
promql,
inputs,
tree,
output_schema,
time_column,
tags,
None,
FusedOutputFormat::PlainRows,
)
}
#[allow(clippy::too_many_arguments)]
fn build(
promql: Arc<str>,
inputs: Vec<(Arc<dyn ExecutionPlan>, RawLabelSource, SeriesOrdering)>,
@@ -818,7 +911,13 @@ impl FusedSeriesExec {
time_column: String,
tags: Vec<String>,
partial_fold: Option<BoundAggregate>,
format: FusedOutputFormat,
) -> DataFusionResult<Self> {
if partial_fold.is_none() && format == FusedOutputFormat::SeriesFrame && inputs.len() > 1 {
return Err(DataFusionError::Plan(
"FusedSeries SeriesFrame output does not yet support multiple sources".into(),
));
}
let raw_schemas = validate_physical_tree(&tree)?;
if inputs.len() != raw_schemas.len() {
return Err(DataFusionError::Plan(format!(
@@ -853,7 +952,19 @@ impl FusedSeriesExec {
}
}
}
let classification = if partial_fold.is_some() {
let frame_schema = if partial_fold.is_none() && format == FusedOutputFormat::SeriesFrame {
let declared = series_frame_output_schema(evaluated_output(&tree)?.as_ref())?;
if declared != output_schema {
return Err(DataFusionError::Plan(
"FusedSeriesExec output schema does not match the evaluated frame layout"
.into(),
));
}
true
} else {
false
};
let classification = if partial_fold.is_some() || frame_schema {
Vec::new()
} else if scalar_root {
let scalar_fields = scalar_root_fields(&output_schema, &output_tags, &time_column)?;
@@ -931,6 +1042,7 @@ impl FusedSeriesExec {
metric: ExecutionPlanMetricsSet::new(),
properties,
partial_fold,
format,
})
}
@@ -953,6 +1065,8 @@ impl DisplayAs for FusedSeriesExec {
write!(f, "FusedSeriesExec: {}\n{}", self.promql, self.tree)?;
if let Some(op) = &self.partial_fold {
write!(f, "\npartial state: {}", render_partial_fold(op))?;
} else if self.format == FusedOutputFormat::SeriesFrame {
write!(f, "\noutput: series_frame")?;
}
Ok(())
}
@@ -1028,6 +1142,7 @@ impl ExecutionPlan for FusedSeriesExec {
self.time_column.clone(),
self.tags.to_vec(),
self.partial_fold.clone(),
self.format,
)?))
}
@@ -1098,6 +1213,30 @@ impl ExecutionPlan for FusedSeriesExec {
if streams.is_empty() {
executor.end_of_input()?;
}
// A frame is encoded from the tree's own evaluated output, and only when the node answers
// values rather than state: partial state is a third output the format does not reach.
let frame_mode =
self.partial_fold.is_none() && self.format == FusedOutputFormat::SeriesFrame;
let eval_schema = if frame_mode {
Some(evaluated_output(&self.tree)?)
} else {
None
};
let frame_builder = eval_schema
.as_ref()
.map(|schema| {
SeriesFrameBuilder::try_new(
schema,
FrameLabelMode::Explicit,
None,
SeriesFrameBuilder::DEFAULT_MAX_ROWS,
SeriesFrameBuilder::DEFAULT_MAX_BYTES,
)
.map_err(|error| {
DataFusionError::Plan(format!("cannot create evaluated frame builder: {error}"))
})
})
.transpose()?;
Ok(Box::pin(FusedSeriesStream {
inputs: streams,
executor: Some(executor),
@@ -1106,6 +1245,9 @@ impl ExecutionPlan for FusedSeriesExec {
tags: Arc::clone(&self.tags),
output_schema: self.output_schema.clone(),
partial_fold,
frame_builder,
eval_schema,
frame_cursor: None,
_export_budget: Arc::clone(&export_budget),
ready: VecDeque::new(),
done: false,
@@ -1226,6 +1368,10 @@ pub struct FusedSeriesStream {
output_schema: SchemaRef,
/// Per-partition aggregate state fold, if this stream exports a partial result.
partial_fold: Option<PartialFold>,
frame_builder: Option<SeriesFrameBuilder>,
eval_schema: Option<Arc<SeriesSchema>>,
/// The evaluated chunk whose frame is still being appended across downstream polls.
frame_cursor: Option<SeriesFrameCursor>,
/// Keeps exported state batches charged for as long as this stream holds them.
_export_budget: Arc<Mutex<Option<MemoryReservation>>>,
/// The batches the executor settled since the last poll of the input.
@@ -1249,11 +1395,18 @@ impl FusedSeriesStream {
})? {
executor.push_raw(*input_id, chunk.clone())?;
}
let mut settled = Vec::new();
while let Some(chunk) = executor.next_chunk()? {
settled.push(chunk);
if self.frame_builder.is_some() {
Ok(())
} else {
let mut settled = Vec::new();
while let Some(chunk) = executor.next_chunk()? {
settled.push(chunk);
}
for chunk in settled {
self.convert(chunk)?;
}
Ok(())
}
self.convert(settled)
}
/// Ends the input, so the tree settles what it was holding.
@@ -1300,6 +1453,8 @@ impl FusedSeriesStream {
self.done = true;
self.executor = None;
self.partial_fold = None;
self.frame_builder = None;
self.frame_cursor = None;
self.inputs.clear();
// This reservation belongs only to this execution. Release it even if the consumer keeps
// the failed stream and the plan alive.
@@ -1311,26 +1466,65 @@ impl FusedSeriesStream {
self.ready.clear();
}
/// Converts settled chunks into query rows or folds them into partial aggregate state.
fn convert(&mut self, chunks: Vec<EvaluatedChunk>) -> DataFusionResult<()> {
/// Converts one settled chunk into the node's output: query rows, an evaluated series frame,
/// or partition-local state for a later aggregate merge.
fn convert(&mut self, chunk: EvaluatedChunk) -> DataFusionResult<()> {
if let Some(fold) = self.partial_fold.as_mut() {
for chunk in chunks {
fold.update(chunk)?;
}
fold.update(chunk)?;
return Ok(());
}
for chunk in &chunks {
if let Some(schema) = self.eval_schema.as_ref()
&& self.frame_builder.is_some()
{
self.frame_cursor = Some(
SeriesFrameCursor::try_new_evaluated(schema, chunk.series, chunk.is_last).map_err(
|error| {
DataFusionError::Execution(format!("cannot frame evaluated chunk: {error}"))
},
)?,
);
} else {
let batch = series_to_record_batch(
&self.classification,
&self.tags,
&self.output_schema,
chunk,
&chunk,
)?;
self.ready.push_back(batch);
}
Ok(())
}
/// Appends only until the current evaluated chunk yields one batch.
fn advance_frame_cursor(&mut self) -> DataFusionResult<Option<RecordBatch>> {
let (Some(cursor), Some(builder)) =
(self.frame_cursor.as_mut(), self.frame_builder.as_mut())
else {
return Ok(None);
};
match cursor.append_to(builder).map_err(|error| {
DataFusionError::Execution(format!("cannot append evaluated frame: {error}"))
})? {
Some(batch) => Ok(Some(batch.with_schema(self.output_schema.clone())?)),
None => {
self.frame_cursor = None;
Ok(None)
}
}
}
fn flush_frames(&mut self) -> DataFusionResult<()> {
if let Some(builder) = self.frame_builder.as_mut()
&& let Some(batch) = builder.flush().map_err(|error| {
DataFusionError::Execution(format!("cannot flush evaluated frame: {error}"))
})?
{
self.ready
.push_back(batch.with_schema(self.output_schema.clone())?);
}
Ok(())
}
/// Publishes the partition's fold only after the input and native execution have drained.
fn finish_fold(&mut self) -> DataFusionResult<()> {
if let Some(mut fold) = self.partial_fold.take() {
@@ -1357,22 +1551,56 @@ impl Stream for FusedSeriesStream {
// 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 this.frame_cursor.is_some() {
match this.advance_frame_cursor() {
Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))),
Ok(None) => {}
Err(error) => {
this.fail();
return Poll::Ready(Some(Err(error)));
}
}
}
if let Some(batch) = this.ready.pop_front() {
return Poll::Ready(Some(Ok(batch)));
}
if !this.done && this.frame_builder.is_some() && this.frame_cursor.is_none() {
let next = this
.executor
.as_mut()
.map_or(Ok(None), NativeExecutor::next_chunk);
match next {
Ok(Some(chunk)) => {
if let Err(error) = this.convert(chunk) {
this.fail();
return Poll::Ready(Some(Err(error)));
}
continue;
}
Ok(None) => {
if this.inputs.is_empty() {
this.executor = None;
}
}
Err(error) => {
this.fail();
return Poll::Ready(Some(Err(error)));
}
}
}
if this.done {
return Poll::Ready(None);
}
let Some((source, input)) = this.inputs.front_mut() else {
match this.pull_one() {
Ok(Some(chunk)) => {
if let Err(error) = this.convert(vec![chunk]) {
if let Err(error) = this.convert(chunk) {
this.fail();
return Poll::Ready(Some(Err(error)));
}
}
Ok(None) => {
if let Err(error) = this.finish_fold() {
if let Err(error) = this.finish_fold().and_then(|()| this.flush_frames()) {
this.fail();
return Poll::Ready(Some(Err(error)));
}
@@ -1388,7 +1616,16 @@ impl Stream for FusedSeriesStream {
};
let (source, polled) = (*source, input.as_mut().poll_next(cx));
match polled {
Poll::Pending => return Poll::Pending,
Poll::Pending => {
if let Err(error) = this.flush_frames() {
this.fail();
return Poll::Ready(Some(Err(error)));
}
if let Some(batch) = this.ready.pop_front() {
return Poll::Ready(Some(Ok(batch)));
}
return Poll::Pending;
}
Poll::Ready(None) => {
this.inputs.pop_front();
if this.inputs.is_empty()
src/promql/src/extension_plan/fused_series/tests.rs
+838 -5
View File
@@ -32,10 +32,11 @@
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::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use common_query::series_frame::{unpack_evaluated_series, unpack_raw_series};
use datafusion::arrow::array::{Array, Float64Array, StringArray, TimestampMillisecondArray};
use datafusion::arrow::buffer::{Buffer, ScalarBuffer};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
@@ -65,17 +66,23 @@ use datafusion::physical_plan::{
use datafusion::prelude::SessionContext;
use datatypes::prelude::ConcreteDataType;
use datatypes::series::frame::{
FrameLabelMode, LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding, LabelDecodeSpec,
LabelDecodeTag, pack_raw_series_chunk, raw_series_frame_schema, series_frame_schema,
FrameContent, FrameLabelMode, LABEL_DECODE_VERSION, LabelDecodeColumn, LabelDecodeEncoding,
LabelDecodeSpec, LabelDecodeTag, SeriesFrameBuilder, SeriesFrameDeclaration,
pack_raw_series_chunk, raw_series_frame_schema, 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, SeriesLabels, SeriesOrdering};
use futures::{Stream, StreamExt, TryStreamExt};
use store_api::region_engine::RawLabelSource;
use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream};
use super::{FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, series_to_record_batch};
mod frame_output;
use super::{
FusedOutputFormat, FusedSeries, FusedSeriesBinding, FusedSeriesStream, OutputColumn,
POLL_BUDGET, series_to_record_batch,
};
use crate::evaluated::EvaluatedChunk;
use crate::extension_plan::aggregate_exchange::state_schema;
use crate::native::execute::{InputId, NativeExecutionState, NativeExecutor};
@@ -85,6 +92,7 @@ use crate::native::plan::{
};
use crate::range_eval::{RangeFunction, RangeWindow};
use crate::vector_match::{BinaryMatching, BinaryOp, Cardinality, FillValues, MatchingLabels};
use common_query::series_frame::series_frame_ordering;
/// The query every tree here is compiled from.
const PROMQL: &str = "rate(x[5m])";
@@ -2496,6 +2504,215 @@ async fn a_ready_batch_is_returned_before_later_input_is_read() {
assert!(total > 1, "the fixture has to outlast the first answer");
}
fn direct_frame_stream(
schema: Arc<SeriesSchema>,
input: Option<SendableRawSeriesStream>,
) -> FusedSeriesStream {
let output_schema = series_frame_schema(&schema, FrameLabelMode::Explicit, None).unwrap();
let builder = SeriesFrameBuilder::try_new(
&schema,
FrameLabelMode::Explicit,
None,
SeriesFrameBuilder::DEFAULT_MAX_ROWS,
SeriesFrameBuilder::DEFAULT_MAX_BYTES,
)
.unwrap();
FusedSeriesStream {
inputs: input.into_iter().map(|input| (0, input)).collect(),
executor: None,
input_ids: Vec::new(),
classification: Arc::from([]),
tags: Arc::from([]),
output_schema,
partial_fold: None,
frame_builder: Some(builder),
eval_schema: Some(schema),
frame_cursor: None,
_export_budget: Arc::new(Mutex::new(None)),
ready: VecDeque::new(),
done: false,
}
}
fn sparse_eval_chunk(rows: usize, value: f64) -> (Arc<SeriesSchema>, EvaluatedChunk) {
let grid =
Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1, (rows * 2) as u64).unwrap()).unwrap());
let schema = Arc::new(SeriesSchema {
fields: Arc::from([FieldSpec { name: VALUE.into() }]),
eval_grid: Some(Arc::clone(&grid)),
});
let axis =
TimeAxis::from_grid_spans(grid, (0..rows).map(|i| GridSpan::new(i * 2, 1)).collect())
.unwrap();
let chunk = EvaluatedChunk::try_new(
Series {
metric: SeriesLabels::try_new(Some(Arc::from(METRIC)), Vec::new()).unwrap(),
time: axis,
fields: SeriesFields::one(Samples::from_float(Float64Array::from(vec![value; rows]))),
},
true,
)
.unwrap();
(schema, chunk)
}
/// A large evaluated chunk is encoded over successive polls rather than queued as all of its
/// physical frame batches at once.
#[tokio::test]
async fn a_large_evaluated_chunk_keeps_a_cursor_between_polls() {
const ROWS: usize = 800;
let value = f64::from_bits(0x4009_21fb_5444_2d18);
let (schema, chunk) = sparse_eval_chunk(ROWS, value);
let mut stream = direct_frame_stream(Arc::clone(&schema), None);
stream
.convert(chunk)
.expect("the evaluated chunk becomes a cursor");
assert!(
stream.ready.is_empty(),
"conversion must not eagerly queue frame batches"
);
assert!(
stream.frame_cursor.is_some(),
"unencoded rows remain owned by the cursor"
);
let first = stream.next().await.unwrap().unwrap();
assert!(first.num_rows() <= SeriesFrameBuilder::DEFAULT_MAX_ROWS);
assert!(
stream.frame_cursor.is_some(),
"the first batch cannot drain 800 rows"
);
assert!(stream.ready.is_empty());
let mut batches = vec![first];
while let Some(batch) = stream.next().await {
batches.push(batch.unwrap());
}
assert!(batches.len() > 1);
let decoded = unpack_evaluated_series(
Box::pin(RecordBatchStreamAdapter::new(
batches[0].schema(),
futures::stream::iter(batches.into_iter().map(Ok)),
)),
None,
)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let mut decoded_rows = Vec::new();
let mut closing_chunks = 0;
for (series, last) in &decoded {
if *last {
closing_chunks += 1;
}
assert_eq!(Some(METRIC), series.metric.name());
for (position, row) in series.time.positions() {
let value = match series.fields.as_slice()[0].sample(row) {
datatypes::series::samples::SampleRef::Float(value) => value,
_ => panic!("the decoded field remains a present Float64"),
};
decoded_rows.push((position, value.to_bits()));
}
}
decoded_rows.sort_unstable_by_key(|(position, _)| *position);
assert_eq!(1, closing_chunks, "the chunk closes exactly once");
assert_eq!(
(0..ROWS)
.map(|i| (i * 2, 0x4009_21fb_5444_2d18))
.collect::<Vec<_>>(),
decoded_rows,
"every sparse position carries the exact expected float bits",
);
assert!(stream.next().await.is_none(), "EOF remains stable");
assert!(stream.next().await.is_none(), "repeated EOF remains stable");
}
/// The cursor and its input are both released when the consumer drops an unfinished frame.
#[tokio::test]
async fn dropping_a_partly_emitted_frame_releases_its_input() {
const ROWS: usize = 800;
let (schema, mut chunk) = sparse_eval_chunk(ROWS, 3.0);
let name = Arc::<str>::from(METRIC);
let name_weak = Arc::downgrade(&name);
chunk
.series
.metric
.set("__name__", Arc::clone(&name))
.unwrap();
drop(name);
let raw = raw_schema();
let raw_chunk = series_chunk("cursor-source", &[(0, 1.0)], true);
let raw_frame = frame_of(&raw, &raw_chunk);
let dropped = Arc::new(AtomicBool::new(false));
let raw_input = Box::pin(StagedStream::new(frame_schema(), vec![raw_frame]).dropped(&dropped));
let input = unpack_raw_series(raw_input, Some(Arc::clone(&raw))).unwrap();
let mut stream = direct_frame_stream(Arc::clone(&schema), Some(input));
stream.convert(chunk).unwrap();
assert!(stream.next().await.unwrap().unwrap().num_rows() > 0);
assert!(stream.frame_cursor.is_some());
assert!(
name_weak.upgrade().is_some(),
"the cursor retains its unencoded label identity"
);
assert!(!dropped.load(Ordering::SeqCst));
assert_eq!(1, stream.inputs.len());
drop(stream);
assert!(dropped.load(Ordering::SeqCst));
assert!(
name_weak.upgrade().is_none(),
"dropping the cursor releases its label identity"
);
}
/// An upstream error clears rows buffered in the frame builder before any batch is delivered.
#[tokio::test]
async fn frame_upstream_error_discards_buffered_rows_and_terminates() {
let (schema, chunk) = sparse_eval_chunk(2, 5.0);
let raw = raw_schema();
let dropped = Arc::new(AtomicBool::new(false));
let input = unpack_raw_series(
Box::pin(
StagedStream::new(frame_schema(), Vec::new())
.error("frame input failed")
.dropped(&dropped),
),
Some(Arc::clone(&raw)),
)
.unwrap();
let mut stream = direct_frame_stream(schema, Some(input));
stream.convert(chunk).unwrap();
assert!(stream.advance_frame_cursor().unwrap().is_none());
assert!(
stream.frame_cursor.is_none(),
"small rows sit buffered in the builder"
);
assert!(stream.ready.is_empty());
let error = stream
.next()
.await
.expect("the upstream error is yielded")
.expect_err("buffered frame rows cannot hide the source error");
assert!(error.to_string().contains("frame input failed"));
assert!(
stream.next().await.is_none(),
"the failed stream terminates once"
);
assert!(
stream.next().await.is_none(),
"EOF after failure remains stable"
);
assert!(
dropped.load(Ordering::SeqCst),
"failure releases the source"
);
assert!(stream.frame_builder.is_none());
assert!(stream.frame_cursor.is_none());
assert!(stream.ready.is_empty());
}
/// 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
@@ -3731,3 +3948,619 @@ async fn a_failed_input_exports_no_partial_state() {
"refused for another reason: {error}"
);
}
/// Frame output is planned from the evaluated root and is emitted without flattening to rows.
#[tokio::test]
async fn series_frame_output_uses_evaluated_codec_schema() {
let tree = rate_tree();
let frame_node = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.expect("evaluated samples have a frame schema");
assert!(
frame_node
.schema()
.fields()
.iter()
.all(|field| field.name() != "ts"),
"frame schema comes from the codec rather than the row output"
);
assert!(format!("{frame_node:?}").contains("FusedSeries"));
let exec = plan_of(frame_node, memory_exec(one_series_frames())).unwrap();
assert_eq!(1, exec.properties().partitioning.partition_count());
let batches = collect(exec.clone(), SessionContext::default().task_ctx())
.await
.unwrap();
assert!(!batches.is_empty());
assert!(batches.iter().all(|batch| batch.schema() == exec.schema()));
assert!(batches.iter().any(|batch| batch.num_rows() > 0));
}
/// Frame output refuses the currently unsupported multi-source physical shape at planning time.
#[test]
fn series_frame_rejects_multiple_sources_during_planning() {
let tree = binary_tree_over_sources(BinaryOp::Add, raw_schema(), raw_schema(), 0, 1);
let err = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(frame_schema()), input_plan(frame_schema())],
binary_output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.unwrap_err();
assert!(err.to_string().contains("multiple sources"));
}
/// A real four-partition input remains four independent frame-output executions.
#[tokio::test]
async fn series_frame_preserves_four_input_partitions() {
let partitions = (0..4)
.map(|index| {
frames(&[series_chunk(
&format!("part-{index}"),
&series_samples((index + 1) as f64),
true,
)])
})
.collect::<Vec<_>>();
let tree = rate_tree();
let node = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.unwrap();
let exec = plan_of(node, memory_exec_partitioned(partitions)).unwrap();
assert_eq!(4, exec.properties().partitioning.partition_count());
let schema = exec.schema();
let mut closes = 0;
for partition in 0..4 {
let stream = exec
.execute(partition, SessionContext::default().task_ctx())
.unwrap();
let decoded = unpack_evaluated_series(stream, None)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
1,
decoded.len(),
"one chain closes in partition {partition}"
);
assert!(decoded[0].1);
closes += decoded.iter().filter(|(_, last)| *last).count();
let values = decoded[0].0.fields.get(0).unwrap();
let value = values.iter().next().unwrap();
let datatypes::series::samples::SampleRef::Float(value) = value else {
panic!("float sample expected")
};
assert!((value - (partition + 1) as f64).abs() < 1e-9);
assert_eq!(schema, exec.schema());
}
assert_eq!(4, closes);
}
/// PlainRows and SeriesFrame encode precisely the same rate values, timestamps and labels.
#[tokio::test]
async fn series_frame_rate_matches_plain_rows_bitwise() {
let input = memory_exec(two_series_frames());
let plain = plan(rate_tree(), Arc::clone(&input)).unwrap();
let plain_rows = rows(
&collect(plain, SessionContext::default().task_ctx())
.await
.unwrap(),
);
let tree = rate_tree();
let node = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.unwrap();
let exec = plan_of(node, input).unwrap();
let batches = collect(exec.clone(), SessionContext::default().task_ctx())
.await
.unwrap();
assert!(batches.iter().all(|batch| batch.schema() == exec.schema()));
let stream = Box::pin(RecordBatchStreamAdapter::new(
exec.schema(),
futures::stream::iter(batches.into_iter().map(Ok)),
));
let decoded = unpack_evaluated_series(stream, None)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let mut frame_rows = Vec::new();
for (series, _) in decoded {
let timestamps = series.time.timestamps().unwrap();
let samples = series.fields.get(0).unwrap();
for (row, timestamp) in timestamps.iter().enumerate() {
let value = samples.iter().nth(row).unwrap();
let datatypes::series::samples::SampleRef::Float(value) = value else {
panic!("float sample expected")
};
frame_rows.push((
*timestamp,
value.to_bits(),
series.metric.get(HOST).map(str::to_owned),
series.metric.name().map(str::to_owned),
));
}
}
let mut expected = plain_rows
.into_iter()
.map(|(time, value, host, name)| (time, value.map(f64::to_bits), host, name))
.collect::<Vec<_>>();
frame_rows.sort();
expected.sort();
assert_eq!(
expected,
frame_rows
.into_iter()
.map(|(time, bits, host, name)| (time, Some(bits), host, name))
.collect::<Vec<_>>()
);
}
/// The partial-state node keeps its state schema and explain annotation, independently of output format.
#[tokio::test]
async fn partial_fold_keeps_plain_rows_state_export() {
let node = partial_node(AggregateKind::Sum);
assert_eq!(FusedOutputFormat::PlainRows, node.format);
let exec = plan_of(node, memory_exec(one_series_frames())).unwrap();
let state = state_schema(&[HOST.to_string()]);
assert_eq!(state.as_ref(), exec.schema().as_ref());
assert!(
displayable(exec.as_ref())
.indent(false)
.to_string()
.contains("partial state:")
);
let batch = state_of(&exec, 0, &state).await;
assert_eq!(state.as_ref(), batch.schema().as_ref());
}
/// `EXPLAIN` says nothing extra about the row table and names the frame output.
///
/// The row table is what a query returns, so its node line has to stay exactly what it was; a frame
/// output answers something else, and the line is the only place a reader can see that.
#[test]
fn the_explain_text_names_the_frame_format_only() {
let rows = fused(rate_tree());
assert_eq!(
"FusedSeries: rate(x[5m])",
Explain(&rows).to_string(),
"the row-table line is what it always was"
);
let framed = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
rows.tree().clone(),
rows.binding().clone(),
rows.inputs().into_iter().cloned().collect(),
Arc::clone(rows.schema()),
FusedOutputFormat::SeriesFrame,
)
.expect("a rate tree is evaluated on a grid");
assert!(
Explain(&framed)
.to_string()
.contains("(output: series_frame)"),
"a frame output names its format: {}",
Explain(&framed)
);
let plain = plan(rate_tree(), memory_exec(Vec::new())).expect("the runtime executes the tree");
let text = displayable(plain.as_ref()).indent(false).to_string();
assert!(
text.starts_with("FusedSeriesExec: rate(x[5m])\n"),
"the row-table node line is unchanged: {text}"
);
}
/// The logical schema follows the format, and a frame schema is the codec's own declaration.
///
/// The frame layout is not assembled here: it is read back through the declaration parser, so what
/// the plan claims is what a consumer of the frame will decode.
#[test]
fn the_logical_schema_follows_the_output_format() {
let rows = fused(rate_tree());
assert_eq!(FusedOutputFormat::PlainRows, rows.format);
assert_eq!(&output_schema(), rows.schema().inner());
let framed = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
rows.tree().clone(),
rows.binding().clone(),
rows.inputs().into_iter().cloned().collect(),
Arc::clone(rows.schema()),
FusedOutputFormat::SeriesFrame,
)
.expect("a rate tree is evaluated on a grid");
assert_eq!(FusedOutputFormat::SeriesFrame, framed.format);
let declaration =
SeriesFrameDeclaration::parse(framed.schema().inner().as_ref()).expect("a frame schema");
assert_eq!(FrameContent::Evaluated, declaration.content());
assert_eq!(FrameLabelMode::Explicit, declaration.labels());
assert_eq!(
SeriesOrdering::Contiguous,
series_frame_ordering(framed.schema().inner().as_ref())
.expect("the frame declares its ordering"),
"computed values promise per-chain contiguity and nothing stronger"
);
}
/// Rebuilding a frame node keeps the format it was built with.
///
/// The format decides the node's schema, so a rebuild that dropped it would answer rows where the
/// plan said frame. `with_exprs_and_inputs` is what the optimizer uses.
#[test]
fn rebuilding_a_frame_node_keeps_its_output_format() {
let rows = fused(rate_tree());
let framed = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
rows.tree().clone(),
rows.binding().clone(),
rows.inputs().into_iter().cloned().collect(),
Arc::clone(rows.schema()),
FusedOutputFormat::SeriesFrame,
)
.expect("a rate tree is evaluated on a grid");
let rebuilt = framed
.with_exprs_and_inputs(
Vec::new(),
framed.inputs().into_iter().cloned().collect::<Vec<_>>(),
)
.expect("the node rebuilds over its own input");
assert_eq!(FusedOutputFormat::SeriesFrame, rebuilt.format);
assert_eq!(framed.schema(), rebuilt.schema());
assert_eq!(Explain(&framed).to_string(), Explain(&rebuilt).to_string());
}
/// A tree whose root is not evaluated samples cannot be framed, and says so when the plan is built.
///
/// A scalar has no series to name, so it is not a series of float samples; an aggregate-state root is
/// refused by the same check for the same reason, and `partial_node` cannot even be built with a
/// frame format. The refusal is a planning error, so a caller learns about it before the query runs
/// rather than from a partition that fails later.
#[test]
fn a_series_frame_refuses_a_root_it_cannot_carry() {
let scalar = scalar_tree(grid(), 1.0);
// A scalar tree binds no raw source, so it is built over no child at all: the refusal below is
// then about the root, not about a source count that was wrong first.
let binding = FusedSeriesBinding {
grid: grid(),
metric_name: Some(Arc::from(METRIC)),
time_index: TIME.to_string(),
tags: vec![HOST.to_string(), NAME.to_string()],
};
let error = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
scalar,
binding,
Vec::new(),
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.expect_err("a scalar is not a series to frame");
assert!(
matches!(error, DataFusionError::Plan(_)),
"refused with {error}"
);
assert!(
error.to_string().contains("scalar"),
"the refusal says what the root is: {error}"
);
// The partial mode is a different output entirely, so a frame format cannot even be asked for.
assert_eq!(
FusedOutputFormat::PlainRows,
partial_node(AggregateKind::Sum).format
);
}
/// A pending input hands the rows a frame already encoded over before it reports pending.
///
/// The builder buffers up to a batch, and a closing row only exists once its chunk is encoded. If a
/// pending input could hold those rows, a consumer waiting for a series to close would wait for input
/// that may never arrive.
///
/// This polls by hand rather than awaiting the stream: `StagedStream` reports pending once and then
/// wakes itself, so an `await` would run straight through the boundary and satisfy every later
/// assertion whether or not the rows were handed over at the right moment. What is checked instead is
/// the poll that *meets* the pending input - it has to return rows, and the input has to have been
/// read only once by then.
#[tokio::test]
async fn a_pending_input_does_not_hold_the_rows_a_frame_buffered() {
let schema = frame_schema();
let batches = one_series_frames();
let handed = Arc::new(AtomicUsize::new(0));
let staged_schema = Arc::clone(&schema);
let staged_handed = Arc::clone(&handed);
let source = staged(Arc::clone(&schema), move || {
Box::pin(
StagedStream::new(Arc::clone(&staged_schema), batches.clone())
// The third read is the one that reports pending. By then the first two chunks have
// filled an evaluation window, so the encoder is holding rows: the boundary is met
// with something buffered, which is the case worth checking.
.pending_after(2)
.counted(&staged_handed),
)
});
let tree = rate_tree();
let node = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(schema)],
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.expect("a rate tree is evaluated on a grid");
let exec = plan_of(node, source).expect("the runtime executes the tree");
let declared = exec.schema();
let mut stream = exec
.execute(0, SessionContext::default().task_ctx())
.expect("the node executes");
// Poll by hand so the boundary is observable: `await` runs through the pending poll, and every
// later assertion would then hold whether or not the rows came out at the right moment.
//
// The pending poll is the one under test. `StagedStream` wakes itself straight after reporting
// pending, so a stream that answered nothing there would still be collected correctly; the only
// way to see the difference is to look at what that single poll returned. Two chunks have been
// read by then and an evaluation window is full, so the encoder is holding rows: it must answer
// them, not pending.
let next = stream.next();
futures::pin_mut!(next);
let mut answered = match futures::poll!(next) {
std::task::Poll::Ready(Some(Ok(batch))) => vec![batch],
std::task::Poll::Ready(Some(Err(error))) => panic!("the frame answered {error}"),
std::task::Poll::Ready(None) => panic!("the frame ended before answering a row"),
std::task::Poll::Pending => panic!(
"the pending input held the rows the frame already encoded: the input handed over {} \
batches",
handed.load(Ordering::SeqCst)
),
};
assert_eq!(
declared.as_ref(),
answered[0].schema().as_ref(),
"a batch is declared as the plan declared it"
);
answered.extend(
stream
.try_collect::<Vec<_>>()
.await
.expect("the stream resumes where the pending poll left it"),
);
let stream = Box::pin(RecordBatchStreamAdapter::new(
declared,
futures::stream::iter(answered.into_iter().map(Ok)),
));
let decoded = unpack_evaluated_series(stream, None)
.expect("the node declares the frame it yields")
.try_collect::<Vec<_>>()
.await
.expect("every batch of the frame decodes");
assert_eq!(
1,
decoded.iter().filter(|(_, last)| *last).count(),
"one series closes exactly once, across the pending boundary"
);
// What the row table answers for the same input is the reference: `rate` decides which grid
// points carry a value, so computing the expected timestamps here would guess at it.
let plain =
plan(rate_tree(), memory_exec(one_series_frames())).expect("the runtime executes a tree");
let mut expected = rows(
&collect(plain, SessionContext::default().task_ctx())
.await
.expect("the row table runs"),
)
.into_iter()
.map(|(time, _, _, _)| time)
.collect::<Vec<_>>();
expected.sort_unstable();
let mut times = Vec::new();
for (series, _) in &decoded {
for timestamp in series.time.timestamps().expect("a grid axis").iter() {
times.push(*timestamp);
}
}
times.sort_unstable();
assert_eq!(
expected, times,
"every row the row table answers is in the frame, none twice"
);
}
/// A closed chain that holds no value still reaches the consumer of a frame.
///
/// The last sample sits past the grid, so no evaluation point has a sample in its window and the
/// window closes the series with a chunk that carries no position. Dropping it would lose the only
/// signal that says this series is done - the row table has no row for it at all, which is exactly
/// why the frame has to keep it.
#[tokio::test]
async fn a_frame_carries_the_closing_chunk_a_row_table_drops() {
let schema = raw_schema();
// Every sample is newer than the last grid point, so `rate` decides nothing at any of them.
let late = [(GRID_START + RANGE_MS + SAMPLE_STEP, 1.0)];
let chunk = chunk_of(
&schema,
"late",
metric(&[(HOST, "late")], Some(METRIC)),
&late.iter().map(|(time, _)| *time).collect::<Vec<_>>(),
vec![late.iter().map(|(_, value)| Some(*value)).collect()],
true,
);
let tree = rate_tree();
let node = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.expect("a rate tree is evaluated on a grid");
let exec = plan_of(node, memory_exec(vec![frame_of(&schema, &chunk)]))
.expect("the runtime executes the tree");
let declared = exec.schema();
let batches = collect(exec, SessionContext::default().task_ctx())
.await
.expect("the frame runs");
assert!(
batches
.iter()
.all(|batch| batch.schema().as_ref() == declared.as_ref()),
"every batch carries the declared schema"
);
let stream = Box::pin(RecordBatchStreamAdapter::new(
declared,
futures::stream::iter(batches.into_iter().map(Ok)),
));
let decoded = unpack_evaluated_series(stream, None)
.expect("the node declares the frame it yields")
.try_collect::<Vec<_>>()
.await
.expect("the frame decodes");
assert_eq!(
1,
decoded.len(),
"the closing chunk is one frame row, not a dropped chunk: {decoded:?}"
);
let (series, is_last) = &decoded[0];
assert!(is_last, "that row is the one that closes the chain");
assert!(
series.time.is_empty(),
"it carries no grid position: {:?}",
series.time.positions().collect::<Vec<_>>()
);
// The row table answers nothing for the same input, which is what makes this chunk worth keeping.
let plain = plan(rate_tree(), memory_exec(vec![frame_of(&schema, &chunk)]))
.expect("the runtime executes a tree");
let rows = rows(
&collect(plain, SessionContext::default().task_ctx())
.await
.expect("the row table runs"),
);
assert!(
rows.is_empty(),
"the row table drops the closing chunk entirely: {rows:?}"
);
}
/// Two chains that close on one label set are both emitted, each closing itself.
///
/// A range function drops the metric name, so two different raw series can reach the output under the
/// same complete labels. A frame carries one chain per chunk and the later one is a further
/// contribution, not a repeat of the first: refusing it as a duplicate identity would drop values the
/// query answers with.
#[tokio::test]
async fn two_chains_that_close_on_one_label_set_each_keep_their_frame() {
// `rate` drops `__name__`, and neither series declares any other tag, so both collapse to the
// same empty label set.
let schema = raw_schema();
let chunk_of_host = |host: &str, scale: f64| {
let samples = series_samples(scale);
chunk_of(
&schema,
host,
metric(&[], None),
&samples.iter().map(|(time, _)| *time).collect::<Vec<_>>(),
vec![samples.iter().map(|(_, value)| Some(*value)).collect()],
true,
)
};
let partitions = vec![
vec![frame_of(&schema, &chunk_of_host("web-1", 1.0))],
vec![frame_of(&schema, &chunk_of_host("web-2", 2.0))],
];
let tree = rate_tree();
let node = FusedSeries::try_new_with_format(
Arc::from(PROMQL),
tree.clone(),
binding(&tree),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
FusedOutputFormat::SeriesFrame,
)
.expect("a rate tree is evaluated on a grid");
let exec =
plan_of(node, memory_exec_partitioned(partitions)).expect("the runtime executes the tree");
assert_eq!(
2,
exec.properties().partitioning.partition_count(),
"the frame keeps the input's partition count"
);
let declared = exec.schema();
let batches = collect(exec, SessionContext::default().task_ctx())
.await
.expect("both chains are legal");
let decoded = unpack_evaluated_series(
Box::pin(RecordBatchStreamAdapter::new(
declared,
futures::stream::iter(batches.into_iter().map(Ok)),
)),
None,
)
.expect("the node declares the frame it yields")
.try_collect::<Vec<_>>()
.await
.expect("the frame decodes");
assert_eq!(
2,
decoded.iter().filter(|(_, last)| *last).count(),
"each chain closes once, in its own partition"
);
assert_eq!(
decoded[0].0.metric, decoded[1].0.metric,
"the two chains carry the one label set"
);
let mut per_chain = Vec::new();
for (series, last) in &decoded {
assert!(*last, "each partition's single chain closes exactly once");
assert_eq!(None, series.metric.name(), "rate drops the metric name");
assert!(
series.metric.tags().is_empty(),
"both chains have empty labels"
);
let mut rows = Vec::new();
for (position, row) in series.time.positions() {
let value = match series.fields.as_slice()[0].sample(row) {
datatypes::series::samples::SampleRef::Float(value) => value,
_ => panic!("each rate grid position has a present value"),
};
rows.push((position, value.to_bits()));
}
rows.sort_unstable_by_key(|(position, _)| *position);
per_chain.push(rows);
}
assert_eq!(2, per_chain.len(), "one closing chain from each partition");
per_chain.sort();
for (rows, value) in per_chain.iter().zip([1.0f64, 2.0]) {
assert_eq!(
(0..GRID_LEN as usize)
.map(|i| (i, value.to_bits()))
.collect::<Vec<_>>(),
*rows,
"each chain preserves every exact grid position and value",
);
}
}
src/promql/src/extension_plan/fused_series/tests/frame_output.rs
+412
View File
@@ -0,0 +1,412 @@
// 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.
use super::*;
use crate::map::MathOp;
use crate::map_label::LabelRewrite;
use crate::native::plan::{BoundMap, MapOp};
fn instant_tree(schema: Arc<SeriesSchema>, grid: Arc<EvalGrid>) -> NativeTree {
let output = NativeType::Samples(Arc::new(SeriesSchema {
fields: schema.fields.clone(),
eval_grid: Some(Arc::clone(&grid)),
}));
NativeTree::try_new(NativeNode {
expr: NativeExpr::Instant {
input: Box::new(NativeNode {
expr: NativeExpr::Input(BoundInput {
source_index: 0,
output: NativeType::Samples(Arc::clone(&schema)),
ordering: SeriesOrdering::Contiguous,
}),
output: NativeType::Samples(Arc::clone(&schema)),
}),
op: BoundInstant {
lookback_ms: 0,
offset_ms: 0,
at_ms: None,
timestamp: false,
output_grid: grid,
},
},
output,
})
.expect("instant selector tree is valid")
}
async fn frame_decode(
exec: Arc<dyn ExecutionPlan>,
) -> Vec<(datatypes::series::samples::Series, bool)> {
let batches = collect(Arc::clone(&exec), SessionContext::default().task_ctx())
.await
.expect("frame output executes");
let stream = Box::pin(RecordBatchStreamAdapter::new(
exec.schema(),
futures::stream::iter(batches.into_iter().map(Ok)),
));
unpack_evaluated_series(stream, None)
.expect("evaluated frame schema decodes")
.try_collect::<Vec<_>>()
.await
.expect("evaluated frame decodes")
}
fn compact_float_rows(
decoded: &[(datatypes::series::samples::Series, bool)],
) -> Vec<(i64, Vec<Option<u64>>)> {
let mut rows = Vec::new();
for (series, _) in decoded {
let timestamps = series.time.timestamps().expect("grid timestamps");
for (row, timestamp) in timestamps.iter().enumerate() {
let values = series
.fields
.as_slice()
.iter()
.map(
|field| match field.iter().nth(row).expect("sample at timestamp") {
datatypes::series::samples::SampleRef::Float(value) => {
Some(value.to_bits())
}
datatypes::series::samples::SampleRef::Missing => None,
_ => panic!("float or null sample expected"),
},
)
.collect();
rows.push((*timestamp, values));
}
}
rows.sort_by_key(|(timestamp, _)| *timestamp);
rows
}
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,
),
]))
}
fn plain_float_rows(batches: &[RecordBatch], fields: &[&str]) -> Vec<(i64, Vec<Option<u64>>)> {
let mut rows = Vec::new();
for batch in batches {
let time = batch
.column_by_name(TIME)
.unwrap()
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap();
let columns = fields
.iter()
.map(|field| {
batch
.column_by_name(field)
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap()
})
.collect::<Vec<_>>();
for row in 0..batch.num_rows() {
rows.push((
time.value(row),
columns
.iter()
.map(|column| column.is_valid(row).then(|| column.value(row).to_bits()))
.collect(),
));
}
}
rows.sort_by_key(|(timestamp, _)| *timestamp);
rows
}
#[tokio::test]
async fn sparse_instant_frame_preserves_positions_nulls_and_nan_across_formats() {
let schema = two_field_schema();
let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1_000, 6).unwrap()).unwrap());
let tree = instant_tree(Arc::clone(&schema), grid);
let chunk = chunk_of(
&schema,
"sparse",
metric(&[(HOST, "sparse")], Some(METRIC)),
&[0, 2_000, 5_000],
vec![
vec![Some(1.0), None, Some(f64::from_bits(0x7ff8_0000_0000_0042))],
vec![None, Some(2.0), None],
],
true,
);
let batches = vec![frame_of(&schema, &chunk)];
let binding = binding_of(&tree, Arc::clone(&schema), vec![HOST.into(), NAME.into()]);
let make_node = |format| {
FusedSeries::try_new_with_format(
Arc::from("instant(x)"),
tree.clone(),
binding.clone(),
vec![input_plan(frame_schema_of(&schema))],
two_field_output_schema().to_dfschema_ref().unwrap(),
format,
)
.expect("instant tree has the declared frame output")
};
let frame_exec = plan_of(
make_node(FusedOutputFormat::SeriesFrame),
memory_exec_with(frame_schema_of(&schema), batches.clone()),
)
.unwrap();
let decoded = frame_decode(frame_exec).await;
assert_eq!(1, decoded.iter().filter(|(_, last)| *last).count());
let expected = vec![
(0, vec![Some(1.0f64.to_bits()), None]),
(2_000, vec![None, Some(2.0f64.to_bits())]),
(5_000, vec![Some(0x7ff8_0000_0000_0042), None]),
];
assert_eq!(expected, compact_float_rows(&decoded));
let plain_exec = plan_of(
make_node(FusedOutputFormat::PlainRows),
memory_exec_with(frame_schema_of(&schema), batches),
)
.unwrap();
let plain_batches = collect(plain_exec, SessionContext::default().task_ctx())
.await
.unwrap();
assert_eq!(expected, plain_float_rows(&plain_batches, &[VALUE, OTHER]));
assert!(decoded.iter().all(|(series, _)| {
series.metric.get(HOST) == Some("sparse") && series.metric.name() == Some(METRIC)
}));
}
#[tokio::test]
async fn abs_rate_frame_matches_plain_rows_with_exact_values_and_labels() {
let rate = rate_tree();
let tree = NativeTree::try_new(NativeNode {
expr: NativeExpr::Map {
input: Box::new(rate.root().clone()),
op: BoundMap {
ops: Arc::from([MapOp::Math(MathOp::Abs)]),
},
},
output: rate.root().output.clone(),
})
.unwrap();
let batches = two_series_frames();
let binding = binding(&tree);
let make_node = |format| {
FusedSeries::try_new_with_format(
Arc::from("abs(rate(x[5m]))"),
tree.clone(),
binding.clone(),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
format,
)
.unwrap()
};
let plain_exec = plan_of(
make_node(FusedOutputFormat::PlainRows),
memory_exec(batches.clone()),
)
.unwrap();
let plain_batches = collect(plain_exec, SessionContext::default().task_ctx())
.await
.unwrap();
let mut expected = rows(&plain_batches)
.into_iter()
.map(|(time, value, host, name)| (time, value.map(f64::to_bits), host, name))
.collect::<Vec<_>>();
expected.sort();
let mut exact = Vec::new();
for (host, scale) in [("web-1", 1.0_f64), ("web-2", 2.0)] {
for position in 0..GRID_LEN {
exact.push((
GRID_START + position as i64 * GRID_STEP,
Some(scale.to_bits()),
Some(host.to_string()),
None,
));
}
}
exact.sort();
assert_eq!(exact, expected);
let frame_exec = plan_of(
make_node(FusedOutputFormat::SeriesFrame),
memory_exec(batches),
)
.unwrap();
let decoded = frame_decode(frame_exec).await;
assert_eq!(2, decoded.iter().filter(|(_, last)| *last).count());
let mut actual = Vec::new();
for (series, _) in &decoded {
assert_eq!(None, series.metric.name());
let host = series.metric.get(HOST).unwrap().to_string();
for (time, values) in compact_float_rows(&[(series.clone(), true)]) {
actual.push((time, values[0], Some(host.clone()), None));
}
}
actual.sort();
assert_eq!(exact, actual);
}
#[tokio::test]
async fn rewritten_identical_labels_keep_disjoint_chains_in_one_encoder() {
let schema = raw_schema();
let semantic_grid =
Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1_000, 6).unwrap()).unwrap());
let base = instant_tree(Arc::clone(&schema), semantic_grid);
let tree = NativeTree::try_new(NativeNode {
expr: NativeExpr::Map {
input: Box::new(base.root().clone()),
op: BoundMap {
ops: Arc::from([
MapOp::Label(Arc::new(
LabelRewrite::replace("host", "shared", "host", ".*").unwrap(),
)),
MapOp::Label(Arc::new(
LabelRewrite::replace("zone", "edge", "host", ".*").unwrap(),
)),
]),
},
},
output: base.root().output.clone(),
})
.unwrap();
let chunks = vec![
chunk_of(
&schema,
"a",
metric(&[(HOST, "a")], Some(METRIC)),
&[0, 2_000],
vec![vec![Some(10.0), Some(12.0)]],
true,
),
chunk_of(
&schema,
"b",
metric(&[(HOST, "b")], Some(METRIC)),
&[1_000, 5_000],
vec![vec![Some(21.0), Some(25.0)]],
true,
),
];
let batches = frames_of(&schema, &chunks);
let binding = binding_of(&tree, Arc::clone(&schema), vec![HOST.into(), NAME.into()]);
let make_node = |format| {
FusedSeries::try_new_with_format(
Arc::from("rewritten instant"),
tree.clone(),
binding.clone(),
vec![input_plan(frame_schema())],
output_schema().to_dfschema_ref().unwrap(),
format,
)
.unwrap()
};
let exec = plan_of(
make_node(FusedOutputFormat::SeriesFrame),
memory_exec(batches.clone()),
)
.unwrap();
let decoded = frame_decode(exec).await;
assert_eq!(
2,
decoded.iter().filter(|(_, last)| *last).count(),
"each chain closes once"
);
let mut groups = Vec::new();
let mut group_rows = Vec::new();
let mut closes = 0;
for (series, last) in &decoded {
assert_eq!(Some("shared"), series.metric.get(HOST));
assert_eq!(Some("edge"), series.metric.get("zone"));
assert_eq!(Some(METRIC), series.metric.name());
let timestamps = series.time.timestamps().unwrap();
let values = series.fields.get(0).unwrap();
for (i, time) in timestamps.iter().enumerate() {
let value = match values.iter().nth(i).unwrap() {
datatypes::series::samples::SampleRef::Float(value) => value.to_bits(),
_ => panic!("float sample expected"),
};
group_rows.push((*time, value));
}
if *last {
closes += 1;
group_rows.sort_by_key(|(time, _)| *time);
groups.push(std::mem::take(&mut group_rows));
}
}
assert_eq!(2, closes, "each chain closes exactly once");
assert!(group_rows.is_empty(), "all output chains close");
groups.sort();
assert_eq!(
vec![
vec![(0, 10.0f64.to_bits()), (2_000, 12.0f64.to_bits())],
vec![(1_000, 21.0f64.to_bits()), (5_000, 25.0f64.to_bits())],
],
groups
);
let plain_exec = plan_of(
make_node(FusedOutputFormat::PlainRows),
memory_exec_with(frame_schema(), batches),
)
.unwrap();
let plain = collect(plain_exec, SessionContext::default().task_ctx())
.await
.unwrap();
let mut plain = rows(&plain)
.into_iter()
.map(|(time, value, host, name)| (time, value.map(f64::to_bits), host, name))
.collect::<Vec<_>>();
plain.sort();
assert_eq!(
vec![
(
0,
Some(10.0f64.to_bits()),
Some("shared".into()),
Some(METRIC.into())
),
(
1_000,
Some(21.0f64.to_bits()),
Some("shared".into()),
Some(METRIC.into())
),
(
2_000,
Some(12.0f64.to_bits()),
Some("shared".into()),
Some(METRIC.into())
),
(
5_000,
Some(25.0f64.to_bits()),
Some("shared".into()),
Some(METRIC.into())
),
],
plain
);
}
tests-integration/Cargo.toml
+1
View File
@@ -109,6 +109,7 @@ otel-arrow-rust.workspace = true
paste.workspace = true
pipeline.workspace = true
plugins.workspace = true
promql.workspace = true
prost.workspace = true
rand.workspace = true
rskafka.workspace = true
tests-integration/tests/http.rs
+212
View File
@@ -46,15 +46,20 @@ use common_frontend::slow_query_event::{
};
use common_memory_manager::OnExhaustedPolicy;
use common_options::plugin_options::StandaloneFlag;
use datafusion::common::{DFSchema, ToDFSchema};
use datafusion::logical_expr::{EmptyRelation, LogicalPlan};
use datafusion::physical_plan::metrics::MetricValue;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
use datatypes::arrow::array::{
Array as _, Float64Array, StringArray, StringViewArray, TimestampMillisecondArray,
};
use datatypes::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
use datatypes::series::samples::{SampleRef, SeriesSchema};
use datatypes::series::time::{EvalGrid, GridSpec};
use flate2::Compression;
use flate2::write::GzEncoder;
use frontend::instance::Instance;
use futures::TryStreamExt;
use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter};
use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter};
use loki_proto::prost_types::Timestamp;
@@ -1855,6 +1860,211 @@ fn fixture_rows(value: impl Fn(f64) -> f64) -> Vec<PromqlRow> {
.collect()
}
/// Replays the captured raw scan through both public FusedSeries output formats.
async fn assert_raw_scan_fused_formats(observer: &NativeFusedObserver, frontend: &Arc<Instance>) {
let table_id = frontend
.catalog_manager()
.table("greptime", "public", "native_demo", None)
.await
.unwrap()
.expect("fixture table exists")
.table_info()
.table_id();
let (_, route) = frontend
.table_metadata_manager()
.table_route_manager()
.get_physical_table_route(table_id)
.await
.unwrap();
assert_eq!(
route.region_routes.len(),
1,
"fixture has exactly one physical region"
);
observer.clear();
let (_, row_schema, _) = promql_output(frontend, "native_demo", true, "0", "2").await;
let origin = observer.take("native_demo");
let children = origin.children();
assert_eq!(children.len(), 1, "the fused origin has exactly one child");
let mut scan_plan = children[0].clone();
if scan_plan
.downcast_ref::<table::table::scan::SeriesScanExec>()
.is_none()
{
let wrapper_children = scan_plan.children();
assert_eq!(
wrapper_children.len(),
1,
"the optional scan wrapper has one child"
);
scan_plan = wrapper_children[0].clone();
}
let scan = scan_plan
.downcast_ref::<table::table::scan::SeriesScanExec>()
.expect("the fused input is a real SeriesScanExec");
assert_eq!(scan.properties().output_partitioning().partition_count(), 4);
let raw_schema = Arc::clone(scan.series_schema());
let label_source = scan.label_source().clone();
let ordering = scan.raw_series_ordering();
let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1_000, 3).unwrap()).unwrap());
let output_schema = Arc::new(SeriesSchema {
fields: Arc::clone(&raw_schema.fields),
eval_grid: Some(Arc::clone(&grid)),
});
let raw_type = promql::native::plan::NativeType::Samples(Arc::clone(&raw_schema));
let root = promql::native::plan::NativeNode {
expr: promql::native::plan::NativeExpr::Instant {
input: Box::new(promql::native::plan::NativeNode {
expr: promql::native::plan::NativeExpr::Input(promql::native::plan::BoundInput {
source_index: 0,
output: raw_type.clone(),
ordering,
}),
output: raw_type,
}),
op: promql::native::plan::BoundInstant {
lookback_ms: 300_000,
offset_ms: 0,
at_ms: None,
timestamp: false,
output_grid: Arc::clone(&grid),
},
},
output: promql::native::plan::NativeType::Samples(Arc::clone(&output_schema)),
};
let tree = promql::native::plan::NativeTree::try_new(root).unwrap();
let binding = promql::extension_plan::FusedSeriesBinding {
grid,
metric_name: Some(Arc::from("native_demo")),
time_index: row_schema
.fields()
.iter()
.find(|field| {
matches!(
field.data_type(),
DataType::Timestamp(TimeUnit::Millisecond, _)
)
})
.expect("PromQL result has a millisecond timestamp field")
.name()
.clone(),
tags: vec!["host".to_string()],
};
let logical_child = LogicalPlan::EmptyRelation(EmptyRelation {
produce_one_row: false,
schema: DFSchema::try_from(scan.schema()).unwrap().into(),
});
let expected = fixture_rows(|value| value);
for format in [
promql::extension_plan::FusedOutputFormat::PlainRows,
promql::extension_plan::FusedOutputFormat::SeriesFrame,
] {
scan_plan.clone().reset_state().unwrap();
let exec = promql::extension_plan::FusedSeries::try_new_with_format(
Arc::from("native_demo"),
tree.clone(),
binding.clone(),
vec![logical_child.clone()],
row_schema.clone().to_dfschema_ref().unwrap(),
format,
)
.unwrap()
.to_execution_plan(vec![(scan_plan.clone(), label_source.clone(), ordering)])
.unwrap();
let context = "standalone raw SeriesScanExec replay";
let mut actual = Vec::new();
let mut closed = 0;
let mut partition_counts = Vec::new();
for partition in 0..4 {
let stream = exec
.execute(
partition,
datafusion::prelude::SessionContext::new().task_ctx(),
)
.unwrap();
if format == promql::extension_plan::FusedOutputFormat::SeriesFrame {
let series = common_query::series_frame::unpack_evaluated_series(stream, None)
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert!(
!series.is_empty(),
"partition {partition} emitted no series frames"
);
partition_counts.push(
series
.iter()
.map(|(s, _)| s.time.timestamps().unwrap().len())
.sum::<usize>(),
);
for (series, last) in series {
assert_eq!(series.metric.name(), Some("native_demo"));
assert_eq!(series.metric.tags().len(), 1);
assert_eq!(series.fields.len(), 1);
let host = series.metric.get("host").unwrap().to_string();
let timestamps = series.time.timestamps().unwrap();
let samples = series.fields.get(0).unwrap();
for (i, timestamp) in timestamps.iter().enumerate() {
let SampleRef::Float(value) = samples.iter().nth(i).unwrap() else {
panic!("expected float sample")
};
actual.push((vec![("host".to_string(), host.clone())], *timestamp, value));
}
closed += usize::from(last);
}
} else {
let batches = stream.try_collect::<Vec<_>>().await.unwrap();
let mut count = 0;
for batch in batches {
let ts = batch
.column_by_name(&binding.time_index)
.unwrap()
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.unwrap();
let host = batch
.column_by_name("host")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let values = batch
.column_by_name("val")
.unwrap()
.as_any()
.downcast_ref::<Float64Array>()
.unwrap();
for i in 0..batch.num_rows() {
actual.push((
vec![("host".to_string(), host.value(i).to_string())],
ts.value(i),
values.value(i),
));
count += 1;
}
}
assert_eq!(row_schema.fields().len(), 3);
assert!(count > 0, "partition {partition} emitted no rows");
partition_counts.push(count);
}
}
assert_eq!(partition_counts.len(), 4);
assert!(
partition_counts.iter().all(|count| *count > 0),
"all four partitions emitted output: {partition_counts:?}"
);
assert_eq!(partition_counts.iter().sum::<usize>(), 24);
if format == promql::extension_plan::FusedOutputFormat::SeriesFrame {
assert_eq!(closed, 8, "one closing frame per host series");
}
actual.sort_unstable_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
assert_rows_eq(&actual, &expected, context);
assert_native_partition_output(&exec, Some(4), 24, context);
}
}
/// Runs one PromQL query through the public frontend handler and returns the plan it executed, the
/// schema it answered with and every row it reported.
async fn promql_output(
@@ -2206,6 +2416,8 @@ pub async fn test_native_promql_route_plans_the_query(store_type: StorageType) {
.await;
}
assert_raw_scan_fused_formats(&observer, frontend).await;
for (query, start, end, expected, scan_partitions) in [
(
"1 + 2",