From 4beb5c7490371dd0763a047cf3a76548340c5a6b Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 20:45:01 +0800 Subject: [PATCH 01/38] feat(query): offer a PromQL statement to the native fused-series route Nothing called `native_compile::compile` or built a `FusedSeries`, so the native tree and its node were library code with tests and no query behind them. This adds the route, behind a per-query flag. The statement is offered to the compiler at the root, where the expression the reader wrote still exists. A shape the compiler refuses - `NotImplemented` - is this route declining the query, and planning continues exactly as it would with the flag off. Anything else is a real error. The physical half is an arm in `PromExtensionPlanner`. Its children are already built by then, so a metric scan arrives as the region scan, and the node asks that scan for its named raw-series binding: `try_to_named_series_scan` resolves the label source and ordering from the scanner itself, which is the only thing that can name them, and returns the leaf the node reads. A scanner that cannot serve a raw frame is reported, not silently re-planned. Two things this needed that the route does not get for free: - The region scan has to be prepared as one partition. `resolve_named_series_binding` refuses a scan whose partition count is not one, and a fresh scan has not set its partitions at all: `ParallelizeScan` runs after extension planning and would spread it, so the node would run one evaluator per partition and split a logical series. Preparing the ranges as one partition is what keeps `FusedSeriesExec` over one series stream, and it is why this route does not scan in parallel yet. - The scan needs the per-series distribution hint. That is what makes a region prepare its series scan, and only that path packs a series' chunks into the frame this node reads. The existing planner asks for it through a sort that `ScanHintRule` turns into the hint; this node has no such sort below it, so the hint is attached to the scan directly. A query the route cannot bind faithfully keeps the existing plan rather than a wrong one: more than one selector, a label matcher that would survive preprocessing, an `__field__` matcher, a native histogram field, a missing table, or a shape whose grid the compiler's statement cannot express. Signed-off-by: discord9 --- Cargo.lock | 1 + src/promql/Cargo.toml | 1 + src/promql/src/extension_plan/planner.rs | 103 +++++- src/query/src/planner.rs | 1 + src/query/src/promql/native_compile.rs | 6 +- src/query/src/promql/planner.rs | 30 +- src/query/src/promql/planner/native.rs | 353 +++++++++++++++++++++ src/query/src/promql/planner/test.rs | 87 ++++- src/query/src/promql/planner/test/delta.rs | 5 + 9 files changed, 581 insertions(+), 6 deletions(-) create mode 100644 src/query/src/promql/planner/native.rs diff --git a/Cargo.lock b/Cargo.lock index 2ef45f7114a..606322e47b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11378,6 +11378,7 @@ dependencies = [ "regex", "snafu 0.8.6", "store-api", + "table", "tokio", "twox-hash", ] diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml index a3d8f5c4e71..31e4d08a9fb 100644 --- a/src/promql/Cargo.toml +++ b/src/promql/Cargo.toml @@ -28,6 +28,7 @@ prost.workspace = true regex.workspace = true snafu.workspace = true store-api.workspace = true +table.workspace = true twox-hash.workspace = true [dev-dependencies] diff --git a/src/promql/src/extension_plan/planner.rs b/src/promql/src/extension_plan/planner.rs index 50cf4543d59..f2d86b76862 100644 --- a/src/promql/src/extension_plan/planner.rs +++ b/src/promql/src/extension_plan/planner.rs @@ -16,15 +16,16 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::catalog::Session; -use datafusion::error::Result as DfResult; +use datafusion::error::{DataFusionError, Result as DfResult}; use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; +use table::table::scan::RegionScanExec; use crate::extension_plan::{ - Absent, EmptyMetric, HistogramFold, InstantManipulate, RangeManipulate, ScalarCalculate, - SeriesDivide, SeriesNormalize, UnionDistinctOn, + Absent, EmptyMetric, FusedSeries, HistogramFold, InstantManipulate, RangeManipulate, + ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, }; pub struct PromExtensionPlanner; @@ -65,8 +66,104 @@ impl ExtensionPlanner for PromExtensionPlanner { ))) } else if let Some(node) = node.as_any().downcast_ref::() { Ok(Some(node.to_execution_plan(physical_inputs[0].clone()))) + } else if let Some(node) = node.as_any().downcast_ref::() { + plan_fused_series(node, physical_inputs) } else { Ok(None) } } } + +/// The label `__name__` is reported under, which is a column of this node's output but not a tag +/// column of any region. +const METRIC_NAME: &str = "__name__"; + +/// Builds the physical node of a [`FusedSeries`] over the raw series leaf its input can provide. +/// +/// The node reads a raw series frame, so what it needs below it is the region scan itself: the +/// scan's own raw-series entrypoint is what packs a series' chunks contiguously and declares their +/// ordering. Anything else in between - a filter, a projection, an exchange - produces a different +/// frame, or no frame at all, and this layer cannot look under it. +/// +/// `Ok(None)` is therefore the answer for a child that is not a region scan: the logical plan is +/// already fixed and DataFusion reports an unplanned extension node normally. A region scan that is +/// found but refuses the raw path is a different case — the logical route chose this node for this +/// scan and nothing here can undo that choice, because the logical plan is not available — so it is +/// reported as an error that names the reason instead of silently producing no plan. +fn plan_fused_series( + fused: &FusedSeries, + physical_inputs: &[Arc], +) -> DfResult>> { + let Some(input) = physical_inputs.first() else { + return Ok(None); + }; + let Some(scan) = input.downcast_ref::() else { + return Ok(None); + }; + + let binding = fused.binding(); + let Some(metric_name) = binding.metric_name.as_deref() else { + return Err(DataFusionError::Plan( + "the FusedSeries node binds no metric name, so the region scan below it cannot be \ + asked for the series of one metric" + .to_string(), + )); + }; + // `__name__` is the metric name rather than a column of the region, so it is not one of the + // logical tags the scan is asked to decode; the frame names each series' own metric. + let logical_tags = binding + .tags + .iter() + .filter(|tag| tag.as_str() != METRIC_NAME) + .cloned() + .collect::>(); + let field_names = binding + .series_schema + .fields + .iter() + .map(|field| field.name.to_string()) + .collect::>(); + + let scan = single_partition_scan(scan)?; + // The scan resolves the label source and the ordering it will declare itself, against its own + // metadata snapshot, and hands both back through the leaf it builds. Naming them here instead + // would be a second copy of a binding only the scanner can make: the physical metadata and + // column ids a `RawLabelSource::PrimaryKey` carries are not reachable from a region scan's + // public surface, so the layer that resolves them is the scanner. + let Some(series_scan) = + scan.try_to_named_series_scan(metric_name, &logical_tags, &field_names)? + else { + return Err(DataFusionError::Plan(format!( + "the region scan under FusedSeries could not serve a raw series frame for metric \ + {metric_name:?}: the scanner refused the raw series path for this scan, so this node \ + has no input to read" + ))); + }; + let label_source = series_scan.label_source().clone(); + let ordering = series_scan.raw_series_ordering(); + + fused + .to_execution_plan(series_scan, label_source, ordering) + .map(Some) +} + +/// The scan as one partition holding every range it already has. +/// +/// A raw-series binding is resolvable only once the scan holds exactly one partition: the label +/// source is decoded against the scanner's own metadata snapshot and the frames it produces carry +/// whole series, so a scan spread over several partitions has no single binding to name. The rule +/// that assigns partition ranges (`ParallelizeScan`) runs *after* physical planning, i.e. after +/// this node is built, so a scan reaching here is unprepared; preparing it here keeps the node from +/// reading one partition of the series it was bound to. A scan that is already prepared — several +/// partitions, assigned by whatever ran before — is passed through as it is, and the binding below +/// declines it if it cannot name one partition. +fn single_partition_scan(scan: &RegionScanExec) -> DfResult> { + if scan.is_partition_set() { + return Ok(Arc::new(scan.clone())); + } + let ranges = scan.get_partition_ranges(); + let prepared = scan + .with_new_partitions(vec![ranges], 1) + .map_err(|error| DataFusionError::External(error.into_inner()))?; + Ok(Arc::new(prepared)) +} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 44cd1653cee..861f923d74b 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -345,6 +345,7 @@ impl DfLogicalPlanner { stmt, &self.engine_state, promql_annotations, + crate::promql::planner::native::native_promql_requested(&query_ctx), ) .await .map_err(BoxedError::new) diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index 5019f4b7a86..001660b9dc5 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -1101,7 +1101,11 @@ fn instant_selection(context: &Context, selector: &VectorSelector) -> Result Result> { +/// +/// Visible to the planner because the binding a [`NativeTree`] is planned against has to name the +/// same grid the tree itself was compiled on, and [`compile`]'s callers therefore read it here +/// rather than deriving the grid a second time. +pub(crate) fn statement_grid(stmt: &EvalStmt) -> Result> { let start_ms = signed_millis_since_epoch(stmt.start)?; let end_ms = signed_millis_since_epoch(stmt.end)?; // The step is narrowed before [`GridSpec`] sees it, so an interval to which no millisecond diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index afe52c7b560..5df731d6a95 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -16,6 +16,7 @@ mod at_modifier; mod function_plans; mod island; mod matching_filters; +pub(crate) mod native; mod set_operator; @@ -292,6 +293,17 @@ pub struct PromPlanner { ctx: PromPlannerContext, /// Optional collector passed to native histogram UDFs. promql_annotations: Option, + /// Whether a statement is first offered to the native fused-series compiler. + /// + /// Read per query from the context extension + /// [`NATIVE_PROMQL_EXTENSION_KEY`](super::planner::native::NATIVE_PROMQL_EXTENSION_KEY), and + /// `false` when the query says nothing: with it off, planning takes the existing path unchanged + /// and nothing below is reached. + /// + /// The engine-wide `QueryOptions::experimental_native_promql` is not the source. It is consumed + /// while `QueryEngineState` is built and not kept, so the planner cannot read it; the extension + /// is what a client sets for one request. + native_fused: bool, } type BinaryFieldPair<'a> = (&'a String, &'a String); @@ -302,22 +314,38 @@ impl PromPlanner { stmt: &EvalStmt, query_engine_state: &QueryEngineState, ) -> Result { - Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None).await + Self::stmt_to_plan_with_annotations(table_provider, stmt, query_engine_state, None, false) + .await } /// Plans a PromQL statement and passes the optional collector to histogram UDFs. + /// + /// `native_fused` enables the native fused-series route: the statement is offered to + /// [`compile`](crate::promql::native_compile::compile) first, and whatever that compiler does + /// not accept is planned by the existing path, unchanged. See + /// [`Self::try_plan_native_fused`]. pub async fn stmt_to_plan_with_annotations( table_provider: DfTableSourceProvider, stmt: &EvalStmt, query_engine_state: &QueryEngineState, promql_annotations: Option, + native_fused: bool, ) -> Result { let mut planner = Self { table_provider, ctx: PromPlannerContext::from_eval_stmt(stmt)?, promql_annotations, + native_fused, }; + // The whole statement is offered to the native compiler first, because that compiler folds + // a range function and the math functions over it into one tree: offered any later, the + // expression has already been split into the nodes those functions would have shared. A + // shape it does not accept is planned by the existing path exactly as if it were not here. + if let Some(plan) = planner.try_plan_native_fused(&stmt.expr).await? { + return planner.strip_tsid_column(plan); + } + let plan = planner .prom_root_expr_to_plan(&stmt.expr, query_engine_state) .await?; diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs new file mode 100644 index 00000000000..21cd4553d15 --- /dev/null +++ b/src/query/src/promql/planner/native.rs @@ -0,0 +1,353 @@ +// 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 native fused-series route of the PromQL planner. +//! +//! [`PromPlanner::try_plan_native_fused`] offers a whole statement expression to the native +//! compiler and, when the compiler accepts it, replaces everything the existing path would plan +//! below the root with one [`FusedSeries`] node over the selector's own table scan. Declining is +//! the normal answer: every shape the compiler does not accept, and every query this route cannot +//! bind to a scan faithfully, keeps the plan the existing planner builds. +//! +//! What the node needs from here is three things the existing planner already resolves: the raw +//! series schema the tree's input was compiled against (the selector's selected field columns), the +//! statement's evaluation grid, and the logical tag names the output carries. The output schema is +//! built to the shape the existing unary pipeline produces — the series fields, then the tags, then +//! the time index — so a node planned here can be consumed by the same code that consumes a +//! `SeriesNormalize` output. + +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; + +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; +use datafusion::common::{DFSchema, DFSchemaRef}; +use datafusion::datasource::DefaultTableSource; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datatypes::series::SeriesOrdering; +use datatypes::series::samples::{FieldSpec, SeriesSchema}; +use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; +use promql_parser::parser::value::ValueType; +use promql_parser::parser::{ + EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector, +}; +use session::context::QueryContextRef; +use snafu::ResultExt; +use store_api::storage::TimeSeriesDistribution; + +use crate::dummy_catalog::DummyTableProvider; +use crate::promql::error::{DataFusionPlanningSnafu, Result}; +use crate::promql::native_compile::{compile, statement_grid}; +use crate::promql::planner::PromPlanner; + +/// The query-context key that turns the native fused-series route on for one query. +/// +/// Named after the engine-wide option it mirrors ([`QueryOptions::experimental_native_promql`]), +/// and read exactly the way the other per-query override in this crate is read +/// ([`remote_dyn_filter_pushdown_enabled_from_extensions`]): a query-context extension, which a +/// client can set per request over the HTTP hints header or as a gRPC context key. The engine-wide +/// option is a query-engine setting that `QueryEngineState` consumes while it is built and does not +/// keep, so it is not readable from the planner's call site; see [`native_promql_requested`]. +/// +/// [`QueryOptions::experimental_native_promql`]: crate::options::QueryOptions::experimental_native_promql +/// [`remote_dyn_filter_pushdown_enabled_from_extensions`]: +/// crate::options::remote_dyn_filter_pushdown_enabled_from_extensions +pub(crate) const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql"; + +/// Whether `query_ctx` asks for the native fused-series route. +/// +/// The route is off unless the query asks for it, so a query that says nothing keeps the plan the +/// existing planner builds. The value is compared literally (`true` in any case) rather than through +/// a shared boolean parser: an unrecognized value leaves the route off, which is the same plan as +/// the flag never having been read, so a typo costs a slow query and never a wrong one. +pub(crate) fn native_promql_requested(query_ctx: &QueryContextRef) -> bool { + query_ctx + .extension(NATIVE_PROMQL_EXTENSION_KEY) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) +} + +/// The ordering this route compiles its trees against. +/// +/// The logical planner cannot read the ordering a scanner will declare — it is decided when the +/// region scanner is prepared, which happens during physical planning — so the route names the +/// strongest ordering it can assume instead of guessing: `KeyThenTime` orders each series' chunks +/// contiguously *and* orders the series by identity, which `Contiguous` (the weaker declaration a +/// two-phase legacy scan makes) also satisfies. Nothing downstream enforces the tree's declared +/// ordering against the frame's own declaration beyond rejecting `Interleaved`, so a scanner that +/// declares `Contiguous` still feeds a tree compiled here. +const NATIVE_ORDERING: SeriesOrdering = SeriesOrdering::KeyThenTime; + +impl PromPlanner { + /// Plans `prom_expr` as one [`FusedSeries`] node, or returns `None` to keep the existing plan. + /// + /// The route is offered a whole expression because the compiler folds a range function and the + /// math functions over it into one tree: offering a sub-expression instead would compile a tree + /// that computes something the reader did not write. Everything else about the node comes from + /// the planner's resolved selector context: the table name, the time index column, the tag + /// columns and (after the `__field__` matcher) the field columns, which are exactly the raw + /// series schema the tree's input has to be compiled against. + /// + /// `None` is not a failure. It says this query cannot take the route *unchanged*, and the + /// caller plans it with the existing planner as if this method did not exist. The two ways that + /// happens are a shape the compiler refuses (`DataFusionError::NotImplemented`, which names + /// what it found) and a query this route cannot bind faithfully — more than one selector, a + /// label matcher the bare scan below the node would not apply, an `__field__` matcher that + /// narrows the field set, a missing table, a native histogram field, or an unusable grid. + pub(super) async fn try_plan_native_fused( + &mut self, + prom_expr: &PromExpr, + ) -> Result> { + if !self.native_fused { + return Ok(None); + } + // Only the statement root is offered: the tree is compiled against the statement's grid, + // and a subquery's expression is evaluated on its own, smaller one. A range-vector root is + // the instant query's raw sample result, which the existing path reports as a matrix rather + // than as one point per grid position. + if prom_expr.value_type() == ValueType::Matrix { + return Ok(None); + } + + // A tree reads exactly one raw input, and the scan below the node is that one selector's, + // so a shape that names no selector or more than one is not offered at all. + let Some(selector) = single_selector(prom_expr) else { + return Ok(None); + }; + + // Resolve the selector the way the existing planner does, and take what is left as the + // rows it selects. The scan below the node is the selector's own scan, unfiltered, so a + // matcher that survives preprocessing selects fewer rows here than the scan reads. + let matchers = self.preprocess_label_matchers(&selector.matchers, &selector.name)?; + if !matchers.matchers.is_empty() { + return Ok(None); + } + // An `__field__` matcher narrows the field columns the existing planner builds its scan + // with; folding every field of the table instead would evaluate a different window. + if self.ctx.field_column_matcher.is_some() { + return Ok(None); + } + // A metric without a table reads as an empty metric, which the existing planner plans. + if self.setup_context().await?.is_some() { + return Ok(None); + } + + let Some(table_name) = self.ctx.table_name.clone() else { + return Ok(None); + }; + let Some(time_index) = self.ctx.time_index_column.clone() else { + return Ok(None); + }; + let fields = self.ctx.field_columns.clone(); + let tags = self.ctx.tag_columns.clone(); + // A series with no field has nothing to fold, and the tree's input would be an empty + // schema, which no window can read. + if fields.is_empty() { + return Ok(None); + } + + // The selector's own table scan, built by the same method the existing path uses: the + // physical planner turns it into the region scan the node reads, and building a scan here + // instead would let the two paths disagree about the columns they read. + let table_ref = self.table_ref()?; + let input = self.create_table_scan_plan(table_ref.clone()).await?; + // The node reads the region scan itself, so the plan below it has to *be* that scan: a + // filter, a projection or a `__tsid` rewrite in between produces a different physical + // child, and the physical layer cannot reach through it. Declining here, while the logical + // plan is still this method's to choose, is what keeps such a query on the existing path. + let LogicalPlan::TableScan(_) = &input else { + return Ok(None); + }; + // A window over a native histogram yields a histogram, which the output schema below + // cannot declare; the existing path evaluates those with its own kernels. + if Self::field_columns_contain_native_histogram(input.schema(), &fields) { + return Ok(None); + } + // A raw series frame is produced by the region's series scan, which a region scan only + // prepares when the request asks for the per-series distribution. Without the hint the + // region scans by row and has no raw frame to hand over, so the node would be planned with + // an input that cannot feed it. + let input = with_per_series_distribution(input); + + let Some(stmt) = self.native_statement(prom_expr) else { + return Ok(None); + }; + let Ok(grid) = statement_grid(&stmt) else { + // The statement's grid is not one a tree can be evaluated on, so there is no grid to + // address the node's output by. + return Ok(None); + }; + + let series_schema = Arc::new(SeriesSchema { + fields: fields + .iter() + .map(|name| FieldSpec { + name: Arc::from(name.as_str()), + }) + .collect::>() + .into(), + eval_grid: None, + }); + let tree = match compile(&stmt, Arc::clone(&series_schema), NATIVE_ORDERING) { + Ok(tree) => tree, + // The compiler refusing a shape is this route declining it: the existing planner gets + // the query exactly as it would have without this flag. + Err(DataFusionError::NotImplemented(_)) => return Ok(None), + Err(error) => return Err(error).context(DataFusionPlanningSnafu), + }; + + let binding = FusedSeriesBinding { + series_schema, + // The tree's own output is evaluated on the statement's grid, and the binding names + // the same one, so a chunk's absolute positions are read on the grid it was computed + // on. + grid, + metric_name: Some(Arc::from(table_name.as_str())), + time_index: time_index.clone(), + tags: tags.clone(), + }; + let output_schema = fused_output_schema(&fields, &tags, &time_index)?; + let node = FusedSeries::try_new( + Arc::from(prom_expr.to_string()), + tree, + binding, + input, + output_schema, + ) + .context(DataFusionPlanningSnafu)?; + + Ok(Some(LogicalPlan::Extension(Extension { + node: Arc::new(node), + }))) + } + + /// The statement being planned, rebuilt from the evaluation numbers the planner holds. + /// + /// The compiler reads the statement's grid, lookback and `@` anchors rather than the planner's + /// mutable state. The route is only offered at the statement root, where `ctx.start` and + /// `ctx.end` are the statement's own bounds: they are rewritten while an anchored window or a + /// subquery is folded, so a nested expression would need the bounds it was actually entered + /// with, which is why no nested expression is offered here. + /// + /// `None` means the bounds are not representable as a [`SystemTime`], which is not a grid a + /// tree can be evaluated on either. + fn native_statement(&self, prom_expr: &PromExpr) -> Option { + let start = UNIX_EPOCH.checked_add(Duration::from_millis(self.ctx.stmt_start as u64))?; + let end = UNIX_EPOCH.checked_add(Duration::from_millis(self.ctx.stmt_end as u64))?; + Some(EvalStmt { + expr: prom_expr.clone(), + start, + end, + interval: Duration::from_millis(self.ctx.interval as u64), + lookback_delta: Duration::from_millis(self.ctx.lookback_delta as u64), + }) + } +} + +/// The columns a [`FusedSeries`] node produces. +/// +/// The shape is the one the existing unary pipeline consumes — the series fields, the tags, then +/// the time index — and each column is declared the way the node's own conversion reads it back: +/// the time column is a non-nullable millisecond timestamp with no timezone, because every +/// evaluated chunk is addressed on the statement's grid and carries one; a field is a nullable +/// `Float64`, because a grid point inside the statement's interval can have no sample; and a tag is +/// a nullable `Utf8`, because a range function drops the metric name and a series can be missing a +/// label entirely. +fn fused_output_schema( + fields: &[String], + tags: &[String], + time_index: &str, +) -> Result { + let mut columns = Vec::with_capacity(fields.len() + tags.len() + 1); + for name in fields { + columns.push(( + None, + Arc::new(Field::new(name, ArrowDataType::Float64, true)), + )); + } + for name in tags { + columns.push((None, Arc::new(Field::new(name, ArrowDataType::Utf8, true)))); + } + columns.push(( + None, + Arc::new(Field::new( + time_index, + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )), + )); + + Ok(Arc::new( + DFSchema::new_with_metadata(columns, Default::default()) + .context(DataFusionPlanningSnafu)?, + )) +} + +/// The selector's own scan, asking the region for the per-series distribution. +/// +/// The distribution is a hint on the *table scan*, which the region server reads when it builds +/// the region scan: `PerSeries` is what makes a region prepare its series scan, whose raw-series +/// entrypoint is the only one that packs a series' chunks into the frame this node reads. The +/// existing planner asks for the same hint indirectly, through a sort on `(series key, time)` that +/// [`ScanHintRule`](crate::optimizer::scan_hint::ScanHintRule) turns into this distribution; the +/// native node has no such sort below it, because it reads raw samples rather than sorted rows, so +/// the hint is attached here. +/// +/// The provider is cloned for this use-site first, exactly as that rule does: a catalog provider can +/// be shared by several scans, and a hint attached to the shared one would change their requests +/// too. A table whose provider is not the region-server provider - a test fixture, or an engine +/// this route cannot read - is left untouched, and the physical route declines it there. +fn with_per_series_distribution(plan: LogicalPlan) -> LogicalPlan { + let LogicalPlan::TableScan(mut scan) = plan else { + return plan; + }; + let Some(source) = scan.source.downcast_ref::() else { + return LogicalPlan::TableScan(scan); + }; + let Some(provider) = source.table_provider.downcast_ref::() else { + return LogicalPlan::TableScan(scan); + }; + let adapter = provider.clone_for_scan(); + adapter.with_distribution(TimeSeriesDistribution::PerSeries); + scan.source = Arc::new(DefaultTableSource::new(Arc::new(adapter))); + LogicalPlan::TableScan(scan) +} + +/// The one selector `expr` folds, if it names exactly one. +/// +/// The walk follows the grammar the native compiler accepts — a call, a range selector, a selector +/// and parentheses — and stops at every other shape, because a shape the compiler cannot compile +/// needs no selector resolved. Two distinct selectors are not one input, and a tree built over the +/// wrong one of them would evaluate a window the query did not write. +fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { + let mut selectors = Vec::new(); + collect_selectors(expr, &mut selectors); + match selectors.as_slice() { + [only] => Some(only), + _ => None, + } +} + +fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>) { + match expr { + PromExpr::VectorSelector(selector) => out.push(selector), + PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => out.push(vs), + PromExpr::Paren(ParenExpr { expr }) => collect_selectors(expr, out), + PromExpr::Call(call) => { + for arg in &call.args.args { + collect_selectors(arg, out); + } + } + _ => {} + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index ae40640b50b..9acdf56b895 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -39,7 +39,7 @@ use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{Extension, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; -use promql::extension_plan::HistogramFold; +use promql::extension_plan::{FusedSeries, HistogramFold}; use promql_parser::label::Labels; use promql_parser::parser; use session::context::QueryContext; @@ -474,6 +474,7 @@ async fn plan_direct_or( table_provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; planner .or_operator( @@ -563,6 +564,7 @@ async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) table_provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let left_context = direct_or_context("lhs", &["job", "k"], "v"); let right_context = direct_or_context("rhs", &["job", "k"], "v"); @@ -635,6 +637,7 @@ async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner ..Default::default() }, promql_annotations: None, + native_fused: false, }; (planner, plan) } @@ -1617,6 +1620,79 @@ async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) assert_eq!(plan.display_indent_schema().to_string(), expected); } +/// Plans `query` with the native fused-series route on or off. +/// +/// The table is the one [`build_test_table_provider`] builds — one tag `tag_0`, the millisecond +/// time index `timestamp`, one field `field_0` — which is the shape the selector context resolves +/// to a raw series schema of one field. +async fn build_native_fused_plan(query: &str, native_fused: bool) -> LogicalPlan { + let eval_stmt = build_eval_stmt(query); + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + PromPlanner::stmt_to_plan_with_annotations( + table_provider, + &eval_stmt, + &build_query_engine_state(), + None, + native_fused, + ) + .await + .unwrap() +} + +fn find_fused_series(plan: &LogicalPlan) -> Option<&FusedSeries> { + if let LogicalPlan::Extension(Extension { node }) = plan + && let Some(fused) = node.as_any().downcast_ref::() + { + return Some(fused); + } + + plan.inputs().into_iter().find_map(find_fused_series) +} + +/// The flag on and a shape the compiler accepts: the statement plans to one native node. +#[tokio::test] +async fn native_fused_rate_plans_to_a_fused_series_node() { + let plan = build_native_fused_plan("rate(some_metric[5m])", true).await; + + let display = plan.display_indent().to_string(); + assert!( + display.contains("FusedSeries: rate(some_metric[5m])"), + "{display}" + ); + // The node reads the selector's own scan; nothing between them packs the frame it reads. + assert!(display.contains("TableScan: some_metric"), "{display}"); + // The node carries the query a reader wrote, not the shape it compiled to. + let fused = find_fused_series(&plan).expect("the plan holds a native node"); + assert_eq!("rate(some_metric[5m])", fused.promql()); +} + +/// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. +#[tokio::test] +async fn native_fused_unsupported_shape_falls_back() { + let plan = build_native_fused_plan("sum(rate(some_metric[5m]))", true).await; + + let display = plan.display_indent().to_string(); + assert!(!display.contains("FusedSeries"), "{display}"); + // Falling back is not an approximation of the existing plan: it is the existing plan. + let without_flag = build_native_fused_plan("sum(rate(some_metric[5m]))", false).await; + assert_eq!(display, without_flag.display_indent().to_string()); +} + +/// The flag off: no statement is offered to the native route at all. +#[tokio::test] +async fn native_fused_flag_off_keeps_the_existing_plan() { + let plan = build_native_fused_plan("rate(some_metric[5m])", false).await; + + let display = plan.display_indent().to_string(); + assert!(!display.contains("FusedSeries"), "{display}"); + assert!(display.contains("PromRangeManipulate"), "{display}"); +} + #[tokio::test] async fn single_abs() { do_single_instant_function_call("abs", "abs").await; @@ -3328,6 +3404,7 @@ async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() { table_provider, ctx: PromPlannerContext::from_eval_stmt(&eval_stmt).unwrap(), promql_annotations: None, + native_fused: false, }; let plan = planner .prom_expr_to_plan(&eval_stmt.expr, &build_query_engine_state()) @@ -4463,6 +4540,7 @@ async fn mixed_histogram_helpers_report_annotations() { &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), &state, Some(collector.clone()), + false, ) .await .unwrap(); @@ -4623,6 +4701,7 @@ async fn mixed_or_topk_bottomk_ignore_native_histograms() { &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), &state, Some(collector.clone()), + false, ) .await .unwrap(); @@ -4960,6 +5039,7 @@ async fn mixed_native_histogram_rate_executes_real_ranges() { ..Default::default() }, promql_annotations: Some(collector.clone()), + native_fused: false, }; let input = LogicalPlan::Extension(Extension { node: Arc::new( @@ -6289,6 +6369,7 @@ async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { table_provider, ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")).unwrap(), promql_annotations: None, + native_fused: false, }; planner.ctx.time_index_column = Some("timestamp".to_string()); planner.ctx.start = 1_000; @@ -7129,6 +7210,7 @@ async fn test_direct_or_numeric_value_types() { table_provider: provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let left_context = direct_or_context("lhs", &["job"], "v"); let right_context = direct_or_context("rhs", &["job"], "v"); @@ -7616,6 +7698,7 @@ async fn malformed_classic_bucket_does_not_drop_native_histogram() { &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), &state, Some(collector.clone()), + false, ) .await .unwrap(); @@ -7699,6 +7782,7 @@ async fn test_mixed_binary_operator_reports_only_dropped_samples() { &operator_eval_stmt(query), &state, Some(annotations.clone()), + false, ) .await .unwrap(); @@ -8360,6 +8444,7 @@ async fn test_mixed_or_value_aliases_do_not_replace_labels() { table_provider, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let left = LogicalPlanBuilder::from(scan(&left)) .project(vec![ diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index 9ca8a8b646c..be0f5b3b28f 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -411,6 +411,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { .await, ctx: PromPlannerContext::default(), promql_annotations: Some(collector.clone()), + native_fused: false, }; let left_context = direct_or_context("lhs", &["job", marker], "v"); let right_context = direct_or_context("rhs", &["job", marker], "v"); @@ -712,6 +713,7 @@ async fn binary_joins_align_only_the_temporality_marker() { .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let joined = planner .join_on_non_field_columns( @@ -751,6 +753,7 @@ async fn binary_joins_align_only_the_temporality_marker() { .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let set = planner .set_op_on_non_field_columns( @@ -806,6 +809,7 @@ async fn binary_joins_align_only_the_temporality_marker() { .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let set = planner .set_op_on_non_field_columns( @@ -918,6 +922,7 @@ async fn binary_joins_align_dictionary_temporality_marker_with_tagless_vector() .await, ctx: PromPlannerContext::default(), promql_annotations: None, + native_fused: false, }; let scan = |name, table: Arc| { LogicalPlanBuilder::scan(name, provider_as_source(table), None) From ff6c5db30708ba08cc438f4cf22a18246e1abd38 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 21:41:03 +0800 Subject: [PATCH 02/38] feat(query): place a fused node region-side, and ask for per-series scans A `FusedSeries` node is built for one metric's series, and a region holds a subset of those series with none split across regions, so the node commutes with the region split when what it reads is per-series. It says so now, and the analyzer may place it inside the merge scan. Three conditions, each drawn from the conversion that has to succeed there: - The node's input is exactly the table scan. The region builds the node from the scan's own raw-series entrypoint and cannot look through a filter, a projection or a sort in between, so a node over anything else has nothing to read where it would be placed. - The node binds a metric name, because the region resolves the series of one named metric. - No node of its tree folds across series. A region holds part of a metric's series, so an aggregate or a binary operation evaluated there would consume a partial input and return a wrong answer without saying so. Every operation the compiler emits today - input, instant, window, map - reads one series at a time, so this refuses nothing buildable; it is the conservative answer for a tree built some other way. A scan below such a node also asks for the per-series distribution now. That is what makes a region prepare its series scan, and only that path packs a series' chunks into the frame the node reads. Both are inert without a `FusedSeries` node, so an ordinary query keeps its plan and its scan request. Not yet reached: the analyzer asks whether a plan encodes before it asks whether a node commutes, and this node has no Substrait serializer, so the plan is expanded at the node and the region still receives the bare scan. A test pins that boundary so it fails the day the node can cross it. Signed-off-by: discord9 --- src/query/src/dist_plan/commutativity.rs | 298 ++++++++++++++++++++++- src/query/src/optimizer/scan_hint.rs | 167 ++++++++++++- 2 files changed, 458 insertions(+), 7 deletions(-) diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index d02b746a8bf..4d6033729e0 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -21,8 +21,9 @@ use datafusion::error::Result as DfResult; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ - EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, + EmptyMetric, FusedSeries, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, }; +use promql::native::plan::{NativeExpr, NativeNode}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use crate::dist_plan::MergeScanLogicalPlan; @@ -240,6 +241,29 @@ impl Categorizer { Commutativity::Commutative } + name if name == FusedSeries::name() => { + let fused = plan.as_any().downcast_ref::().unwrap(); + // A region builds this node from the region scan itself: the conversion takes that + // scan's raw-series entrypoint and cannot look through anything in between, so a + // node whose input is not that scan has nothing to read where it would be placed + // and stays where it was planned. + if !matches!(fused.inputs().as_slice(), [LogicalPlan::TableScan(_)]) { + return Commutativity::NonCommutative; + } + // The region asks the scan for the series of one named metric, so a node that binds + // no name cannot be built there at all. + if fused.binding().metric_name.is_none() { + return Commutativity::NonCommutative; + } + // A region holds a subset of the metric's series and no series is split across + // regions, so the node may be evaluated there only when every node of its tree + // reads one series at a time. + if reads_across_series(fused.tree().root()) { + return Commutativity::NonCommutative; + } + + Commutativity::Commutative + } name if name == SeriesNormalize::name() || name == InstantManipulate::name() || name == RangeManipulate::name() => @@ -346,19 +370,60 @@ impl Categorizer { } } +/// Whether any node of `tree` reads more than one series at a time. +/// +/// Everything the native compiler emits today reads one series: an instant selection, a window and +/// a point-wise map each select, fold or map the samples of the series they were handed. An +/// aggregate folds a label set's series together and a binary operation pairs them, so both consume +/// several series at once. Neither has a compiler producer yet, which is why refusing them costs +/// nothing today: a tree built some other way is exactly what this check is for, because evaluating +/// it one region at a time would fold a partial input and answer a wrong query silently. +fn reads_across_series(node: &NativeNode) -> bool { + match &node.expr { + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, + NativeExpr::Instant { input, .. } + | NativeExpr::Window { input, .. } + | NativeExpr::Map { input, .. } => reads_across_series(input), + NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } => true, + } +} + #[cfg(test)] mod tests { use std::collections::{BTreeMap, BTreeSet}; + use std::time::{Duration, UNIX_EPOCH}; + use datafusion::arrow::datatypes::{ + DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit, + }; + use datafusion::common::{DFSchema, DFSchemaRef}; + use datafusion::datasource::DefaultTableSource; use datafusion_common::Column; + use datafusion_common::config::ConfigOptions; use datafusion_expr::LogicalPlanBuilder; use datafusion_expr::expr::ScalarFunction; + use datafusion_expr::{Extension, LogicalPlan}; use datafusion_functions::core::coalesce; + use datafusion_optimizer::analyzer::AnalyzerRule; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{FieldSpec, SeriesSchema}; + use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; use promql::functions::{ NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile, }; + use promql::native::plan::{ + AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, Grouping, + NativeType, + }; + use promql_parser::parser; + use store_api::storage::RegionId; + use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use super::*; + use crate::dist_plan::DistPlannerAnalyzer; + use crate::optimizer::test_util::mock_table_provider; + use crate::promql::native_compile::{compile, statement_grid}; + use crate::query_engine::DefaultSerializer; #[test] fn series_divide_by_tsid_is_commutative() { @@ -393,6 +458,237 @@ mod tests { )); } } + + /// The statement one `FusedSeries` test node is bound to. + fn native_statement() -> promql_parser::parser::EvalStmt { + promql_parser::parser::EvalStmt { + expr: parser::parse("some_metric").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + } + } + + /// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to. + /// + /// The tree is the simplest shape the compiler accepts, because what these tests change is + /// where the node may be placed, not what it computes. + fn fused_series_plan(input: LogicalPlan) -> LogicalPlan { + fused_series_plan_named(input, Some("some_metric")) + } + + /// The same node, with the metric name it binds chosen by the caller. + fn fused_series_plan_named(input: LogicalPlan, metric_name: Option<&str>) -> LogicalPlan { + let stmt = native_statement(); + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile( + &stmt, + Arc::clone(&series_schema), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let binding = FusedSeriesBinding { + series_schema, + grid: statement_grid(&stmt).unwrap(), + metric_name: metric_name.map(Arc::from), + time_index: "ts".to_string(), + tags: Vec::new(), + }; + let output_schema: DFSchemaRef = Arc::new( + DFSchema::new_with_metadata( + vec![ + ( + None, + Arc::new(Field::new("v0", ArrowDataType::Float64, true)), + ), + ( + None, + Arc::new(Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )), + ), + ], + Default::default(), + ) + .unwrap(), + ); + + LogicalPlan::Extension(Extension { + node: Arc::new( + FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + input, + output_schema, + ) + .unwrap(), + ), + }) + } + + /// The region's own view of the metric's table: a scan whose provider is the region provider. + fn table_scan_plan() -> LogicalPlan { + let provider = mock_table_provider(RegionId::new(1, 1)); + LogicalPlanBuilder::scan( + "some_metric", + Arc::new(DefaultTableSource::new(Arc::new(provider))), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn holds_fused_series(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(extension) = plan + && extension.node.as_any().is::() + { + return true; + } + plan.inputs().into_iter().any(holds_fused_series) + } + + #[test] + fn fused_series_over_its_scan_is_commutative() { + let plan = fused_series_plan(table_scan_plan()); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::Commutative + )); + } + + /// The region builds the node from the scan's own raw-series entrypoint, so an input that is + /// not that scan leaves the node with nothing to read where it would be placed. + #[test] + fn fused_series_over_anything_but_a_scan_is_not_commutative() { + let scan = table_scan_plan(); + let input = LogicalPlanBuilder::from(scan) + .filter(datafusion_expr::col("k0").eq(datafusion_expr::lit("a"))) + .unwrap() + .build() + .unwrap(); + let plan = fused_series_plan(input); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::NonCommutative + )); + } + + /// The region asks the scan for one named metric's series, so a node binding no name cannot + /// be built there. + #[test] + fn fused_series_that_binds_no_metric_name_is_not_commutative() { + let plan = fused_series_plan_named(table_scan_plan(), None); + + assert!(matches!( + Categorizer::check_plan(&plan, None).unwrap(), + Commutativity::NonCommutative + )); + } + + /// A tree built by something other than the compiler can hold a node that folds several + /// series. The check is a property of the tree, so it is exercised on the node itself: the + /// compiler emits no such node today, and nothing else would build a tree containing one. + #[test] + fn a_tree_that_folds_across_series_is_not_per_series() { + let raw = NativeType::Samples(Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + })); + let input = NativeNode { + expr: NativeExpr::Input(BoundInput { + output: raw.clone(), + ordering: SeriesOrdering::KeyThenTime, + }), + output: raw, + }; + let group = Grouping::By(Arc::from(["k0".into()])); + let aggregate = NativeNode { + expr: NativeExpr::Aggregate { + input: Box::new(input), + op: BoundAggregate { + kind: AggregateKind::Sum, + group: group.clone(), + stage: AggregateStage::Partial, + }, + }, + output: NativeType::AggregateState(Arc::new(AggregateStateSchema { + group, + accumulators: Arc::from([]), + grid: statement_grid(&native_statement()).unwrap(), + })), + }; + assert!(reads_across_series(&aggregate)); + + // The same walk over what the compiler emits stays per-series. + let single_selector = compile( + &native_statement(), + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + assert!(!reads_across_series(single_selector.root())); + } + + /// Blocks placement, and pins the blocker rather than the intent. + /// + /// `PlanRewriter::should_expand` asks two questions about a node before it lets the descent + /// continue, and the Substrait one comes first: unless the *whole* root plan encodes, every + /// subtree that fails to encode is expanded on, whichever way `check_plan` answers. A + /// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers + /// `NotImplemented` for every name it does not list — so its subtree does not encode, the + /// boundary is drawn at the node, and the region receives the bare scan instead of the node. + /// + /// The commutativity arm above is necessary and not sufficient. This test fails the day the + /// node gains a serializer branch (or the analyzer stops asking), which is when the placement + /// assertion this one replaces can be written. + #[test] + fn fused_series_does_not_encode_so_the_analyzer_keeps_it_on_the_frontend() { + let plan = fused_series_plan(table_scan_plan()); + + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "the fused node encodes now, so it can be placed inside the merge scan" + ); + + let analyzed = DistPlannerAnalyzer {} + .analyze(plan, &ConfigOptions::default()) + .unwrap(); + let mut remote_input_holds_fused = false; + analyzed + .apply(|node| { + if let LogicalPlan::Extension(extension) = node + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + remote_input_holds_fused = holds_fused_series(merge_scan.input()); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert!( + !remote_input_holds_fused, + "the analyzer placed the node despite the encoding gate: {analyzed}" + ); + } } pub type Transformer = Arc Option>; diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index 4685889c998..0a1334dd9c0 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -26,7 +26,7 @@ use datafusion_expr::expr::Sort; use datafusion_expr::{Expr, LogicalPlan, utils}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; use datatypes::arrow::datatypes::{DataType, TimeUnit as ArrowTimeUnit}; -use promql::extension_plan::{InstantManipulate, SeriesDivide, SeriesNormalize}; +use promql::extension_plan::{FusedSeries, InstantManipulate, SeriesDivide, SeriesNormalize}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{TimeSeriesDistribution, TimeSeriesRowSelector}; @@ -37,6 +37,9 @@ use crate::dummy_catalog::DummyTableProvider; /// - the nearest order requirement to the leaf table scan node as ordering hint. /// - the group by columns when all aggregate functions are `last_value` as /// time series row selector hint. +/// - the per-series distribution when the leaf table scan node is the input of a +/// [`FusedSeries`] node, which reads the raw series frames a region only packs +/// for that distribution. /// /// [`ScanRequest`]: store_api::storage::ScanRequest #[derive(Debug)] @@ -68,6 +71,7 @@ impl ScanHintRule { fn set_hints( plan: LogicalPlan, rewriter: &mut ScanHintRewriter, + fused_series_input: bool, ) -> Result> { let LogicalPlan::TableScan(mut table_scan) = plan else { return Ok(Transformed::no(plan)); @@ -89,8 +93,10 @@ impl ScanHintRule { }; let use_last_row = rewriter.inside_single_evaluation && filters_preserve_last_row; - let has_hint = - rewriter.order_expr.is_some() || rewriter.ts_row_selector.is_some() || use_last_row; + let has_hint = rewriter.order_expr.is_some() + || rewriter.ts_row_selector.is_some() + || use_last_row + || fused_series_input; if !has_hint { return Ok(Transformed::no(LogicalPlan::TableScan(table_scan))); } @@ -101,6 +107,14 @@ impl ScanHintRule { // local as well as the new LastRow hint. let adapter = original.clone_for_scan(); Self::apply_hints(&adapter, rewriter); + if fused_series_input { + // The scan below a `FusedSeries` node reads packs of raw series, which a region only + // produces for a per-series request; the existing PromQL pipeline asks for the same + // distribution through a sort on `(series key, time)`, which the native node does not + // have below it. Applying it after the order hints makes the plan shape the deciding + // one: the node's own input is what reads the frames. + adapter.with_distribution(TimeSeriesDistribution::PerSeries); + } if use_last_row { // Apply the instant-derived hint after the aggregate hint. Both // select LastRow today, and this ordering preserves the existing @@ -275,6 +289,12 @@ struct ScanHintRewriter { ts_stack: Vec, Column)>>, inside_single_evaluation: bool, single_evaluation_stack: Vec, + /// Whether the node being visited is the input of a [`FusedSeries`] node. + /// + /// The flag only describes that one edge: it is set while a `FusedSeries` node is visited, so + /// the node visited next — that node's own input — sees it, and a node further down does not. + fused_series_input: bool, + fused_series_input_stack: Vec, } impl TreeNodeRewriter for ScanHintRewriter { @@ -285,6 +305,13 @@ impl TreeNodeRewriter for ScanHintRewriter { self.ts_stack.push(self.ts_row_selector.clone()); self.single_evaluation_stack .push(self.inside_single_evaluation); + self.fused_series_input_stack.push(self.fused_series_input); + // Read the edge this node sits on before recording what this node is, so a `FusedSeries` + // node hands the flag to its own input and to nothing below it. + let fused_series_input = std::mem::replace( + &mut self.fused_series_input, + matches!(&node, LogicalPlan::Extension(extension) if extension.node.as_any().is::()), + ); if let LogicalPlan::Sort(sort) = &node { self.order_expr = Some(sort.expr.clone()); @@ -321,7 +348,7 @@ impl TreeNodeRewriter for ScanHintRewriter { } } - ScanHintRule::set_hints(node, self) + ScanHintRule::set_hints(node, self, fused_series_input) } fn f_up(&mut self, node: LogicalPlan) -> Result> { @@ -334,6 +361,9 @@ impl TreeNodeRewriter for ScanHintRewriter { if let Some(previous) = self.single_evaluation_stack.pop() { self.inside_single_evaluation = previous; } + if let Some(previous) = self.fused_series_input_stack.pop() { + self.fused_series_input = previous; + } Ok(Transformed::no(node)) } } @@ -446,6 +476,7 @@ impl ScanHintRewriter { mod test { use std::collections::HashMap; use std::sync::Arc; + use std::time::{Duration, UNIX_EPOCH}; use datafusion::functions_aggregate::first_last::last_value_udaf; use datafusion::functions_aggregate::min_max::max_udaf; @@ -454,21 +485,26 @@ mod test { use datafusion::logical_expr::{WindowFrame, WindowFunctionDefinition}; use datafusion::prelude::JoinType; use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_common::{DFSchema, DFSchemaRef}; use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, Cast, WindowFunctionParams, }; use datafusion_expr::expr_fn::scalar_subquery; use datafusion_expr::{Extension, LogicalPlan, LogicalPlanBuilder, col, lit}; use datafusion_optimizer::OptimizerContext; - use datatypes::arrow::datatypes::DataType; + use datatypes::arrow::datatypes::{DataType, Field, TimeUnit}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::ColumnSchema; - use promql::extension_plan::RangeManipulate; + use datatypes::series::SeriesOrdering; + use datatypes::series::samples::{FieldSpec, SeriesSchema}; + use promql::extension_plan::{FusedSeries, FusedSeriesBinding, RangeManipulate}; + use promql_parser::parser; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use store_api::storage::{RegionId, TimeSeriesRowSelector}; use super::*; + use crate::promql::native_compile::{compile, statement_grid}; fn scan_requests(plan: &LogicalPlan) -> Vec { scan_requests_with_names(plan) @@ -1520,4 +1556,123 @@ mod test { Some(TimeSeriesDistribution::PerSeries) ); } + + /// A `FusedSeries` node over `input`, bound to the tree a bare selector compiles to. + /// + /// The rule reads the plan shape, not what the node computes, so the simplest shape the + /// compiler accepts is enough for these tests. + fn fused_series(input: LogicalPlan) -> LogicalPlan { + let stmt = promql_parser::parser::EvalStmt { + expr: parser::parse("some_metric").unwrap(), + start: UNIX_EPOCH, + end: UNIX_EPOCH + .checked_add(Duration::from_secs(100_000)) + .unwrap(), + interval: Duration::from_secs(5), + lookback_delta: Duration::from_secs(1), + }; + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile( + &stmt, + Arc::clone(&series_schema), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + let binding = FusedSeriesBinding { + series_schema, + grid: statement_grid(&stmt).unwrap(), + metric_name: Some("some_metric".into()), + time_index: "ts".to_string(), + tags: Vec::new(), + }; + let output_schema: DFSchemaRef = Arc::new( + DFSchema::new_with_metadata( + vec![ + (None, Arc::new(Field::new("v0", DataType::Float64, true))), + ( + None, + Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )), + ), + ], + Default::default(), + ) + .unwrap(), + ); + + LogicalPlan::Extension(Extension { + node: Arc::new( + FusedSeries::try_new( + Arc::from("some_metric"), + tree, + binding, + input, + output_schema, + ) + .unwrap(), + ), + }) + } + + /// The region only packs raw series frames for a per-series request, and the node below is + /// what reads them, so the node's own input asks for that distribution. + #[test] + fn a_scan_under_fused_series_asks_for_the_per_series_distribution() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let plan = fused_series(scan_plan(provider.clone(), "t")); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!( + scan_requests(&rewritten)[0].distribution, + Some(TimeSeriesDistribution::PerSeries) + ); + // The hint belongs to this use-site only, never to the shared catalog provider. + assert_eq!(provider.scan_request().distribution, None); + } + + /// The ordinary path: a scan with no fused node above it is left exactly as it was. + #[test] + fn a_scan_not_under_fused_series_keeps_its_request() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let plan = scan_plan(provider.clone(), "t"); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!(scan_requests(&rewritten)[0].distribution, None); + assert_eq!(provider.scan_request().distribution, None); + } + + /// Only the node's own input is a raw series reader; a scan further down feeds something else + /// and keeps the request it had. + #[test] + fn a_scan_below_fused_series_input_keeps_its_request() { + let provider = Arc::new(mock_table_provider(RegionId::new(1, 1))); + let scan = scan_plan(provider.clone(), "t"); + let input = LogicalPlanBuilder::from(scan) + .project(vec![col("ts"), col("v0")]) + .unwrap() + .build() + .unwrap(); + let plan = fused_series(input); + + let rewritten = ScanHintRule + .rewrite(plan, &OptimizerContext::default()) + .unwrap() + .data; + + assert_eq!(scan_requests(&rewritten)[0].distribution, None); + } } From ba9f36fdd08f5f796b40dab2590a2e7bbde57437 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 21:50:52 +0800 Subject: [PATCH 03/38] docs(query): say what the fused placement does not prove The arm above answers that a fused node commutes with the region split, and it listed "no series is split across regions" among its reasons. That is not established anywhere. A region evaluates the whole tree over the series its own scan returns, so a series whose samples reach two regions is evaluated per fragment, and a range function over a fragment is not the range function over the series. What the placement does have is three proofs about the node - its input is the scan, it binds a metric name, and no node of its tree folds across series. What it does not have is any check that the series stays in one region: the storage partition key is whatever `partition_column_names()` names and may be a column that varies within one series, and this categorizer runs before region selection, so there is no selected-region vector here to consult either. So the comment now says that plainly, names the deployment the placement is correct under, and says what would make it unconditional. A single-region deployment satisfies it; a partitioned table over several regions would answer wrongly rather than fail. Signed-off-by: discord9 --- src/query/src/dist_plan/commutativity.rs | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 4d6033729e0..8d26605cece 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -255,13 +255,28 @@ impl Categorizer { if fused.binding().metric_name.is_none() { return Commutativity::NonCommutative; } - // A region holds a subset of the metric's series and no series is split across - // regions, so the node may be evaluated there only when every node of its tree - // reads one series at a time. + // A region evaluates the tree over whatever series its own scan returns, so this + // is only answerable when every node of the tree reads one series at a time. An + // aggregate or a binary operation would fold or pair across series, and one region + // holds only part of them. if reads_across_series(fused.tree().root()) { return Commutativity::NonCommutative; } + // TEMPORARY, and this is the one thing here that is not proved: the placement is + // only correct while each series reaches exactly one region. A region evaluates the + // whole tree over the series its scan returns, so a series whose samples are split + // across regions is evaluated per fragment, and a range function over a fragment is + // not the range function over the series. Nothing checks that a series stays in one + // region: the storage partition key is whatever `partition_column_names()` names and + // may be a column that varies within one series, and this categorizer runs before + // region selection, so it has no selected-region vector to consult either. A single + // region deployment satisfies it; a partitioned table over several regions does not, + // and would return a wrong answer rather than an error. + // + // What would make this unconditional is a partition key the label set determines, + // checked against the table's physical partition expressions. Until then the + // placement is a proof of concept over one region. Commutativity::Commutative } name if name == SeriesNormalize::name() From da27d115ca15a43908fff5e8406a32ce1d5dc5fb Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 21 Oct 2026 22:54:53 +0800 Subject: [PATCH 04/38] feat(query): let a region input travel as a plan when regions are in process A plan holding a node with no Substrait serializer was always cut at that node, because the analyzer asks whether a plan encodes before it asks whether a node commutes with the region split. That is right for a region reached over the wire: the plan is encoded into the request and has to decode on the other side. Standalone does not use that path. Its region server is the one built in the same process, its query handler passes the request straight to `RegionServer::handle_read`, and the plan stays a `LogicalPlan` from the frontend to the region. Encoding it is work nobody reads, and refusing to place a node because it cannot be encoded refuses a plan that would have travelled intact. So the engine now knows whether its regions are reached in process, and the encodability gate applies only when they are not. The capability is set where it can be stated truthfully - the standalone bootstrap, which creates both the frontend and the region server it reads - and it defaults to false everywhere else, so an ordinary build and every other deployment keeps today's behaviour byte for byte. Skipping the gate skips one early return and adds no reason to expand: `Categorizer::check_plan` still runs and still decides, and a plan that does not encode for any other reason still expands at the node that caused it. Two tests hold those two halves apart. Signed-off-by: discord9 --- src/cmd/src/standalone.rs | 3 + src/frontend/src/instance/builder.rs | 21 ++- src/query/src/dist_plan/analyzer.rs | 99 +++++++++-- src/query/src/dist_plan/analyzer/test.rs | 146 +++++++-------- src/query/src/dist_plan/commutativity.rs | 215 +++++++++++++++++++---- src/query/src/query_engine.rs | 36 +++- src/query/src/query_engine/state.rs | 37 +++- 7 files changed, 424 insertions(+), 133 deletions(-) diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index ca0f3b74057..ec8305efa1c 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -734,6 +734,9 @@ impl StartCommand { procedure_executor.clone(), process_manager, ) + // The region servers this frontend reads are the ones built above, in this process, so a + // region input can travel as a logical plan instead of Substrait bytes. + .with_in_process_region_transport() .with_local_file_access(local_file_access); plugins::setup_frontend_plugins_post_build(&mut plugins, &plugin_opts, &fe_instance) diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 624ccb2f9cd..918655d0c5e 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -78,6 +78,9 @@ pub struct FrontendBuilder { procedure_executor: ProcedureExecutorRef, process_manager: ProcessManagerRef, local_file_access: LocalFileAccess, + /// Whether this frontend reaches the regions of its datanodes in process. Declared by whoever + /// assembles the frontend, because only it knows what its node manager talks to. + in_process_region_transport: bool, } impl FrontendBuilder { @@ -105,6 +108,7 @@ impl FrontendBuilder { procedure_executor, process_manager, local_file_access: LocalFileAccess::Disabled, + in_process_region_transport: false, } } @@ -157,6 +161,20 @@ impl FrontendBuilder { self } + /// Declares that every region this frontend queries is read in this very process. + /// + /// The caller must own the node manager this frontend was built with, and know that its + /// `DatanodeManager` hands region reads to a `RegionServer` in this process — as standalone + /// startup does, since it passes the datanode's own region server to the node manager it + /// installs. The distributed planner then hands a region a `LogicalPlan` instead of Substrait + /// bytes, which `RegionServer::handle_read` accepts. A frontend whose node manager reaches + /// remote datanodes must leave it off: those datanodes only accept Substrait bytes + /// (`RegionServer::handle_remote_read`), so a plan that does not encode would fail there. + pub fn with_in_process_region_transport(mut self) -> Self { + self.in_process_region_transport = true; + self + } + pub fn options(&self) -> &FrontendOptions { &self.options } @@ -315,7 +333,7 @@ impl FrontendBuilder { let mut query_options = self.options.query.clone(); query_options.enable_per_region_metrics = self.options.logging.enable_per_region_metrics; - let query_engine = QueryEngineFactory::try_new_with_plugins( + let query_engine = QueryEngineFactory::try_new_with_plugins_and_region_transport( self.catalog_manager.clone(), Some(partition_manager.clone()), Some(region_query_handler.clone()), @@ -323,6 +341,7 @@ impl FrontendBuilder { Some(procedure_service_handler), Some(Arc::new(flow_service)), true, + self.in_process_region_transport, plugins.clone(), query_options, ) diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index a39514a0e43..955cc6cbe55 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -106,8 +106,35 @@ impl ExtensionOptions for DistPlannerOptions { } } -#[derive(Debug)] -pub struct DistPlannerAnalyzer; +/// Rewrites a logical plan into a distributed plan: the part above a merge scan runs on the +/// frontend, and every `MergeScan` carries the sub-plan that runs in a region. +#[derive(Debug, Default)] +pub struct DistPlannerAnalyzer { + /// Whether every region a rewritten plan reads is reached in process. + /// + /// A region input normally has to be encoded to Substrait before it can leave the frontend, and + /// [`PlanRewriter::should_expand`] cuts the plan at the first node that does not encode. When + /// regions are reached in process, that input travels as a `LogicalPlan` instead, so nothing is + /// encoded and no node has to be cut for encoding's sake. + /// + /// Whoever builds the query engine sets this, once, from the transport it wired. It is neither a + /// query-time decision nor a property of a region. + in_process_region_transport: bool, +} + +impl DistPlannerAnalyzer { + /// Creates the analyzer. + /// + /// `in_process_region_transport` must only be `true` for an engine whose region reads are served + /// by region servers in the same process. Its plans are then handed to + /// `RegionServer::handle_read` as `LogicalPlan`s; the remote read path + /// (`RegionServer::handle_remote_read`) decodes Substrait bytes and cannot do that. + pub fn new(in_process_region_transport: bool) -> Self { + Self { + in_process_region_transport, + } + } +} impl AnalyzerRule for DistPlannerAnalyzer { fn name(&self) -> &str { @@ -250,12 +277,15 @@ fn unwrap_dictionary_literals(plan: LogicalPlan) -> DfResult DfResult { + let in_process_region_transport = self.in_process_region_transport; // Use the subquery-aware transform so expression subqueries (including // nested scalar subqueries left in place by DataFusion 55's // `enable_physical_uncorrelated_scalar_subquery`) are visited at every // depth and their table scans get wrapped in `MergeScan`. - let plan = plan.transform_with_subqueries(&Self::inspect_plan_with_subquery)?; - let mut rewriter = PlanRewriter::new(&plan.data); + let plan = plan.transform_with_subqueries(&|plan| { + Self::inspect_plan_with_subquery(plan, in_process_region_transport) + })?; + let mut rewriter = PlanRewriter::new(&plan.data, in_process_region_transport); let result = plan.data.rewrite(&mut rewriter)?.data; Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } @@ -267,7 +297,10 @@ impl DistPlannerAnalyzer { Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } - fn inspect_plan_with_subquery(plan: LogicalPlan) -> DfResult> { + fn inspect_plan_with_subquery( + plan: LogicalPlan, + in_process_region_transport: bool, + ) -> DfResult> { // Workaround for https://github.com/GreptimeTeam/greptimedb/issues/5469 and https://github.com/GreptimeTeam/greptimedb/issues/5799 // FIXME(yingwen): Remove the `Limit` plan once we update DataFusion. if let LogicalPlan::Limit(_) | LogicalPlan::Distinct(_) = &plan { @@ -277,7 +310,12 @@ impl DistPlannerAnalyzer { let exprs = plan .expressions_consider_join() .into_iter() - .map(|e| e.transform(&Self::transform_subquery).map(|x| x.data)) + .map(|e| { + e.transform(|subquery_expr| { + Self::transform_subquery(subquery_expr, in_process_region_transport) + }) + .map(|x| x.data) + }) .collect::>>()?; // Some plans that are special treated (should not call `with_new_exprs` on them) @@ -289,27 +327,33 @@ impl DistPlannerAnalyzer { } } - fn transform_subquery(expr: Expr) -> DfResult> { + fn transform_subquery( + expr: Expr, + in_process_region_transport: bool, + ) -> DfResult> { match expr { Expr::Exists(exists) => Ok(Transformed::yes(Expr::Exists(Exists { - subquery: Self::handle_subquery(exists.subquery)?, + subquery: Self::handle_subquery(exists.subquery, in_process_region_transport)?, negated: exists.negated, }))), Expr::InSubquery(in_subquery) => Ok(Transformed::yes(Expr::InSubquery(InSubquery { expr: in_subquery.expr, - subquery: Self::handle_subquery(in_subquery.subquery)?, + subquery: Self::handle_subquery(in_subquery.subquery, in_process_region_transport)?, negated: in_subquery.negated, }))), Expr::ScalarSubquery(scalar_subquery) => Ok(Transformed::yes(Expr::ScalarSubquery( - Self::handle_subquery(scalar_subquery)?, + Self::handle_subquery(scalar_subquery, in_process_region_transport)?, ))), _ => Ok(Transformed::no(expr)), } } - fn handle_subquery(subquery: Subquery) -> DfResult { - let mut rewriter = PlanRewriter::new(&subquery.subquery); + fn handle_subquery( + subquery: Subquery, + in_process_region_transport: bool, + ) -> DfResult { + let mut rewriter = PlanRewriter::new(&subquery.subquery, in_process_region_transport); let mut rewrote_subquery = subquery .subquery .as_ref() @@ -400,7 +444,13 @@ enum RewriterStatus { struct PlanRewriter { /// Whether the whole plan this rewriter walks can be encoded to Substrait. /// Used by [`PlanRewriter::should_expand`] to skip the per-node encoding check. + /// + /// Only read when that gate applies, i.e. when `in_process_region_transport` is false. With + /// in-process regions the question has no answer worth computing and the field stays `false`. whole_plan_encodable: bool, + /// Whether the regions this plan reads are reached in process, so an input crosses to a region + /// as a `LogicalPlan` rather than as Substrait bytes. See [`DistPlannerAnalyzer`]. + in_process_region_transport: bool, /// Current level in the tree level: usize, /// Simulated stack for the `rewrite` recursion @@ -459,11 +509,16 @@ struct PlanRewriter { impl PlanRewriter { /// `plan` must be the root of the tree that is about to be rewritten, otherwise /// [`PlanRewriter::should_expand`] may skip a check it has to perform. - fn new(plan: &LogicalPlan) -> Self { + fn new(plan: &LogicalPlan, in_process_region_transport: bool) -> Self { Self { - whole_plan_encodable: DFLogicalSubstraitConvertor - .encode(plan, DefaultSerializer) - .is_ok(), + // Only a plan that crosses to a region as Substrait bytes has to encode at all, and + // encoding is the expensive part of a pass. With in-process regions the answer is never + // read, so it is not computed. + whole_plan_encodable: !in_process_region_transport + && DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .is_ok(), + in_process_region_transport, ..Default::default() } } @@ -495,7 +550,17 @@ impl PlanRewriter { // is the dominant cost on deep PromQL plans. // When the root does not encode, the plan holds at least one node that has to stay // on the frontend and only the per-node check locates it. - if !self.whole_plan_encodable + // + // This question only exists for a plan that has to cross to a region as Substrait bytes. + // With in-process regions the same `LogicalPlan` is handed to `RegionServer::handle_read` + // never having been encoded, so encodability says nothing about where a node may run, and + // asking it would cut the plan at the first node with no serializer branch (`FusedSeries`) + // before `check_plan` could decide. Skipping it only removes a reason to expand: every + // check below still runs, and a node they call `NonCommutative`/`Unimplemented`/ + // `Unsupported` still expands, so a node no serializer lists only rides to a region when + // the categorizer places it there. + if !self.in_process_region_transport + && !self.whole_plan_encodable && let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer) { debug!( diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 694f2bcde59..c0a4683b1c9 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -68,6 +68,12 @@ fn collect_merge_scan_remote_dyn_filter_producer_ids( producer_ids.extend(producer_id_list); } +/// The analyzer as an engine whose regions are reached over the wire builds it: every plan a region +/// receives from it has to be encodable. +fn analyzer() -> DistPlannerAnalyzer { + DistPlannerAnalyzer::default() +} + struct MergeScanRemoteDynFilterProducerIdCollector<'a> { producer_ids: &'a mut Vec, } @@ -395,9 +401,7 @@ fn frontend_only_histogram_folds_stay_above_merge_scan() { "{name}" ); - let result = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); + let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap(); let result_text = result.to_string(); assert!( result_text.contains("HistogramFold:"), @@ -442,9 +446,7 @@ fn unencodable_unnest_stays_above_merge_scan() { .is_err() ); - let result = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); + let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap(); let result_text = result.to_string(); assert!(result_text.contains("Unnest:"), "{result_text}"); @@ -481,9 +483,7 @@ fn dictionary_literals_are_unwrapped_before_pushdown() { .build() .unwrap(); - let result = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); + let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap(); assert!( result .to_string() @@ -578,7 +578,7 @@ fn expand_proj_sort_proj() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.number", @@ -625,7 +625,7 @@ fn expand_proj_sort_partial_proj() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.number", @@ -661,7 +661,7 @@ fn expand_sort_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, t.pk3, t.ts, t.number", @@ -699,7 +699,7 @@ fn expand_sort_alias_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: something", @@ -738,7 +738,7 @@ fn expand_sort_alias_conflict_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan.clone(), &config); + let result = analyzer().analyze(plan.clone(), &config); assert!(result.is_err(), "Expected error for ambiguous alias"); assert!(format!("{result:?}").contains("AmbiguousReference")); @@ -746,7 +746,7 @@ fn expand_sort_alias_conflict_limit() { config.extensions.insert(DistPlannerOptions { allow_query_fallback: true, }); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", @@ -780,7 +780,7 @@ fn expand_sort_alias_conflict_but_not_really_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1", @@ -818,7 +818,7 @@ fn expand_limit_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: t.pk1 ASC NULLS LAST", @@ -853,7 +853,7 @@ fn expand_sort_limit_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: t.pk1 ASC NULLS LAST", @@ -898,7 +898,7 @@ fn expand_proj_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: min(t.number)", @@ -943,7 +943,7 @@ fn expand_proj_alias_fake_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: pk1, pk2, min(t.number)", @@ -987,7 +987,7 @@ fn expand_proj_alias_aliased_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: pk42, pk43, min(t.number)", @@ -1022,7 +1022,7 @@ fn expand_expr_over_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: substr(t.pk1,Int64(1),Int64(1)), t.pk2, min(t.number)", @@ -1057,7 +1057,7 @@ fn expand_part_col_aggr_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: min(max(t.number))", @@ -1090,7 +1090,7 @@ fn expand_step_aggr_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Aggregate: groupBy=[[]], aggr=[[min(max(t.number))]]", @@ -1135,7 +1135,7 @@ fn expand_part_col_aggr_part_col_aggr() { assert_eq!(expected_original, plan.to_string()); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, min(max(t.number))", @@ -1167,7 +1167,7 @@ fn expand_step_aggr_proj() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: min(t.number)", @@ -1225,7 +1225,7 @@ fn expand_complex_col_req_sort_pql() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2", @@ -1271,7 +1271,7 @@ fn expand_proj_sort_step_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Aggregate: groupBy=[[]], aggr=[[min(t.number)]]", @@ -1312,7 +1312,7 @@ fn expand_proj_sort_limit_step_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Aggregate: groupBy=[[]], aggr=[[min(t.number)]]", " Projection: t.number", @@ -1354,7 +1354,7 @@ fn expand_proj_limit_step_aggr_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: min(t.number) ASC NULLS LAST", " Aggregate: groupBy=[[]], aggr=[[min(t.number)]]", @@ -1396,7 +1396,7 @@ fn expand_proj_sort_part_col_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Aggregate: groupBy=[[t.pk1, t.pk2]], aggr=[[min(t.number)]]", @@ -1438,7 +1438,7 @@ fn expand_proj_sort_limit_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Aggregate: groupBy=[[t.pk1, t.pk2]], aggr=[[min(t.number)]]", " Projection: t.number, t.pk1, t.pk2", @@ -1480,7 +1480,7 @@ fn expand_proj_part_col_aggr_limit_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Sort: t.pk2 ASC NULLS LAST", " Projection: t.pk1, t.pk2, min(t.number)", @@ -1522,7 +1522,7 @@ fn expand_proj_part_col_aggr_sort_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, min(t.number)", " Limit: skip=0, fetch=10", @@ -1557,9 +1557,7 @@ fn test_simplify_select_now_expression() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); common_telemetry::info!("Analyzed plan: {}", result); @@ -1615,7 +1613,7 @@ fn sibling_merge_scans_have_unique_remote_dyn_filter_producer_ids() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let mut producer_ids = Vec::new(); collect_merge_scan_remote_dyn_filter_producer_id_list(&result, &mut producer_ids); @@ -1670,7 +1668,7 @@ fn pre_merge_scan_optimizer_eliminates_projected_false_filter() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_eq!("EmptyRelation: rows=0", result.to_string()); } @@ -1718,7 +1716,7 @@ fn test_simplify_now_expression() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let plan_str = result.to_string(); common_telemetry::info!("Analyzed plan: {}", plan_str); @@ -1757,7 +1755,7 @@ fn expand_proj_limit_part_col_aggr_sort() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); // Pre-MergeScan optimizer intentionally excludes PushDownLimit, so the // remote plan shows an explicit Limit node instead of `fetch=10` on // TableScan. @@ -1802,7 +1800,7 @@ fn expand_proj_limit_sort_part_col_aggr() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); // Pre-MergeScan optimizer intentionally excludes PushDownLimit, so the // remote plan shows an explicit Limit node instead of `fetch=10` on // TableScan. @@ -1840,7 +1838,7 @@ fn expand_step_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Projection: t.pk1, min(t.number)", @@ -1873,7 +1871,7 @@ fn expand_step_aggr_avg_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=10", " Projection: t.pk1, avg(t.number)", @@ -1906,7 +1904,7 @@ fn expand_part_col_aggr_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.pk1, t.pk2, min(t.number)", " Limit: skip=0, fetch=10", @@ -1940,7 +1938,7 @@ fn transform_simple_projection_filter() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Distinct:", " MergeScan [is_placeholder=false]", @@ -1968,7 +1966,7 @@ fn transform_aggregator() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = "Projection: avg(t.number)\ \n MergeScan [is_placeholder=false, remote_input=[\ \nAggregate: groupBy=[[]], aggr=[[avg(t.number)]]\ @@ -1994,7 +1992,7 @@ fn transform_distinct_order() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: t.number", " MergeScan [is_placeholder=false, remote_input=[ @@ -2022,7 +2020,7 @@ fn transform_single_limit() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = "Projection: t.number\ \n MergeScan [is_placeholder=false, remote_input=[ Limit: skip=0, fetch=1 @@ -2063,7 +2061,7 @@ fn transform_unalighed_join_with_alias() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Limit: skip=0, fetch=1", " LeftSemi Join: t.number = right.number", @@ -2099,7 +2097,7 @@ fn transform_subquery_sort_alias() { .build() .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: a.pk1, a.pk2, a.pk3, a.ts, a.number", " MergeSort: a.number ASC NULLS LAST", @@ -2130,7 +2128,7 @@ fn transform_sort_subquery_alias() { .build() .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ "Projection: a.pk1, a.pk2, a.pk3, a.ts, a.number", " MergeSort: a.number ASC NULLS LAST", @@ -2169,7 +2167,7 @@ fn remote_dyn_filter_producer_ids_do_not_collide_between_subquery_and_outer_plan .unwrap() .build() .unwrap(); - let rewritten = DistPlannerAnalyzer {}.try_push_down(outer_plan).unwrap(); + let rewritten = analyzer().try_push_down(outer_plan).unwrap(); let mut producer_ids = BTreeSet::new(); collect_merge_scan_remote_dyn_filter_producer_ids(&rewritten, &mut producer_ids); @@ -2201,7 +2199,7 @@ fn date_bin_ts_group_by() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let expected = [ r#"Projection: date_bin(IntervalDayTime("IntervalDayTime { days: 0, milliseconds: 60000 }"),t.ts), min(t.number)"#, @@ -2249,9 +2247,7 @@ fn test_last_value_order_by() { try_encode_decode_substrait(&plan, ctx.state()); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: last_value(t.ts) ORDER BY [t.ts ASC NULLS FIRST]", @@ -2312,9 +2308,7 @@ fn test_last_value_no_order_by() { try_encode_decode_substrait(&split.lower_state, ctx.state()); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: last_value(t.ts)", @@ -2344,9 +2338,7 @@ fn test_table_scan_projection() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: t.ts", " MergeScan [is_placeholder=false, remote_input=[", @@ -2380,9 +2372,7 @@ fn test_table_scan_cast_projection_pushdown() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: ts", " MergeScan [is_placeholder=false, remote_input=[", @@ -2426,9 +2416,7 @@ fn test_cast_filter_simplified_after_pushdown() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let expected = [ "Projection: ts", " MergeScan [is_placeholder=false, remote_input=[", @@ -2483,7 +2471,7 @@ fn test_join_side_local_filter_pushdown_into_merge_scan() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let predicate = col("t1.pk1").eq(lit("v")); @@ -2555,7 +2543,7 @@ fn test_left_join_left_side_filter_pushdown_into_merge_scan() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let plan_str = result.to_string(); @@ -2611,7 +2599,7 @@ fn test_join_cross_table_predicate_not_pushed_to_single_side() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let plan_str = result.to_string(); @@ -2668,9 +2656,7 @@ fn scheduled_now_yields_stable_literal_in_remote_plan() { scheduled_time: Some(scheduled_time), }); - let result = DistPlannerAnalyzer {} - .analyze(plan.clone(), &config) - .unwrap(); + let result = analyzer().analyze(plan.clone(), &config).unwrap(); let result_str = result.to_string(); common_telemetry::info!("Analyzed plan with scheduled time: {}", result_str); @@ -2706,9 +2692,7 @@ fn scheduled_now_yields_stable_literal_in_remote_plan() { // Absent the extension (default config), the same plan simplifies to a // wall-clock literal — the remote_input contains a variable timestamp. let default_config = ConfigOptions::default(); - let wall_result = DistPlannerAnalyzer {} - .analyze(plan, &default_config) - .unwrap(); + let wall_result = analyzer().analyze(plan, &default_config).unwrap(); let wall_str = wall_result.to_string(); // The wall-clock result must NOT contain the scheduled literal. assert!( @@ -2769,7 +2753,7 @@ fn scheduled_current_timestamp_filter_folds_before_remote_pushdown() { scheduled_time: Some(scheduled_time), }); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); assert_remote_table_scan_filters_are_safe(&result); let result_str = result.to_string(); @@ -2816,7 +2800,7 @@ fn scheduled_none_falls_back_to_wall_clock() { scheduled_time: None, }); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let result_str = result.to_string(); // Must still simplify now() — remotely a literal, top is now(). @@ -2884,7 +2868,7 @@ fn expand_nested_scalar_subquery() { .unwrap(); let config = ConfigOptions::default(); - let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); + let result = analyzer().analyze(plan, &config).unwrap(); let result_str = result.to_string(); // The main scan and both subquery levels must each be wrapped in a MergeScan. diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 8d26605cece..c8d56cc9a04 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -417,7 +417,7 @@ mod tests { use datafusion_common::config::ConfigOptions; use datafusion_expr::LogicalPlanBuilder; use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{Extension, LogicalPlan}; + use datafusion_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion_functions::core::coalesce; use datafusion_optimizer::analyzer::AnalyzerRule; use datatypes::series::SeriesOrdering; @@ -572,6 +572,124 @@ mod tests { plan.inputs().into_iter().any(holds_fused_series) } + /// An extension node with no `ExtensionSerializer` branch: neither one under its own name nor + /// one it shares, since the serializer lists names and not shapes. It stands in for the nodes + /// that cannot reach a region as bytes for their own reasons, which the encoding gate must keep + /// honouring. The categorizer refuses it, as it refuses every name it does not know. + #[derive(Debug, PartialEq, Eq, Hash)] + struct UnserializableNode { + input: LogicalPlan, + schema: DFSchemaRef, + } + + impl PartialOrd for UnserializableNode { + fn partial_cmp(&self, other: &Self) -> Option { + let ordering = self.input.partial_cmp(&other.input); + match ordering { + Some(std::cmp::Ordering::Equal) if self.schema != other.schema => None, + ordering => ordering, + } + } + } + + impl UnserializableNode { + fn new(input: LogicalPlan) -> Self { + let schema = input.schema().clone(); + Self { input, schema } + } + } + + impl UserDefinedLogicalNodeCore for UnserializableNode { + fn name(&self) -> &str { + "UnserializableNode" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "UnserializableNode") + } + + fn with_exprs_and_inputs( + &self, + _exprs: Vec, + mut inputs: Vec, + ) -> DfResult { + Ok(Self::new(inputs.swap_remove(0))) + } + } + + /// The plan that node builds over `input`. + fn unserializable_plan(input: LogicalPlan) -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(UnserializableNode::new(input)), + }) + } + + fn holds_unserializable_node(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(extension) = plan + && extension.node.as_any().is::() + { + return true; + } + plan.inputs().into_iter().any(holds_unserializable_node) + } + + fn holds_table_scan(plan: &LogicalPlan) -> bool { + if matches!(plan, LogicalPlan::TableScan(_)) { + return true; + } + plan.inputs().into_iter().any(holds_table_scan) + } + + /// The merge scan `plan` is, if it is one. + fn as_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { + if let LogicalPlan::Extension(extension) = plan { + extension + .node + .as_any() + .downcast_ref::() + } else { + None + } + } + + /// The first merge scan in `plan`'s visible tree. A merge scan's remote input is not one of its + /// children, so this does not look into one. + fn first_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { + if let Some(merge_scan) = as_merge_scan(plan) { + return Some(merge_scan); + } + plan.inputs().into_iter().find_map(first_merge_scan) + } + + fn analyze(analyzer: DistPlannerAnalyzer, plan: LogicalPlan) -> LogicalPlan { + analyzer.analyze(plan, &ConfigOptions::default()).unwrap() + } + + /// Whether the remote input of any merge scan in `plan` holds a node `wanted` accepts. + fn any_remote_input_holds(plan: &LogicalPlan, wanted: fn(&LogicalPlan) -> bool) -> bool { + let mut holds = false; + plan.apply(|node| { + if let Some(merge_scan) = as_merge_scan(node) { + holds |= wanted(merge_scan.input()); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + holds + } + #[test] fn fused_series_over_its_scan_is_commutative() { let plan = fused_series_plan(table_scan_plan()); @@ -659,49 +777,82 @@ mod tests { assert!(!reads_across_series(single_selector.root())); } - /// Blocks placement, and pins the blocker rather than the intent. + /// The encoding gate is what kept the node on the frontend, and it only exists for a plan that + /// has to be encoded. /// - /// `PlanRewriter::should_expand` asks two questions about a node before it lets the descent - /// continue, and the Substrait one comes first: unless the *whole* root plan encodes, every - /// subtree that fails to encode is expanded on, whichever way `check_plan` answers. A - /// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers - /// `NotImplemented` for every name it does not list — so its subtree does not encode, the - /// boundary is drawn at the node, and the region receives the bare scan instead of the node. + /// `PlanRewriter::should_expand` asked two questions about a node before it let the descent + /// continue, and the Substrait one came first: unless the *whole* root plan encoded, every + /// subtree that failed to encode was expanded on, whichever way `check_plan` answered. A + /// `FusedSeries` node has no serializer branch — `ExtensionSerializer` answers `NotImplemented` + /// for every name it does not list — so over the wire the boundary is still drawn at the node + /// and the region receives the bare scan. /// - /// The commutativity arm above is necessary and not sufficient. This test fails the day the - /// node gains a serializer branch (or the analyzer stops asking), which is when the placement - /// assertion this one replaces can be written. + /// A plan handed to `RegionServer::handle_read` as a `LogicalPlan` is never encoded, so there + /// the gate has nothing to say and the commutativity arm decides: the node rides inside the + /// merge scan. #[test] - fn fused_series_does_not_encode_so_the_analyzer_keeps_it_on_the_frontend() { + fn fused_series_rides_to_the_region_only_when_regions_are_reached_in_process() { let plan = fused_series_plan(table_scan_plan()); assert!( DFLogicalSubstraitConvertor .encode(&plan, DefaultSerializer) .is_err(), - "the fused node encodes now, so it can be placed inside the merge scan" + "the fused node encodes now, so the encoding gate decides nothing" + ); + assert!( + DFLogicalSubstraitConvertor + .encode(&table_scan_plan(), DefaultSerializer) + .is_ok(), + "the scan under the fused node encodes, so the fused node is the only reason this plan does not" ); - let analyzed = DistPlannerAnalyzer {} - .analyze(plan, &ConfigOptions::default()) - .unwrap(); - let mut remote_input_holds_fused = false; - analyzed - .apply(|node| { - if let LogicalPlan::Extension(extension) = node - && let Some(merge_scan) = extension - .node - .as_any() - .downcast_ref::() - { - remote_input_holds_fused = holds_fused_series(merge_scan.input()); - } - Ok(TreeNodeRecursion::Continue) - }) - .unwrap(); + // Over the wire the plan has to encode, so the boundary is drawn at the node. + let encoded = analyze(DistPlannerAnalyzer::new(false), plan.clone()); assert!( - !remote_input_holds_fused, - "the analyzer placed the node despite the encoding gate: {analyzed}" + !any_remote_input_holds(&encoded, holds_fused_series), + "the analyzer placed the node although its plan has to encode: {encoded}" + ); + + // In process nothing is encoded, so the categorizer places the node. + let in_process = analyze(DistPlannerAnalyzer::new(true), plan); + assert!( + any_remote_input_holds(&in_process, holds_fused_series), + "the analyzer kept the node on the frontend although regions are reached in process: {in_process}" + ); + } + + /// Skipping the encoding gate removes one reason to expand and adds none: a node that does not + /// encode for its own reasons, and that the categorizer refuses, still draws the boundary at + /// itself when regions are reached in process. + /// + /// Without this, the skip would amount to placing every unencodable node region-side. + #[test] + fn a_plan_that_does_not_encode_for_another_reason_still_expands_at_that_node() { + let plan = unserializable_plan(table_scan_plan()); + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "the stand-in node has a serializer branch now, so it no longer proves anything" + ); + + let analyzed = analyze(DistPlannerAnalyzer::new(true), plan); + + assert!( + !any_remote_input_holds(&analyzed, holds_unserializable_node), + "the analyzer placed a node the categorizer refuses: {analyzed}" + ); + let inputs = analyzed.inputs(); + let [input] = inputs.as_slice() else { + panic!("expected the node to keep one input in: {analyzed}"); + }; + // `expand` keeps a projection over the merge scan to restore the schema of what it wrapped. + let merge_scan = first_merge_scan(input) + .unwrap_or_else(|| panic!("expected the boundary right below the node in: {analyzed}")); + assert!( + holds_table_scan(merge_scan.input()), + "expected the scan under the node to be the region input in: {analyzed}" ); } } diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index c512b9c1e5e..1c773ad099e 100644 --- a/src/query/src/query_engine.rs +++ b/src/query/src/query_engine.rs @@ -189,7 +189,7 @@ impl QueryEngineFactory { plugins: Plugins, options: QueryOptions, ) -> datafusion::error::Result { - let state = Arc::new(QueryEngineState::try_new( + Self::try_new_with_plugins_and_region_transport( catalog_manager, partition_rule_manager, region_query_handler, @@ -197,6 +197,40 @@ impl QueryEngineFactory { procedure_service_handler, flow_service_handler, with_dist_planner, + false, + plugins, + options, + ) + } + + /// [`QueryEngineFactory::try_new_with_plugins`], additionally stating how the engine reaches + /// the regions it queries. + /// + /// `in_process_region_transport` reaches `DistPlannerAnalyzer`: when it is `true`, a region + /// input may travel as a `LogicalPlan` instead of Substrait bytes, which only a region server in + /// the same process can accept. + #[allow(clippy::too_many_arguments)] + pub fn try_new_with_plugins_and_region_transport( + catalog_manager: CatalogManagerRef, + partition_rule_manager: Option, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + in_process_region_transport: bool, + plugins: Plugins, + options: QueryOptions, + ) -> datafusion::error::Result { + let state = Arc::new(QueryEngineState::try_new_with_region_transport( + catalog_manager, + partition_rule_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + in_process_region_transport, plugins.clone(), options, )?); diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index e2095c33530..3ee6d2db3af 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -147,6 +147,39 @@ impl QueryEngineState { with_dist_planner: bool, plugins: Plugins, options: QueryOptionsNew, + ) -> DfResult { + Self::try_new_with_region_transport( + catalog_list, + partition_rule_manager, + region_query_handler, + table_mutation_handler, + procedure_service_handler, + flow_service_handler, + with_dist_planner, + false, + plugins, + options, + ) + } + + /// [`QueryEngineState::try_new`], additionally stating how the engine reaches its regions. + /// + /// `in_process_region_transport` reaches [`DistPlannerAnalyzer`]: when it is `true` the + /// distributed planner may hand a region a `LogicalPlan` instead of Substrait bytes, which only + /// the in-process read path accepts. Every other caller builds an engine whose regions are + /// reached over the wire, which is what [`QueryEngineState::try_new`] states. + #[allow(clippy::too_many_arguments)] + pub(crate) fn try_new_with_region_transport( + catalog_list: CatalogManagerRef, + partition_rule_manager: Option, + region_query_handler: Option, + table_mutation_handler: Option, + procedure_service_handler: Option, + flow_service_handler: Option, + with_dist_planner: bool, + in_process_region_transport: bool, + plugins: Plugins, + options: QueryOptionsNew, ) -> DfResult { let total_memory = get_total_memory_bytes().max(0) as u64; let memory_pool_size = options.memory_pool_size.resolve(total_memory) as usize; @@ -209,7 +242,9 @@ impl QueryEngineState { )), ); if with_dist_planner { - analyzer.rules.push(Arc::new(DistPlannerAnalyzer)); + analyzer.rules.push(Arc::new(DistPlannerAnalyzer::new( + in_process_region_transport, + ))); analyzer.rules.push(Arc::new(JsonSchemaConcretizeRule)); } analyzer.rules.push(Arc::new(FixStateUdafOrderingAnalyzer)); From 612160078912ebbcc54483149b9b5b2ab6fbb8f9 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 12:10:56 +0800 Subject: [PATCH 05/38] fix(query): keep ordinary plans on the encoding path Signed-off-by: discord9 --- src/query/src/dist_plan/analyzer.rs | 41 ++++++------ src/query/src/planner.rs | 3 +- src/query/src/promql/planner.rs | 5 +- src/query/src/promql/planner/native.rs | 86 +++++++++++++------------- src/query/src/promql/planner/test.rs | 38 ++++++++++++ src/query/src/query_engine/state.rs | 6 ++ 6 files changed, 113 insertions(+), 66 deletions(-) diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 955cc6cbe55..76020bf262b 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -37,7 +37,7 @@ use datafusion_optimizer::push_down_filter::PushDownFilter; use datafusion_optimizer::rewrite_set_comparison::RewriteSetComparison; use datafusion_optimizer::scalar_subquery_to_join::ScalarSubqueryToJoin; use datafusion_optimizer::simplify_expressions::SimplifyExpressions; -use promql::extension_plan::SeriesDivide; +use promql::extension_plan::{FusedSeries, SeriesDivide}; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::metadata::TableType; use table::table::adapter::DfTableProviderAdapter; @@ -113,9 +113,9 @@ pub struct DistPlannerAnalyzer { /// Whether every region a rewritten plan reads is reached in process. /// /// A region input normally has to be encoded to Substrait before it can leave the frontend, and - /// [`PlanRewriter::should_expand`] cuts the plan at the first node that does not encode. When - /// regions are reached in process, that input travels as a `LogicalPlan` instead, so nothing is - /// encoded and no node has to be cut for encoding's sake. + /// [`PlanRewriter::should_expand`] cuts the plan at the first node that does not encode. In + /// process, a root `FusedSeries` plan travels as a `LogicalPlan` instead and skips that gate; + /// other plans retain the ordinary encoding checks. /// /// Whoever builds the query engine sets this, once, from the transport it wired. It is neither a /// query-time decision nor a property of a region. @@ -445,12 +445,11 @@ struct PlanRewriter { /// Whether the whole plan this rewriter walks can be encoded to Substrait. /// Used by [`PlanRewriter::should_expand`] to skip the per-node encoding check. /// - /// Only read when that gate applies, i.e. when `in_process_region_transport` is false. With - /// in-process regions the question has no answer worth computing and the field stays `false`. + /// Used for the ordinary standalone fast path when the root plan is encodable. It is not + /// needed when the root itself is the in-process `FusedSeries` exception. whole_plan_encodable: bool, - /// Whether the regions this plan reads are reached in process, so an input crosses to a region - /// as a `LogicalPlan` rather than as Substrait bytes. See [`DistPlannerAnalyzer`]. - in_process_region_transport: bool, + /// In-process transport skips the encodability gate only when the rewritten root itself is FusedSeries. + skip_encodability_gate: bool, /// Current level in the tree level: usize, /// Simulated stack for the `rewrite` recursion @@ -506,19 +505,23 @@ struct PlanRewriter { new_child_plan: Option, } +fn is_fused_series_plan(plan: &LogicalPlan) -> bool { + matches!(plan, LogicalPlan::Extension(extension) if extension.node.as_any().is::()) +} + impl PlanRewriter { /// `plan` must be the root of the tree that is about to be rewritten, otherwise /// [`PlanRewriter::should_expand`] may skip a check it has to perform. fn new(plan: &LogicalPlan, in_process_region_transport: bool) -> Self { + let skip_encodability_gate = in_process_region_transport && is_fused_series_plan(plan); Self { - // Only a plan that crosses to a region as Substrait bytes has to encode at all, and - // encoding is the expensive part of a pass. With in-process regions the answer is never - // read, so it is not computed. - whole_plan_encodable: !in_process_region_transport + // Preserve the ordinary whole-plan fast path unless this is the in-process root + // exception: no node in that tree needs an encoding check. + skip_encodability_gate, + whole_plan_encodable: !skip_encodability_gate && DFLogicalSubstraitConvertor .encode(plan, DefaultSerializer) .is_ok(), - in_process_region_transport, ..Default::default() } } @@ -551,15 +554,13 @@ impl PlanRewriter { // When the root does not encode, the plan holds at least one node that has to stay // on the frontend and only the per-node check locates it. // - // This question only exists for a plan that has to cross to a region as Substrait bytes. - // With in-process regions the same `LogicalPlan` is handed to `RegionServer::handle_read` - // never having been encoded, so encodability says nothing about where a node may run, and - // asking it would cut the plan at the first node with no serializer branch (`FusedSeries`) - // before `check_plan` could decide. Skipping it only removes a reason to expand: every + // For the in-process root `FusedSeries` exception, encodability says nothing about where + // a node may run, and asking it would cut the plan at the first node with no serializer + // branch before `check_plan` could decide. Skipping it only removes a reason to expand: every // check below still runs, and a node they call `NonCommutative`/`Unimplemented`/ // `Unsupported` still expands, so a node no serializer lists only rides to a region when // the categorizer places it there. - if !self.in_process_region_transport + if !self.skip_encodability_gate && !self.whole_plan_encodable && let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer) { diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 861f923d74b..c11517841e0 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -345,7 +345,8 @@ impl DfLogicalPlanner { stmt, &self.engine_state, promql_annotations, - crate::promql::planner::native::native_promql_requested(&query_ctx), + self.engine_state.experimental_native_promql() + && crate::promql::planner::native::native_promql_requested(&query_ctx), ) .await .map_err(BoxedError::new) diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5df731d6a95..7fba4db8816 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -300,9 +300,8 @@ pub struct PromPlanner { /// `false` when the query says nothing: with it off, planning takes the existing path unchanged /// and nothing below is reached. /// - /// The engine-wide `QueryOptions::experimental_native_promql` is not the source. It is consumed - /// while `QueryEngineState` is built and not kept, so the planner cannot read it; the extension - /// is what a client sets for one request. + /// The engine-wide `QueryOptions::experimental_native_promql` is applied while `QueryEngineState` + /// is built; this planner reads the per-query request extension instead. native_fused: bool, } diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 21cd4553d15..df97d4241a9 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -32,7 +32,6 @@ use std::time::{Duration, UNIX_EPOCH}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; -use datafusion::datasource::DefaultTableSource; use datafusion::error::DataFusionError; use datafusion::logical_expr::{Extension, LogicalPlan}; use datatypes::series::SeriesOrdering; @@ -44,9 +43,7 @@ use promql_parser::parser::{ }; use session::context::QueryContextRef; use snafu::ResultExt; -use store_api::storage::TimeSeriesDistribution; -use crate::dummy_catalog::DummyTableProvider; use crate::promql::error::{DataFusionPlanningSnafu, Result}; use crate::promql::native_compile::{compile, statement_grid}; use crate::promql::planner::PromPlanner; @@ -57,8 +54,8 @@ use crate::promql::planner::PromPlanner; /// and read exactly the way the other per-query override in this crate is read /// ([`remote_dyn_filter_pushdown_enabled_from_extensions`]): a query-context extension, which a /// client can set per request over the HTTP hints header or as a gRPC context key. The engine-wide -/// option is a query-engine setting that `QueryEngineState` consumes while it is built and does not -/// keep, so it is not readable from the planner's call site; see [`native_promql_requested`]. +/// option is applied while building query-engine state; this planner reads the per-query request +/// extension at the call site; see [`native_promql_requested`]. /// /// [`QueryOptions::experimental_native_promql`]: crate::options::QueryOptions::experimental_native_promql /// [`remote_dyn_filter_pushdown_enabled_from_extensions`]: @@ -160,7 +157,7 @@ impl PromPlanner { // physical planner turns it into the region scan the node reads, and building a scan here // instead would let the two paths disagree about the columns they read. let table_ref = self.table_ref()?; - let input = self.create_table_scan_plan(table_ref.clone()).await?; + let mut input = self.create_table_scan_plan(table_ref.clone()).await?; // The node reads the region scan itself, so the plan below it has to *be* that scan: a // filter, a projection or a `__tsid` rewrite in between produces a different physical // child, and the physical layer cannot reach through it. Declining here, while the logical @@ -168,17 +165,39 @@ impl PromPlanner { let LogicalPlan::TableScan(_) = &input else { return Ok(None); }; + let offset = selector + .offset + .as_ref() + .map(|offset| match offset { + promql_parser::parser::Offset::Pos(duration) => duration.as_millis() as i64, + promql_parser::parser::Offset::Neg(duration) => -(duration.as_millis() as i64), + }) + .unwrap_or_default(); + let (saved_start, saved_end, saved_range) = (self.ctx.start, self.ctx.end, self.ctx.range); + self.ctx.range = selector_range_millis(prom_expr); + // Keep the fused scan's bounds identical to the legacy selector scan, including @ and offset. + let time_filter = (|| { + let anchored = self.at_modifier_offset(&selector.at, &selector.offset)?; + if anchored.is_some() { + self.ctx.end = saved_start; + } + self.build_time_index_filter(anchored.unwrap_or(offset), &input.schema()) + })(); + self.ctx.start = saved_start; + self.ctx.end = saved_end; + self.ctx.range = saved_range; + let time_filter = time_filter?; + if let Some(filter) = time_filter { + let LogicalPlan::TableScan(scan) = &mut input else { + unreachable!(); + }; + scan.filters.push(filter); + } // A window over a native histogram yields a histogram, which the output schema below // cannot declare; the existing path evaluates those with its own kernels. if Self::field_columns_contain_native_histogram(input.schema(), &fields) { return Ok(None); } - // A raw series frame is produced by the region's series scan, which a region scan only - // prepares when the request asks for the per-series distribution. Without the hint the - // region scans by row and has no raw frame to hand over, so the node would be planned with - // an input that cannot feed it. - let input = with_per_series_distribution(input); - let Some(stmt) = self.native_statement(prom_expr) else { return Ok(None); }; @@ -293,42 +312,25 @@ fn fused_output_schema( )) } -/// The selector's own scan, asking the region for the per-series distribution. -/// -/// The distribution is a hint on the *table scan*, which the region server reads when it builds -/// the region scan: `PerSeries` is what makes a region prepare its series scan, whose raw-series -/// entrypoint is the only one that packs a series' chunks into the frame this node reads. The -/// existing planner asks for the same hint indirectly, through a sort on `(series key, time)` that -/// [`ScanHintRule`](crate::optimizer::scan_hint::ScanHintRule) turns into this distribution; the -/// native node has no such sort below it, because it reads raw samples rather than sorted rows, so -/// the hint is attached here. -/// -/// The provider is cloned for this use-site first, exactly as that rule does: a catalog provider can -/// be shared by several scans, and a hint attached to the shared one would change their requests -/// too. A table whose provider is not the region-server provider - a test fixture, or an engine -/// this route cannot read - is left untouched, and the physical route declines it there. -fn with_per_series_distribution(plan: LogicalPlan) -> LogicalPlan { - let LogicalPlan::TableScan(mut scan) = plan else { - return plan; - }; - let Some(source) = scan.source.downcast_ref::() else { - return LogicalPlan::TableScan(scan); - }; - let Some(provider) = source.table_provider.downcast_ref::() else { - return LogicalPlan::TableScan(scan); - }; - let adapter = provider.clone_for_scan(); - adapter.with_distribution(TimeSeriesDistribution::PerSeries); - scan.source = Arc::new(DefaultTableSource::new(Arc::new(adapter))); - LogicalPlan::TableScan(scan) -} - /// The one selector `expr` folds, if it names exactly one. /// /// The walk follows the grammar the native compiler accepts — a call, a range selector, a selector /// and parentheses — and stops at every other shape, because a shape the compiler cannot compile /// needs no selector resolved. Two distinct selectors are not one input, and a tree built over the /// wrong one of them would evaluate a window the query did not write. +fn selector_range_millis(expr: &PromExpr) -> Option { + match expr { + PromExpr::MatrixSelector(matrix) => Some(matrix.range.as_millis() as i64), + PromExpr::Paren(paren) => selector_range_millis(&paren.expr), + PromExpr::Call(call) => call + .args + .args + .iter() + .find_map(|arg| selector_range_millis(arg)), + _ => None, + } +} + fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { let mut selectors = Vec::new(); collect_selectors(expr, &mut selectors); diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 9acdf56b895..0f782f909e4 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1693,6 +1693,44 @@ async fn native_fused_flag_off_keeps_the_existing_plan() { assert!(display.contains("PromRangeManipulate"), "{display}"); } +/// The native selector scan is bounded to the range selector window, not the table retention. +#[tokio::test] +async fn native_fused_rate_scan_has_bounded_time_filter() { + let plan = build_native_fused_plan("rate(some_metric[5m])", true).await; + let display = plan.display_indent().to_string(); + + assert!( + display.contains("timestamp >= TimestampMillisecond(-299999, None)") + && display.contains("timestamp <= TimestampMillisecond(100000000, None)"), + "native scan should include the bounded five-minute statement window: {display}" + ); +} + +/// An anchored native range scan uses the anchor, while preserving the outer statement window. +#[tokio::test] +async fn native_fused_offset_rate_scan_uses_offset_window_bounds() { + let plan = build_native_fused_plan("rate(some_metric[5m] offset 10s)", true).await; + let display = plan.display_indent().to_string(); + + assert!( + display.contains("timestamp >= TimestampMillisecond(-309999, None)") + && display.contains("timestamp <= TimestampMillisecond(99990000, None)"), + "native scan should account for the selector offset: {display}" + ); +} + +#[tokio::test] +async fn native_fused_anchored_rate_scan_uses_anchor_bounds() { + let plan = build_native_fused_plan("rate(some_metric[5m] @ 300)", true).await; + let display = plan.display_indent().to_string(); + + assert!( + display.contains("timestamp >= TimestampMillisecond(1, None)") + && display.contains("timestamp <= TimestampMillisecond(300000, None)"), + "native scan should cover the five-minute window ending at the anchor: {display}" + ); +} + #[tokio::test] async fn single_abs() { do_single_instant_function_call("abs", "abs").await; diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 3ee6d2db3af..80e7574bd13 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -99,6 +99,7 @@ pub struct QueryEngineState { table_functions: Arc>>>, extension_rules: Vec>, plugins: Plugins, + experimental_native_promql: bool, } impl fmt::Debug for QueryEngineState { @@ -332,6 +333,7 @@ impl QueryEngineState { table_functions: Arc::new(RwLock::new(HashMap::new())), extension_rules, plugins, + experimental_native_promql: options.experimental_native_promql, scalar_functions: Arc::new(RwLock::new(HashMap::new())), }) } @@ -480,6 +482,10 @@ impl QueryEngineState { self.df_context.register_udwf(func); } + pub(crate) fn experimental_native_promql(&self) -> bool { + self.experimental_native_promql + } + pub fn catalog_manager(&self) -> &CatalogManagerRef { &self.catalog_manager } From 64f039eaf4bd506f1a91cdbee8a6bf45b6e40d46 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 12:23:56 +0800 Subject: [PATCH 06/38] test(query): run the encoding checks for both transports Signed-off-by: discord9 --- src/query/src/dist_plan/analyzer/test.rs | 90 +++++++++++++++++------- 1 file changed, 65 insertions(+), 25 deletions(-) diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index c0a4683b1c9..b512655dd2c 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -74,6 +74,32 @@ fn analyzer() -> DistPlannerAnalyzer { DistPlannerAnalyzer::default() } +#[test] +fn ordinary_encodable_plan_keeps_in_process_encodability_gate() { + let table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .build() + .unwrap(); + + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_ok() + ); + + let in_process_rewriter = PlanRewriter::new(&plan, true); + assert!(in_process_rewriter.whole_plan_encodable); + assert!(!in_process_rewriter.skip_encodability_gate); + + let wire_rewriter = PlanRewriter::new(&plan, false); + assert!(wire_rewriter.whole_plan_encodable); + assert!(!wire_rewriter.skip_encodability_gate); +} + struct MergeScanRemoteDynFilterProducerIdCollector<'a> { producer_ids: &'a mut Vec, } @@ -401,22 +427,29 @@ fn frontend_only_histogram_folds_stay_above_merge_scan() { "{name}" ); - let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap(); - let result_text = result.to_string(); - assert!( - result_text.contains("HistogramFold:"), - "{name}\n{result_text}" - ); + for analyzer in [ + DistPlannerAnalyzer::default(), + DistPlannerAnalyzer::new(true), + ] { + let result = analyzer + .analyze(plan.clone(), &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!( + result_text.contains("HistogramFold:"), + "{name}\n{result_text}" + ); - let remote_input = find_merge_scan(&result).unwrap().input(); - let remote_text = remote_input.to_string(); - assert!( - !remote_text.contains("HistogramFold:"), - "{name}\n{remote_text}" - ); - DFLogicalSubstraitConvertor - .encode(remote_input, DefaultSerializer) - .unwrap(); + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!( + !remote_text.contains("HistogramFold:"), + "{name}\n{remote_text}" + ); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + } } } @@ -446,17 +479,24 @@ fn unencodable_unnest_stays_above_merge_scan() { .is_err() ); - let result = analyzer().analyze(plan, &ConfigOptions::default()).unwrap(); - let result_text = result.to_string(); - assert!(result_text.contains("Unnest:"), "{result_text}"); + for analyzer in [ + DistPlannerAnalyzer::default(), + DistPlannerAnalyzer::new(true), + ] { + let result = analyzer + .analyze(plan.clone(), &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!(result_text.contains("Unnest:"), "{result_text}"); - let remote_input = find_merge_scan(&result).unwrap().input(); - let remote_text = remote_input.to_string(); - assert!(!remote_text.contains("Unnest:"), "{remote_text}"); - assert!(remote_text.contains("make_array"), "{remote_text}"); - DFLogicalSubstraitConvertor - .encode(remote_input, DefaultSerializer) - .unwrap(); + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!(!remote_text.contains("Unnest:"), "{remote_text}"); + assert!(remote_text.contains("make_array"), "{remote_text}"); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + } } #[test] From 67a3d22c5b997568b31b3659e30093bd9f4cd3cf Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 12:32:22 +0800 Subject: [PATCH 07/38] fix(query): keep a fused route to series a region owns Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 54 +++++++++++++++++++++++++- src/query/src/promql/planner/test.rs | 51 +++++++++++++++++++++++- 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index df97d4241a9..7f3eff70458 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -162,9 +162,19 @@ impl PromPlanner { // filter, a projection or a `__tsid` rewrite in between produces a different physical // child, and the physical layer cannot reach through it. Declining here, while the logical // plan is still this method's to choose, is what keeps such a query on the existing path. - let LogicalPlan::TableScan(_) = &input else { + let LogicalPlan::TableScan(scan) = &input else { return Ok(None); }; + let table = self.table_from_source(&scan.source)?; + // FusedSeries evaluates each region independently, so partition keys must identify a series. + if !table + .table_info() + .meta + .partition_column_names() + .all(|column| tags.iter().any(|tag| tag == column)) + { + return Ok(None); + } let offset = selector .offset .as_ref() @@ -353,3 +363,45 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>) _ => {} } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use session::context::QueryContextBuilder; + + use super::{NATIVE_PROMQL_EXTENSION_KEY, native_promql_requested}; + + #[test] + fn native_promql_requested_is_false_when_extension_is_absent() { + let query_ctx = Arc::new(QueryContextBuilder::default().build()); + + assert!(!native_promql_requested(&query_ctx)); + } + + #[test] + fn native_promql_requested_accepts_true_case_insensitively() { + for value in ["true", "TRUE", "TrUe"] { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), value.to_string()) + .build(), + ); + + assert!(native_promql_requested(&query_ctx), "value: {value}"); + } + } + + #[test] + fn native_promql_requested_rejects_false_and_unrecognized_values() { + for value in ["false", "FALSE", "yes", "1", " true ", ""] { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), value.to_string()) + .build(), + ); + + assert!(!native_promql_requested(&query_ctx), "value: {value:?}"); + } + } +} diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 0f782f909e4..7ed504c42d0 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -898,6 +898,16 @@ async fn build_test_table_provider( table_name_tuples: &[(String, String)], num_tag: usize, num_field: usize, +) -> DfTableSourceProvider { + build_test_table_provider_with_partition_columns(table_name_tuples, num_tag, num_field, vec![]) + .await +} + +async fn build_test_table_provider_with_partition_columns( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, + partition_key_indices: Vec, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); for (schema_name, table_name) in table_name_tuples { @@ -928,6 +938,7 @@ async fn build_test_table_provider( let table_meta = TableMetaBuilder::empty() .schema(schema) .primary_key_indices((0..num_tag).collect()) + .partition_key_indices(partition_key_indices.clone()) .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) .next_column_id(1024) .build() @@ -1626,13 +1637,21 @@ async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) /// time index `timestamp`, one field `field_0` — which is the shape the selector context resolves /// to a raw series schema of one field. async fn build_native_fused_plan(query: &str, native_fused: bool) -> LogicalPlan { - let eval_stmt = build_eval_stmt(query); let table_provider = build_test_table_provider( &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], 1, 1, ) .await; + build_native_fused_plan_with_provider(query, native_fused, table_provider).await +} + +async fn build_native_fused_plan_with_provider( + query: &str, + native_fused: bool, + table_provider: DfTableSourceProvider, +) -> LogicalPlan { + let eval_stmt = build_eval_stmt(query); PromPlanner::stmt_to_plan_with_annotations( table_provider, &eval_stmt, @@ -1671,6 +1690,36 @@ async fn native_fused_rate_plans_to_a_fused_series_node() { assert_eq!("rate(some_metric[5m])", fused.promql()); } +#[tokio::test] +async fn native_fused_time_partitioned_table_falls_back() { + let table_provider = build_test_table_provider_with_partition_columns( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![1], // timestamp + ) + .await; + let plan = + build_native_fused_plan_with_provider("rate(some_metric[2m])", true, table_provider).await; + let display = plan.display_indent().to_string(); + assert!(!display.contains("FusedSeries"), "{display}"); + assert!(display.contains("PromRangeManipulate"), "{display}"); +} + +#[tokio::test] +async fn native_fused_tag_partitioned_table_still_plans_fused_series() { + let table_provider = build_test_table_provider_with_partition_columns( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![0], // tag_0 + ) + .await; + let plan = + build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; + assert!(plan.display_indent().to_string().contains("FusedSeries")); +} + /// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. #[tokio::test] async fn native_fused_unsupported_shape_falls_back() { From 0d29dc98e85ae160163595792b7f19054962e692 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 13:09:26 +0800 Subject: [PATCH 08/38] fix(query): emit the fused output in the existing column order Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 20 ++++++++++---------- src/query/src/promql/planner/test.rs | 8 ++++++++ 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 7f3eff70458..7baff0abdfe 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -23,8 +23,8 @@ //! What the node needs from here is three things the existing planner already resolves: the raw //! series schema the tree's input was compiled against (the selector's selected field columns), the //! statement's evaluation grid, and the logical tag names the output carries. The output schema is -//! built to the shape the existing unary pipeline produces — the series fields, then the tags, then -//! the time index — so a node planned here can be consumed by the same code that consumes a +//! built to the shape the existing unary pipeline produces — the time index, then the series +//! fields, then the tags — so a node planned here can be consumed by the same code that consumes a //! `SeriesNormalize` output. use std::sync::Arc; @@ -298,6 +298,14 @@ fn fused_output_schema( time_index: &str, ) -> Result { let mut columns = Vec::with_capacity(fields.len() + tags.len() + 1); + columns.push(( + None, + Arc::new(Field::new( + time_index, + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )), + )); for name in fields { columns.push(( None, @@ -307,14 +315,6 @@ fn fused_output_schema( for name in tags { columns.push((None, Arc::new(Field::new(name, ArrowDataType::Utf8, true)))); } - columns.push(( - None, - Arc::new(Field::new( - time_index, - ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), - false, - )), - )); Ok(Arc::new( DFSchema::new_with_metadata(columns, Default::default()) diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 7ed504c42d0..9464e900ad0 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1688,6 +1688,14 @@ async fn native_fused_rate_plans_to_a_fused_series_node() { // The node carries the query a reader wrote, not the shape it compiled to. let fused = find_fused_series(&plan).expect("the plan holds a native node"); assert_eq!("rate(some_metric[5m])", fused.promql()); + assert_eq!( + vec!["timestamp", "field_0", "tag_0"], + plan.schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>() + ); } #[tokio::test] From 129faa4e49c8ff027ee5609e2031e3dd9e9b226e Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 13:30:32 +0800 Subject: [PATCH 09/38] chore(query): drop a needless borrow on the scan schema Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 7baff0abdfe..1b9106aa7bf 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -191,7 +191,7 @@ impl PromPlanner { if anchored.is_some() { self.ctx.end = saved_start; } - self.build_time_index_filter(anchored.unwrap_or(offset), &input.schema()) + self.build_time_index_filter(anchored.unwrap_or(offset), input.schema()) })(); self.ctx.start = saved_start; self.ctx.end = saved_end; From 5527c751d2f98e96953f2914bdeb5c00913aa99a Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 14:29:34 +0800 Subject: [PATCH 10/38] fix(query): decline the native route where no raw scan exists Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 16 ++++ src/query/src/promql/planner/test.rs | 102 +++++++++++++++++++++++-- 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 1b9106aa7bf..7dd6c85cd9d 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -166,6 +166,22 @@ impl PromPlanner { return Ok(None); }; let table = self.table_from_source(&scan.source)?; + let table_meta = &table.table_info().meta; + // The region's raw-series preflight refuses these, and the physical layer turns that refusal into an error, so decline this route here instead. + let time_index_type = table_meta + .schema + .timestamp_column() + .map(|column| &column.data_type); + if time_index_type + != Some(&datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype()) + || table_meta + .options + .extra_options + .get(store_api::mito_engine_options::APPEND_MODE_KEY) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) + { + return Ok(None); + } // FusedSeries evaluates each region independently, so partition keys must identify a series. if !table .table_info() diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 9464e900ad0..b9dd3f1de8e 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -908,6 +908,25 @@ async fn build_test_table_provider_with_partition_columns( num_tag: usize, num_field: usize, partition_key_indices: Vec, +) -> DfTableSourceProvider { + build_test_table_provider_with_options( + table_name_tuples, + num_tag, + num_field, + partition_key_indices, + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .await +} + +async fn build_test_table_provider_with_options( + table_name_tuples: &[(String, String)], + num_tag: usize, + num_field: usize, + partition_key_indices: Vec, + timestamp_type: ConcreteDataType, + append_mode: bool, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); for (schema_name, table_name) in table_name_tuples { @@ -920,12 +939,8 @@ async fn build_test_table_provider_with_partition_columns( )); } columns.push( - ColumnSchema::new( - "timestamp".to_string(), - ConcreteDataType::timestamp_millisecond_datatype(), - false, - ) - .with_time_index(true), + ColumnSchema::new("timestamp".to_string(), timestamp_type.clone(), false) + .with_time_index(true), ); for i in 0..num_field { columns.push(ColumnSchema::new( @@ -941,6 +956,17 @@ async fn build_test_table_provider_with_partition_columns( .partition_key_indices(partition_key_indices.clone()) .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) .next_column_id(1024) + .options(table::requests::TableOptions { + extra_options: if append_mode { + HashMap::from([( + store_api::mito_engine_options::APPEND_MODE_KEY.to_string(), + "true".to_string(), + )]) + } else { + HashMap::new() + }, + ..Default::default() + }) .build() .unwrap(); let table_info = TableInfoBuilder::default() @@ -1714,6 +1740,70 @@ async fn native_fused_time_partitioned_table_falls_back() { assert!(display.contains("PromRangeManipulate"), "{display}"); } +#[tokio::test] +async fn native_fused_non_millisecond_time_index_falls_back() { + let table_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_microsecond_datatype(), + false, + ) + .await; + let native = + build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; + let fallback_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_microsecond_datatype(), + false, + ) + .await; + let without_flag = + build_native_fused_plan_with_provider("rate(some_metric[5m])", false, fallback_provider) + .await; + assert!(!native.display_indent().to_string().contains("FusedSeries")); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string() + ); +} + +#[tokio::test] +async fn native_fused_append_mode_table_falls_back() { + let table_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ) + .await; + let native = + build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; + let fallback_provider = build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ) + .await; + let without_flag = + build_native_fused_plan_with_provider("rate(some_metric[5m])", false, fallback_provider) + .await; + assert!(!native.display_indent().to_string().contains("FusedSeries")); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string() + ); +} + #[tokio::test] async fn native_fused_tag_partitioned_table_still_plans_fused_series() { let table_provider = build_test_table_provider_with_partition_columns( From 8efec617433f5bc0b290dbcfad94639eb96f2410 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 17:31:51 +0800 Subject: [PATCH 11/38] fix(promql): keep the fused scan's parallelism Signed-off-by: discord9 --- src/promql/src/extension_plan/planner.rs | 44 +++-- src/query/src/optimizer/parallelize_scan.rs | 171 +------------------- src/table/src/table/scan.rs | 100 +++++++++++- 3 files changed, 119 insertions(+), 196 deletions(-) diff --git a/src/promql/src/extension_plan/planner.rs b/src/promql/src/extension_plan/planner.rs index f2d86b76862..701ca4be01d 100644 --- a/src/promql/src/extension_plan/planner.rs +++ b/src/promql/src/extension_plan/planner.rs @@ -67,7 +67,11 @@ impl ExtensionPlanner for PromExtensionPlanner { } else if let Some(node) = node.as_any().downcast_ref::() { Ok(Some(node.to_execution_plan(physical_inputs[0].clone()))) } else if let Some(node) = node.as_any().downcast_ref::() { - plan_fused_series(node, physical_inputs) + plan_fused_series( + node, + physical_inputs, + session.config_options().execution.target_partitions, + ) } else { Ok(None) } @@ -90,9 +94,13 @@ const METRIC_NAME: &str = "__name__"; /// found but refuses the raw path is a different case — the logical route chose this node for this /// scan and nothing here can undo that choice, because the logical plan is not available — so it is /// reported as an error that names the reason instead of silently producing no plan. +/// The fused leaf needs a prepared scan: range assignment runs after physical planning and cannot +/// see a `SeriesScanExec`. Preparation keeps the scan's parallel layout, and the series-producer +/// contract guarantees that one series stays in one scan partition. fn plan_fused_series( fused: &FusedSeries, physical_inputs: &[Arc], + target_partitions: usize, ) -> DfResult>> { let Some(input) = physical_inputs.first() else { return Ok(None); @@ -124,7 +132,18 @@ fn plan_fused_series( .map(|field| field.name.to_string()) .collect::>(); - let scan = single_partition_scan(scan)?; + let scan = if scan.is_partition_set() { + Arc::new(scan.clone()) + } else { + let partitions = table::table::scan::assign_partition_ranges( + scan.get_partition_ranges(), + target_partitions.max(1), + ); + Arc::new( + scan.with_new_partitions(partitions, target_partitions.max(1)) + .map_err(|error| DataFusionError::External(error.into_inner()))?, + ) + }; // The scan resolves the label source and the ordering it will declare itself, against its own // metadata snapshot, and hands both back through the leaf it builds. Naming them here instead // would be a second copy of a binding only the scanner can make: the physical metadata and @@ -146,24 +165,3 @@ fn plan_fused_series( .to_execution_plan(series_scan, label_source, ordering) .map(Some) } - -/// The scan as one partition holding every range it already has. -/// -/// A raw-series binding is resolvable only once the scan holds exactly one partition: the label -/// source is decoded against the scanner's own metadata snapshot and the frames it produces carry -/// whole series, so a scan spread over several partitions has no single binding to name. The rule -/// that assigns partition ranges (`ParallelizeScan`) runs *after* physical planning, i.e. after -/// this node is built, so a scan reaching here is unprepared; preparing it here keeps the node from -/// reading one partition of the series it was bound to. A scan that is already prepared — several -/// partitions, assigned by whatever ran before — is passed through as it is, and the binding below -/// declines it if it cannot name one partition. -fn single_partition_scan(scan: &RegionScanExec) -> DfResult> { - if scan.is_partition_set() { - return Ok(Arc::new(scan.clone())); - } - let ranges = scan.get_partition_ranges(); - let prepared = scan - .with_new_partitions(vec![ranges], 1) - .map_err(|error| DataFusionError::External(error.into_inner()))?; - Ok(Arc::new(prepared)) -} diff --git a/src/query/src/optimizer/parallelize_scan.rs b/src/query/src/optimizer/parallelize_scan.rs index 8e0113e69df..1bc07e34a93 100644 --- a/src/query/src/optimizer/parallelize_scan.rs +++ b/src/query/src/optimizer/parallelize_scan.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::BinaryHeap; use std::sync::Arc; use common_telemetry::debug; @@ -22,7 +21,6 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{DataFusionError, Result}; -use store_api::region_engine::PartitionRange; use table::table::scan::RegionScanExec; #[derive(Debug)] @@ -71,7 +69,7 @@ impl ParallelizeScan { // assign ranges to each partition let mut partition_ranges = - Self::assign_partition_range(ranges, expected_partition_num); + table::table::scan::assign_partition_ranges(ranges, expected_partition_num); debug!( "Assign {total_range_num} ranges to {expected_partition_num} partitions" ); @@ -116,176 +114,13 @@ impl ParallelizeScan { Ok(result) } - - /// Distribute [`PartitionRange`]s to each partition. - /// - /// Currently we assign ranges to partitions according to their rows so each partition - /// has similar number of rows. This method always return `expected_partition_num` partitions. - pub fn assign_partition_range( - mut ranges: Vec, - expected_partition_num: usize, - ) -> Vec> { - if ranges.is_empty() { - // Returns a single partition with no range. - return vec![vec![]; expected_partition_num]; - } - - if ranges.len() == 1 { - let mut vec = vec![vec![]; expected_partition_num]; - vec[0] = ranges; - return vec; - } - - // Sort ranges by number of rows in descending order. - ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows)); - let mut partition_ranges = vec![vec![]; expected_partition_num]; - - #[derive(Eq, PartialEq)] - struct HeapNode { - num_rows: usize, - partition_idx: usize, - } - - impl Ord for HeapNode { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - // Reverse for min-heap. - self.num_rows.cmp(&other.num_rows).reverse() - } - } - - impl PartialOrd for HeapNode { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } - } - - let mut part_heap = - BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode { - num_rows: 0, - partition_idx, - })); - - // Assigns the range to the partition with the smallest number of rows. - for range in ranges { - let mut node = part_heap.pop().unwrap(); - let partition_idx = node.partition_idx; - node.num_rows += range.num_rows; - partition_ranges[partition_idx].push(range); - part_heap.push(node); - } - - partition_ranges - } } #[cfg(test)] mod test { use common_time::Timestamp; use common_time::timestamp::TimeUnit; - - use super::*; - - #[test] - fn test_assign_partition_range() { - let ranges = vec![ - PartitionRange { - start: Timestamp::new(0, TimeUnit::Second), - end: Timestamp::new(10, TimeUnit::Second), - num_rows: 100, - identifier: 1, - }, - PartitionRange { - start: Timestamp::new(10, TimeUnit::Second), - end: Timestamp::new(20, TimeUnit::Second), - num_rows: 200, - identifier: 2, - }, - PartitionRange { - start: Timestamp::new(20, TimeUnit::Second), - end: Timestamp::new(30, TimeUnit::Second), - num_rows: 150, - identifier: 3, - }, - PartitionRange { - start: Timestamp::new(30, TimeUnit::Second), - end: Timestamp::new(40, TimeUnit::Second), - num_rows: 250, - identifier: 4, - }, - ]; - - // assign to 2 partitions - let expected_partition_num = 2; - let result = - ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num); - let expected = vec![ - vec![ - PartitionRange { - start: Timestamp::new(30, TimeUnit::Second), - end: Timestamp::new(40, TimeUnit::Second), - num_rows: 250, - identifier: 4, - }, - PartitionRange { - start: Timestamp::new(0, TimeUnit::Second), - end: Timestamp::new(10, TimeUnit::Second), - num_rows: 100, - identifier: 1, - }, - ], - vec![ - PartitionRange { - start: Timestamp::new(10, TimeUnit::Second), - end: Timestamp::new(20, TimeUnit::Second), - num_rows: 200, - identifier: 2, - }, - PartitionRange { - start: Timestamp::new(20, TimeUnit::Second), - end: Timestamp::new(30, TimeUnit::Second), - num_rows: 150, - identifier: 3, - }, - ], - ]; - assert_eq!(result, expected); - - // assign 4 ranges to 5 partitions. - let expected_partition_num = 5; - let result = ParallelizeScan::assign_partition_range(ranges, expected_partition_num); - let expected = vec![ - vec![PartitionRange { - start: Timestamp::new(30, TimeUnit::Second), - end: Timestamp::new(40, TimeUnit::Second), - num_rows: 250, - identifier: 4, - }], - vec![PartitionRange { - start: Timestamp::new(0, TimeUnit::Second), - end: Timestamp::new(10, TimeUnit::Second), - num_rows: 100, - identifier: 1, - }], - vec![PartitionRange { - start: Timestamp::new(10, TimeUnit::Second), - end: Timestamp::new(20, TimeUnit::Second), - num_rows: 200, - identifier: 2, - }], - vec![], - vec![PartitionRange { - start: Timestamp::new(20, TimeUnit::Second), - end: Timestamp::new(30, TimeUnit::Second), - num_rows: 150, - identifier: 3, - }], - ]; - assert_eq!(result, expected); - - // assign 0 ranges to 5 partitions. Should return 5 empty ranges. - let result = ParallelizeScan::assign_partition_range(vec![], 5); - assert_eq!(result.len(), 5); - } + use store_api::region_engine::PartitionRange; #[test] fn test_assign_unbalance_partition_range() { @@ -319,7 +154,7 @@ mod test { // assign to 2 partitions let expected_partition_num = 2; let result = - ParallelizeScan::assign_partition_range(ranges.clone(), expected_partition_num); + table::table::scan::assign_partition_ranges(ranges.clone(), expected_partition_num); let expected = vec![ vec![PartitionRange { start: Timestamp::new(30, TimeUnit::Second), diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 486b1e8d6da..38548fa0853 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{BinaryHeap, HashSet}; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; @@ -109,6 +109,66 @@ impl std::fmt::Debug for RegionScanExec { } } +/// Distribute [`PartitionRange`]s to each partition. +/// +/// Currently we assign ranges to partitions according to their rows so each partition +/// has similar number of rows. This method always return `expected_partition_num` partitions. +pub fn assign_partition_ranges( + mut ranges: Vec, + expected_partition_num: usize, +) -> Vec> { + if ranges.is_empty() { + // Returns a single partition with no range. + return vec![vec![]; expected_partition_num]; + } + + if ranges.len() == 1 { + let mut vec = vec![vec![]; expected_partition_num]; + vec[0] = ranges; + return vec; + } + + // Sort ranges by number of rows in descending order. + ranges.sort_by_key(|b| std::cmp::Reverse(b.num_rows)); + let mut partition_ranges = vec![vec![]; expected_partition_num]; + + #[derive(Eq, PartialEq)] + struct HeapNode { + num_rows: usize, + partition_idx: usize, + } + + impl Ord for HeapNode { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Reverse for min-heap. + self.num_rows.cmp(&other.num_rows).reverse() + } + } + + impl PartialOrd for HeapNode { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + let mut part_heap = + BinaryHeap::from_iter((0..expected_partition_num).map(|partition_idx| HeapNode { + num_rows: 0, + partition_idx, + })); + + // Assigns the range to the partition with the smallest number of rows. + for range in ranges { + let mut node = part_heap.pop().unwrap(); + let partition_idx = node.partition_idx; + node.num_rows += range.num_rows; + partition_ranges[partition_idx].push(range); + part_heap.push(node); + } + + partition_ranges +} + impl RegionScanExec { pub fn new( scanner: RegionScannerRef, @@ -434,6 +494,7 @@ impl RegionScanExec { /// Resolves authorized logical names against this prepared scanner's physical metadata, /// then selects a native raw-series leaf and returns the binding it was validated against. + /// Prepared scans may contain multiple partitions; the producer keeps each series in one partition. /// /// Approval must be based on the returned [`RawSeriesRequest`], because that binding, not the /// requested names, describes the approved output. See [`Self::try_to_named_series_scan`] @@ -474,10 +535,6 @@ impl RegionScanExec { } let scanner = self.scanner.lock().unwrap(); let metadata = scanner.metadata(); - if scanner.properties().num_partitions().max(1) != 1 { - return None; - } - let primary_key_ids = metadata .primary_key_columns() .map(|column| column.column_id) @@ -1302,6 +1359,39 @@ mod test { ) } + #[test] + fn assign_partition_ranges_preserves_requested_partition_count() { + let ranges = (0..4) + .map(|identifier| PartitionRange { + start: common_time::Timestamp::new( + identifier as i64, + common_time::timestamp::TimeUnit::Second, + ), + end: common_time::Timestamp::new( + identifier as i64 + 1, + common_time::timestamp::TimeUnit::Second, + ), + num_rows: 1, + identifier, + }) + .collect(); + assert_eq!(assign_partition_ranges(ranges, 4).len(), 4); + } + + #[test] + fn named_series_binding_accepts_multiple_prepared_partitions() { + let (_, plan) = dynamic_filter_fixture(5687); + let prepared = plan + .with_new_partitions(vec![vec![], vec![]], 2) + .expect("prepare two scanner partitions"); + assert!(prepared.is_partition_set()); + assert!( + prepared + .resolve_named_series_binding("cpu", &[], &["a".to_string()]) + .is_some() + ); + } + #[test] fn named_series_scan_requires_prepared_single_partition() { let (_, plan) = dynamic_filter_fixture(5686); From 1ced0f7fdb278d7433405cafbe62ed5b29b48361 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 22 Oct 2026 18:05:46 +0800 Subject: [PATCH 12/38] chore(query): share the range assignment with every caller Signed-off-by: discord9 --- src/cmd/src/datanode/scanbench.rs | 5 ++--- src/query/src/optimizer/parallelize_scan.rs | 9 ++++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 9fbde16fb7b..87a8281e170 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -47,7 +47,7 @@ use mito2::sst::file_ref::FileReferenceManager; use moka::future::CacheBuilder; use object_store::manager::ObjectStoreManager; use object_store::util::normalize_dir; -use query::optimizer::parallelize_scan::ParallelizeScan; +use query::optimizer::parallelize_scan::assign_partition_ranges; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt}; use sqlparser::ast::ExprWithAlias as SqlExprWithAlias; @@ -995,8 +995,7 @@ impl ScanbenchCommand { let all_ranges: Vec<_> = original_partitions.into_iter().flatten().collect(); // Distribute ranges across partitions - let mut partitions = - ParallelizeScan::assign_partition_range(all_ranges, self.parallelism); + let mut partitions = assign_partition_ranges(all_ranges, self.parallelism); // Sort ranges within each partition by start time ascending for partition in &mut partitions { diff --git a/src/query/src/optimizer/parallelize_scan.rs b/src/query/src/optimizer/parallelize_scan.rs index 1bc07e34a93..a0b0545d4e0 100644 --- a/src/query/src/optimizer/parallelize_scan.rs +++ b/src/query/src/optimizer/parallelize_scan.rs @@ -23,6 +23,8 @@ use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{DataFusionError, Result}; use table::table::scan::RegionScanExec; +pub use table::table::scan::assign_partition_ranges; + #[derive(Debug)] pub struct ParallelizeScan; @@ -69,7 +71,7 @@ impl ParallelizeScan { // assign ranges to each partition let mut partition_ranges = - table::table::scan::assign_partition_ranges(ranges, expected_partition_num); + assign_partition_ranges(ranges, expected_partition_num); debug!( "Assign {total_range_num} ranges to {expected_partition_num} partitions" ); @@ -122,6 +124,8 @@ mod test { use common_time::timestamp::TimeUnit; use store_api::region_engine::PartitionRange; + use super::assign_partition_ranges; + #[test] fn test_assign_unbalance_partition_range() { let ranges = vec![ @@ -153,8 +157,7 @@ mod test { // assign to 2 partitions let expected_partition_num = 2; - let result = - table::table::scan::assign_partition_ranges(ranges.clone(), expected_partition_num); + let result = assign_partition_ranges(ranges.clone(), expected_partition_num); let expected = vec![ vec![PartitionRange { start: Timestamp::new(30, TimeUnit::Second), From 638d48502477a55818f4106d56de1fa3c673f400 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 08:36:30 +0800 Subject: [PATCH 13/38] fix(query): decline the native route where it cannot reach a region Signed-off-by: discord9 --- src/query/src/dist_plan/commutativity.rs | 3 + src/query/src/planner.rs | 26 +++++++-- src/query/src/promql/planner.rs | 16 ++++++ src/query/src/promql/planner/native.rs | 35 ++++++++++++ src/query/src/promql/planner/test.rs | 71 +++++++++++++++++++++++- src/query/src/query_engine/state.rs | 9 +++ 6 files changed, 154 insertions(+), 6 deletions(-) diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index c8d56cc9a04..ce8466cd68f 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -400,6 +400,9 @@ fn reads_across_series(node: &NativeNode) -> bool { | NativeExpr::Window { input, .. } | NativeExpr::Map { input, .. } => reads_across_series(input), NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } => true, + // A vector operation reads the whole vector at each grid point (`topk`, `absent`, + // `histogram_quantile`, ...), so it is not a per-series node either. + NativeExpr::Vector { .. } => true, } } diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index c11517841e0..de45e55e8d8 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -316,7 +316,12 @@ impl DfLogicalPlanner { } #[tracing::instrument(skip_all)] - async fn plan_pql(&self, stmt: &EvalStmt, query_ctx: QueryContextRef) -> Result { + async fn plan_pql( + &self, + stmt: &EvalStmt, + query_ctx: QueryContextRef, + native_route_allowed: bool, + ) -> Result { let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector); if let Some(collector) = &promql_annotations { @@ -345,7 +350,9 @@ impl DfLogicalPlanner { stmt, &self.engine_state, promql_annotations, - self.engine_state.experimental_native_promql() + native_route_allowed + && self.engine_state.experimental_native_promql() + && self.engine_state.in_process_region_transport() && crate::promql::planner::native::native_promql_requested(&query_ctx), ) .await @@ -463,8 +470,16 @@ impl DfLogicalPlanner { alias: eval.alias.clone(), }; let stmt = QueryLanguageParser::parse_promql(&prom_query, &query_ctx)?; - - self.plan(&stmt, query_ctx).await + // A hybrid CTE registers this plan as a table the outer SQL query reads, so the + // `FusedSeries` node would sit under a SQL root. The in-process encodability + // exception is scoped to a plan whose root is the node, so here the node would be + // cut off and left in the frontend, where the physical planner cannot build it - + // a query that used to answer would fail instead. The route is therefore not + // offered for an embedded statement. + let QueryStatement::Promql(eval, _alias) = &stmt else { + unreachable!("a TQL eval parses to a PromQL statement"); + }; + self.plan_pql(eval, query_ctx, false).await } Tql::Explain(_) => UnimplementedSnafu { operation: "TQL EXPLAIN in CTEs", @@ -701,7 +716,8 @@ impl LogicalPlanner for DfLogicalPlanner { async fn plan(&self, stmt: &QueryStatement, query_ctx: QueryContextRef) -> Result { match stmt { QueryStatement::Sql(stmt) => self.plan_sql(stmt, query_ctx).await, - QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx).await, + // The statement's own plan is what a region receives, so the native route is offered. + QueryStatement::Promql(stmt, _alias) => self.plan_pql(stmt, query_ctx, true).await, } } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 7fba4db8816..2de6a87f241 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -5505,6 +5505,22 @@ impl PromPlanner { .map(|idx| schema.field(idx).data_type()) } + /// Whether every field column carries a type the region's raw series preflight serves. + /// + /// That preflight accepts numeric columns and native histograms only, so anything else - a + /// string, a boolean, a timestamp - has no raw framing. A route that offered it would fail + /// after the logical plan is fixed; declining keeps the existing plan. + fn field_columns_are_raw_series_servable( + schema: &DFSchemaRef, + field_columns: &[String], + ) -> bool { + field_columns.iter().all(|field| { + Self::field_column_type(schema, field).is_some_and(|data_type| { + data_type.is_numeric() || data_type == &Self::native_histogram_arrow_type() + }) + }) + } + fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool { Self::field_column_type(schema, field_column) .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 7dd6c85cd9d..31b083b08ea 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -224,6 +224,14 @@ impl PromPlanner { if Self::field_columns_contain_native_histogram(input.schema(), &fields) { return Ok(None); } + // The region's raw series preflight serves a field only when it is numeric or a native + // histogram, and the physical layer turns that refusal into a planning error because the + // logical plan is already fixed by then. A table whose selected fields are, say, a string + // and a boolean therefore has to be declined here, while the existing plan is still the + // fallback - `stringsfield` over such a table answered normally before this route existed. + if !Self::field_columns_are_raw_series_servable(input.schema(), &fields) { + return Ok(None); + } let Some(stmt) = self.native_statement(prom_expr) else { return Ok(None); }; @@ -250,6 +258,14 @@ impl PromPlanner { Err(DataFusionError::NotImplemented(_)) => return Ok(None), Err(error) => return Err(error).context(DataFusionPlanningSnafu), }; + // A vector operation (`sort`, `absent`, `topk`, `histogram_quantile`, ...) reads the whole + // vector at each grid point rather than one series at a time, so it cannot be evaluated by + // one executor per scan partition. Declining here is the only place that still works: once + // the `FusedSeries` node exists the logical plan is fixed, and the categorizer merely + // declining to push it down would leave a node the physical layer must build anyway. + if reads_across_series(tree.root()) { + return Ok(None); + } let binding = FusedSeriesBinding { series_schema, @@ -357,6 +373,25 @@ fn selector_range_millis(expr: &PromExpr) -> Option { } } +/// Whether any node of `tree` reads the whole input vector at a point rather than one series. +/// +/// Such a node folds or pairs across series (`Aggregate`, `Binary`) or reorders/reads a whole +/// vector (`Vector`: `sort`, `absent`, `topk`, `limitk`, `histogram_quantile`, ...). This node runs +/// one executor per scan partition, and a partition holds only part of a metric's series, so those +/// nodes would answer from a fragment. +fn reads_across_series(node: &promql::native::plan::NativeNode) -> bool { + use promql::native::plan::NativeExpr; + match &node.expr { + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, + NativeExpr::Instant { input, .. } + | NativeExpr::Window { input, .. } + | NativeExpr::Map { input, .. } => reads_across_series(input), + NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } | NativeExpr::Vector { .. } => { + true + } + } +} + fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { let mut selectors = Vec::new(); collect_selectors(expr, &mut selectors); diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index b9dd3f1de8e..c01d2d40f4f 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -916,6 +916,23 @@ async fn build_test_table_provider_with_partition_columns( partition_key_indices, ConcreteDataType::timestamp_millisecond_datatype(), false, + ConcreteDataType::float64_datatype(), + ) + .await +} + +/// A table whose single value field has the given type. +async fn build_test_table_provider_with_field_type( + field_type: ConcreteDataType, +) -> DfTableSourceProvider { + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + field_type, ) .await } @@ -927,6 +944,7 @@ async fn build_test_table_provider_with_options( partition_key_indices: Vec, timestamp_type: ConcreteDataType, append_mode: bool, + field_type: ConcreteDataType, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); for (schema_name, table_name) in table_name_tuples { @@ -945,7 +963,7 @@ async fn build_test_table_provider_with_options( for i in 0..num_field { columns.push(ColumnSchema::new( format!("field_{i}"), - ConcreteDataType::float64_datatype(), + field_type.clone(), true, )); } @@ -1749,6 +1767,7 @@ async fn native_fused_non_millisecond_time_index_falls_back() { vec![], ConcreteDataType::timestamp_microsecond_datatype(), false, + ConcreteDataType::float64_datatype(), ) .await; let native = @@ -1760,6 +1779,7 @@ async fn native_fused_non_millisecond_time_index_falls_back() { vec![], ConcreteDataType::timestamp_microsecond_datatype(), false, + ConcreteDataType::float64_datatype(), ) .await; let without_flag = @@ -1781,6 +1801,7 @@ async fn native_fused_append_mode_table_falls_back() { vec![], ConcreteDataType::timestamp_millisecond_datatype(), true, + ConcreteDataType::float64_datatype(), ) .await; let native = @@ -1792,6 +1813,7 @@ async fn native_fused_append_mode_table_falls_back() { vec![], ConcreteDataType::timestamp_millisecond_datatype(), true, + ConcreteDataType::float64_datatype(), ) .await; let without_flag = @@ -1819,6 +1841,53 @@ async fn native_fused_tag_partitioned_table_still_plans_fused_series() { } /// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. +/// A value field the region's raw scan cannot serve keeps the existing plan. +/// +/// The preflight serves numeric columns and native histograms only, and the physical layer turns +/// its refusal into a planning error once the logical plan is fixed. A bare selector over a string +/// field answered normally before this route existed, so it must keep doing that. +#[tokio::test] +async fn native_fused_non_numeric_field_falls_back() { + for field_type in [ + ConcreteDataType::string_datatype(), + ConcreteDataType::boolean_datatype(), + ] { + let native = build_native_fused_plan_with_provider( + "some_metric", + true, + build_test_table_provider_with_field_type(field_type.clone()).await, + ) + .await; + let without_flag = build_native_fused_plan_with_provider( + "some_metric", + false, + build_test_table_provider_with_field_type(field_type.clone()).await, + ) + .await; + assert!( + find_fused_series(&native).is_none(), + "{field_type} must not reach the native route" + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "{field_type} must keep the existing plan" + ); + } +} + +/// A whole-vector operation is not a per-series tree, so it stays on the existing plan. +#[tokio::test] +async fn native_fused_whole_vector_operation_falls_back() { + for query in ["sort(some_metric)", "absent(some_metric)"] { + let native = build_native_fused_plan(query, true).await; + assert!( + find_fused_series(&native).is_none(), + "`{query}` reads the whole vector and must not be placed region-side" + ); + } +} + #[tokio::test] async fn native_fused_unsupported_shape_falls_back() { let plan = build_native_fused_plan("sum(rate(some_metric[5m]))", true).await; diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 80e7574bd13..0b37a9a4324 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -100,6 +100,9 @@ pub struct QueryEngineState { extension_rules: Vec>, plugins: Plugins, experimental_native_promql: bool, + /// Whether this engine's regions are reached in process, which is what lets a region take a + /// `LogicalPlan` instead of Substrait bytes. + in_process_region_transport: bool, } impl fmt::Debug for QueryEngineState { @@ -334,6 +337,7 @@ impl QueryEngineState { extension_rules, plugins, experimental_native_promql: options.experimental_native_promql, + in_process_region_transport, scalar_functions: Arc::new(RwLock::new(HashMap::new())), }) } @@ -486,6 +490,11 @@ impl QueryEngineState { self.experimental_native_promql } + /// Whether a region read is served in this process, so a region can take a `LogicalPlan`. + pub(crate) fn in_process_region_transport(&self) -> bool { + self.in_process_region_transport + } + pub fn catalog_manager(&self) -> &CatalogManagerRef { &self.catalog_manager } From de97d868007d23762557166b4f28c1d276608564 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 12:46:54 +0800 Subject: [PATCH 14/38] fix(query): only offer the native route to an engine that serves frames Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 23 +++++ src/query/src/promql/planner/test.rs | 131 +++++++++++++++++++++---- 2 files changed, 135 insertions(+), 19 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 31b083b08ea..97647201858 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -30,6 +30,7 @@ use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; +use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; @@ -167,6 +168,14 @@ impl PromPlanner { }; let table = self.table_from_source(&scan.source)?; let table_meta = &table.table_info().meta; + // An engine without a raw-series scanner has no frame to give this node, and by the time + // the physical planner asks, the logical plan is already this node: the refusal there is a + // planning error, not a fallback. So the engine is checked here, where the existing plan is + // still the answer. A file-engine table is the concrete case: its scanner is a + // `SinglePartitionScanner`, which leaves `raw_series_schema` at the default `None`. + if !supports_native_raw_series(&table_meta.engine) { + return Ok(None); + } // The region's raw-series preflight refuses these, and the physical layer turns that refusal into an error, so decline this route here instead. let time_index_type = table_meta .schema @@ -373,6 +382,20 @@ fn selector_range_millis(expr: &PromExpr) -> Option { } } +/// Whether `engine` can serve raw series frames to the node this route builds. +/// +/// A whitelist, not a probe: the logical planner cannot ask a region what its scanner supports +/// without reading the table, and a false negative only costs the query its native route. Mito's +/// series scanner implements the raw path; every other engine in the workspace leaves +/// `raw_series_schema` at its default `None`, so only those two are named. This is deliberately not +/// [`supports_pk_dictionary_encoding`]: that answers a different question and includes the metric +/// engine, whose logical tables a bare scan cannot reach anyway. +/// +/// [`supports_pk_dictionary_encoding`]: table::table::scan +fn supports_native_raw_series(engine: &str) -> bool { + matches!(engine, MITO_ENGINE | MITO2_ENGINE) +} + /// Whether any node of `tree` reads the whole input vector at a point rather than one series. /// /// Such a node folds or pairs across series (`Aggregate`, `Binary`) or reorders/reads a whole diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index c01d2d40f4f..eef10150057 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -18,7 +18,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use catalog::RegisterTableRequest; use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; use common_base::Plugins; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE, MITO_ENGINE}; use common_query::native_histogram::{ CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, }; @@ -894,6 +894,18 @@ async fn optimize_and_create_physical_plan( (optimized, physical) } +/// The engine a fixture gets when the test does not care where the scan is served. +/// +/// `TableMeta::empty()` leaves it blank, and the existing plan assertions are written against that +/// blank engine; naming it here keeps the native-route fixtures explicit instead. +const DEFAULT_ENGINE: &str = ""; + +/// The engine the native route accepts, because its scanner serves raw series frames. +/// +/// It also decides whether a tag column is dictionary encoded, so a fixture that sets it sees +/// `Dictionary(UInt32, Utf8)` where a blank-engine fixture sees `Utf8`. +const NATIVE_ENGINE: &str = MITO_ENGINE; + async fn build_test_table_provider( table_name_tuples: &[(String, String)], num_tag: usize, @@ -917,6 +929,7 @@ async fn build_test_table_provider_with_partition_columns( ConcreteDataType::timestamp_millisecond_datatype(), false, ConcreteDataType::float64_datatype(), + DEFAULT_ENGINE, ) .await } @@ -933,6 +946,22 @@ async fn build_test_table_provider_with_field_type( ConcreteDataType::timestamp_millisecond_datatype(), false, field_type, + DEFAULT_ENGINE, + ) + .await +} + +/// A one-tag, one-field fixture on the engine the native route accepts. +async fn native_engine_provider(partition_key_indices: Vec) -> DfTableSourceProvider { + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + partition_key_indices, + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, ) .await } @@ -945,6 +974,7 @@ async fn build_test_table_provider_with_options( timestamp_type: ConcreteDataType, append_mode: bool, field_type: ConcreteDataType, + engine: &str, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); for (schema_name, table_name) in table_name_tuples { @@ -970,6 +1000,7 @@ async fn build_test_table_provider_with_options( let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() .schema(schema) + .engine(engine.to_string()) .primary_key_indices((0..num_tag).collect()) .partition_key_indices(partition_key_indices.clone()) .value_indices((num_tag + 1..num_tag + 1 + num_field).collect()) @@ -1681,10 +1712,15 @@ async fn do_single_instant_function_call(fn_name: &'static str, plan_name: &str) /// time index `timestamp`, one field `field_0` — which is the shape the selector context resolves /// to a raw series schema of one field. async fn build_native_fused_plan(query: &str, native_fused: bool) -> LogicalPlan { - let table_provider = build_test_table_provider( + let table_provider = build_test_table_provider_with_options( &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], 1, 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, ) .await; build_native_fused_plan_with_provider(query, native_fused, table_provider).await @@ -1744,18 +1780,20 @@ async fn native_fused_rate_plans_to_a_fused_series_node() { #[tokio::test] async fn native_fused_time_partitioned_table_falls_back() { - let table_provider = build_test_table_provider_with_partition_columns( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - vec![1], // timestamp - ) - .await; - let plan = + let table_provider = native_engine_provider(vec![1]).await; + let native = build_native_fused_plan_with_provider("rate(some_metric[2m])", true, table_provider).await; - let display = plan.display_indent().to_string(); + let fallback_provider = native_engine_provider(vec![1]).await; + let without_flag = + build_native_fused_plan_with_provider("rate(some_metric[2m])", false, fallback_provider) + .await; + let display = native.display_indent().to_string(); assert!(!display.contains("FusedSeries"), "{display}"); - assert!(display.contains("PromRangeManipulate"), "{display}"); + assert_eq!( + display, + without_flag.display_indent().to_string(), + "a time-partitioned table must keep the existing plan, not an approximation of it" + ); } #[tokio::test] @@ -1768,6 +1806,7 @@ async fn native_fused_non_millisecond_time_index_falls_back() { ConcreteDataType::timestamp_microsecond_datatype(), false, ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, ) .await; let native = @@ -1780,6 +1819,7 @@ async fn native_fused_non_millisecond_time_index_falls_back() { ConcreteDataType::timestamp_microsecond_datatype(), false, ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, ) .await; let without_flag = @@ -1802,6 +1842,7 @@ async fn native_fused_append_mode_table_falls_back() { ConcreteDataType::timestamp_millisecond_datatype(), true, ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, ) .await; let native = @@ -1814,6 +1855,7 @@ async fn native_fused_append_mode_table_falls_back() { ConcreteDataType::timestamp_millisecond_datatype(), true, ConcreteDataType::float64_datatype(), + NATIVE_ENGINE, ) .await; let without_flag = @@ -1828,19 +1870,64 @@ async fn native_fused_append_mode_table_falls_back() { #[tokio::test] async fn native_fused_tag_partitioned_table_still_plans_fused_series() { - let table_provider = build_test_table_provider_with_partition_columns( - &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], - 1, - 1, - vec![0], // tag_0 - ) - .await; + let table_provider = native_engine_provider(vec![0]).await; let plan = build_native_fused_plan_with_provider("rate(some_metric[5m])", true, table_provider).await; assert!(plan.display_indent().to_string().contains("FusedSeries")); } /// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. +/// A table whose engine has no raw-series scanner keeps the existing plan. +/// +/// A file-engine table answers an ordinary read with a `SinglePartitionScanner`, whose +/// `raw_series_schema` is the trait's default `None`. Offering the native route to such a table +/// would build the node and then fail in physical planning, because the logical plan is fixed by +/// then; declining here is what keeps the query answering. +#[tokio::test] +async fn native_fused_engine_without_raw_series_falls_back() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + FILE_ENGINE, + ) + .await, + ) + .await; + let without_flag = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + FILE_ENGINE, + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_none(), + "a file-engine table must not reach the native route" + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "a file-engine table must keep the existing plan" + ); +} + /// A value field the region's raw scan cannot serve keeps the existing plan. /// /// The preflight serves numeric columns and native histograms only, and the physical layer turns @@ -1881,10 +1968,16 @@ async fn native_fused_non_numeric_field_falls_back() { async fn native_fused_whole_vector_operation_falls_back() { for query in ["sort(some_metric)", "absent(some_metric)"] { let native = build_native_fused_plan(query, true).await; + let without_flag = build_native_fused_plan(query, false).await; assert!( find_fused_series(&native).is_none(), "`{query}` reads the whole vector and must not be placed region-side" ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan, not an approximation of it" + ); } } From ff1288f1785df36d09fdbfd3199f0d33a16bf4fe Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 13:12:51 +0800 Subject: [PATCH 15/38] fix(query): match the raw scan's own field and engine rules Signed-off-by: discord9 --- src/query/src/promql/planner.rs | 11 ++++--- src/query/src/promql/planner/native.rs | 22 +++++++++----- src/query/src/promql/planner/test.rs | 42 ++++++++++++++++++++++++-- 3 files changed, 62 insertions(+), 13 deletions(-) diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 2de6a87f241..07d4e61dd90 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -5507,16 +5507,19 @@ impl PromPlanner { /// Whether every field column carries a type the region's raw series preflight serves. /// - /// That preflight accepts numeric columns and native histograms only, so anything else - a - /// string, a boolean, a timestamp - has no raw framing. A route that offered it would fail - /// after the logical plan is fixed; declining keeps the existing plan. + /// The preflight decides on `ConcreteDataType::is_numeric`, which excludes the decimals, so + /// the Arrow type is converted rather than asked: Arrow's own `is_numeric` includes + /// `Decimal128`/`Decimal256` (and `Float16`, which has no column type here), and a decimal + /// field would then pass here and fail once the logical plan is fixed. Anything else - a + /// string, a boolean, a timestamp - has no raw framing either. fn field_columns_are_raw_series_servable( schema: &DFSchemaRef, field_columns: &[String], ) -> bool { field_columns.iter().all(|field| { Self::field_column_type(schema, field).is_some_and(|data_type| { - data_type.is_numeric() || data_type == &Self::native_histogram_arrow_type() + datatypes::prelude::ConcreteDataType::from_arrow_type(data_type).is_numeric() + || data_type == &Self::native_histogram_arrow_type() }) }) } diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 97647201858..88e06e33f7a 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -30,7 +30,7 @@ use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; -use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE}; +use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; @@ -385,15 +385,23 @@ fn selector_range_millis(expr: &PromExpr) -> Option { /// Whether `engine` can serve raw series frames to the node this route builds. /// /// A whitelist, not a probe: the logical planner cannot ask a region what its scanner supports -/// without reading the table, and a false negative only costs the query its native route. Mito's -/// series scanner implements the raw path; every other engine in the workspace leaves -/// `raw_series_schema` at its default `None`, so only those two are named. This is deliberately not -/// [`supports_pk_dictionary_encoding`]: that answers a different question and includes the metric -/// engine, whose logical tables a bare scan cannot reach anyway. +/// without reading the table, and a false negative only costs the query its native route. +/// +/// Mito's series scanner implements the raw path. The metric engine is named too because a logical +/// region reads through its physical region and hands back that same scanner - its own raw-series +/// test is `test_raw_series_scan_uses_logical_metric_route_and_bound_metadata` - and the ordinary +/// metric rewrite puts a filter and an alias between the scan and this node, which the bare-scan +/// check above already refuses. A logical table reaches that check with a bare scan only when the +/// physical table cannot be resolved, and there the logical scan is served by the same scanner. +/// +/// Every other engine in the workspace leaves `raw_series_schema` at its default `None`: a +/// file-engine table, for one, is read by a `SinglePartitionScanner` that cannot frame a series. +/// +/// This is deliberately not [`supports_pk_dictionary_encoding`]: that answers a different question. /// /// [`supports_pk_dictionary_encoding`]: table::table::scan fn supports_native_raw_series(engine: &str) -> bool { - matches!(engine, MITO_ENGINE | MITO2_ENGINE) + matches!(engine, MITO_ENGINE | MITO2_ENGINE | METRIC_ENGINE) } /// Whether any node of `tree` reads the whole input vector at a point rather than one series. diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index eef10150057..7eddd3b81db 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -18,7 +18,9 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use catalog::RegisterTableRequest; use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; use common_base::Plugins; -use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE, MITO_ENGINE}; +use common_catalog::consts::{ + DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE, +}; use common_query::native_histogram::{ CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, }; @@ -946,7 +948,9 @@ async fn build_test_table_provider_with_field_type( ConcreteDataType::timestamp_millisecond_datatype(), false, field_type, - DEFAULT_ENGINE, + // The native engine, so the field guard below is what refuses this table: on the blank + // engine the engine gate would answer first and this test would prove nothing. + NATIVE_ENGINE, ) .await } @@ -1877,6 +1881,37 @@ async fn native_fused_tag_partitioned_table_still_plans_fused_series() { } /// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. +/// A metric-engine logical table is served by its physical region's scanner, so it is admitted. +/// +/// The ordinary metric rewrite puts a filter and an alias between the scan and this node, which the +/// bare-scan check refuses - that is why such a query usually falls back. This test pins the one +/// shape that reaches the node bare (the physical table could not be resolved), so that the engine +/// whitelist does not refuse a table whose region really does serve frames. +#[tokio::test] +async fn native_fused_metric_engine_bare_scan_is_admitted() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + METRIC_ENGINE, + ) + .await, + ) + .await; + assert_eq!( + find_fused_series(&native).map(|fused| fused.promql().to_string()), + Some("rate(some_metric[5m])".to_string()), + "a metric-engine table serves frames, so a bare scan over one keeps the native route" + ); +} + /// A table whose engine has no raw-series scanner keeps the existing plan. /// /// A file-engine table answers an ordinary read with a `SinglePartitionScanner`, whose @@ -1938,6 +1973,9 @@ async fn native_fused_non_numeric_field_falls_back() { for field_type in [ ConcreteDataType::string_datatype(), ConcreteDataType::boolean_datatype(), + // Arrow's `is_numeric` includes the decimals while the region's preflight does not, so + // asking Arrow would let a decimal field through to a planning error. + ConcreteDataType::decimal128_datatype(10, 2), ] { let native = build_native_fused_plan_with_provider( "some_metric", From 8e55a489520587018cfd6a0801cb174aa76f3d30 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 13:26:39 +0800 Subject: [PATCH 16/38] refactor(promql): keep one cross-series predicate for both callers Signed-off-by: discord9 --- src/promql/src/native/plan.rs | 22 +++++++++++++ src/query/src/dist_plan/commutativity.rs | 16 ++-------- src/query/src/promql/planner/native.rs | 39 ++++++++---------------- src/query/src/promql/planner/test.rs | 39 ++++++++++++++++++------ 4 files changed, 66 insertions(+), 50 deletions(-) diff --git a/src/promql/src/native/plan.rs b/src/promql/src/native/plan.rs index fd308c6d9fb..2d6ee4090cc 100644 --- a/src/promql/src/native/plan.rs +++ b/src/promql/src/native/plan.rs @@ -347,6 +347,28 @@ pub struct NativeNode { pub output: NativeType, } +/// Whether `node` reads the whole input vector at a point rather than one series at a time. +/// +/// Three kinds do: an `Aggregate` folds every series of its group, a `Binary` pairs two vectors' +/// series by label, and a `Vector` operation (`sort`, `absent`, `topk`, `limitk`, +/// `histogram_quantile`, ...) reorders or reads the whole vector at a grid point. Everything else - +/// an input, a scalar, a selection, a window, a map - answers one series from one series' rows. +/// +/// Both a region-side placement and the native route's admission ask this, and they must agree: a +/// tree admitted here but refused there becomes a node the physical layer cannot build, and one +/// refused here but admitted there would be pushed to a region that holds part of the metric. +pub fn reads_across_series(node: &NativeNode) -> bool { + match &node.expr { + NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, + NativeExpr::Instant { input, .. } + | NativeExpr::Window { input, .. } + | NativeExpr::Map { input, .. } => reads_across_series(input), + NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } | NativeExpr::Vector { .. } => { + true + } + } +} + /// A bound evaluation tree. /// /// The tree itself carries no run-time state: every execution builds its own, so one compiled plan diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index ce8466cd68f..17114e5dc55 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -23,7 +23,7 @@ use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use promql::extension_plan::{ EmptyMetric, FusedSeries, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize, }; -use promql::native::plan::{NativeExpr, NativeNode}; +use promql::native::plan::reads_across_series; use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME; use crate::dist_plan::MergeScanLogicalPlan; @@ -393,18 +393,6 @@ impl Categorizer { /// several series at once. Neither has a compiler producer yet, which is why refusing them costs /// nothing today: a tree built some other way is exactly what this check is for, because evaluating /// it one region at a time would fold a partial input and answer a wrong query silently. -fn reads_across_series(node: &NativeNode) -> bool { - match &node.expr { - NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, - NativeExpr::Instant { input, .. } - | NativeExpr::Window { input, .. } - | NativeExpr::Map { input, .. } => reads_across_series(input), - NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } => true, - // A vector operation reads the whole vector at each grid point (`topk`, `absent`, - // `histogram_quantile`, ...), so it is not a per-series node either. - NativeExpr::Vector { .. } => true, - } -} #[cfg(test)] mod tests { @@ -431,7 +419,7 @@ mod tests { }; use promql::native::plan::{ AggregateKind, AggregateStage, AggregateStateSchema, BoundAggregate, BoundInput, Grouping, - NativeType, + NativeExpr, NativeNode, NativeType, }; use promql_parser::parser; use store_api::storage::RegionId; diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 88e06e33f7a..6dabf979649 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -30,7 +30,7 @@ use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; -use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; +use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; @@ -38,6 +38,7 @@ use datafusion::logical_expr::{Extension, LogicalPlan}; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; +use promql::native::plan::reads_across_series; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector, @@ -387,12 +388,15 @@ fn selector_range_millis(expr: &PromExpr) -> Option { /// A whitelist, not a probe: the logical planner cannot ask a region what its scanner supports /// without reading the table, and a false negative only costs the query its native route. /// -/// Mito's series scanner implements the raw path. The metric engine is named too because a logical -/// region reads through its physical region and hands back that same scanner - its own raw-series -/// test is `test_raw_series_scan_uses_logical_metric_route_and_bound_metadata` - and the ordinary -/// metric rewrite puts a filter and an alias between the scan and this node, which the bare-scan -/// check above already refuses. A logical table reaches that check with a bare scan only when the -/// physical table cannot be resolved, and there the logical scan is served by the same scanner. +/// Mito's series scanner implements the raw path. +/// +/// The metric engine is deliberately absent. A logical metric table does read through its physical +/// region, which is a mito scanner, but the planner reaches this node with a bare scan only when it +/// could not resolve that physical table: the ordinary rewrite puts a `__table_id` filter and an +/// alias between the scan and this node, which the bare-scan check above already refuses. A logical +/// table whose physical table is missing is not a table this route can serve - its `TableInfo` +/// carries no `append_mode` either, so the preflight's own check could not be made here - and +/// naming the engine would admit exactly that unverifiable case. /// /// Every other engine in the workspace leaves `raw_series_schema` at its default `None`: a /// file-engine table, for one, is read by a `SinglePartitionScanner` that cannot frame a series. @@ -401,26 +405,7 @@ fn selector_range_millis(expr: &PromExpr) -> Option { /// /// [`supports_pk_dictionary_encoding`]: table::table::scan fn supports_native_raw_series(engine: &str) -> bool { - matches!(engine, MITO_ENGINE | MITO2_ENGINE | METRIC_ENGINE) -} - -/// Whether any node of `tree` reads the whole input vector at a point rather than one series. -/// -/// Such a node folds or pairs across series (`Aggregate`, `Binary`) or reorders/reads a whole -/// vector (`Vector`: `sort`, `absent`, `topk`, `limitk`, `histogram_quantile`, ...). This node runs -/// one executor per scan partition, and a partition holds only part of a metric's series, so those -/// nodes would answer from a fragment. -fn reads_across_series(node: &promql::native::plan::NativeNode) -> bool { - use promql::native::plan::NativeExpr; - match &node.expr { - NativeExpr::Input(_) | NativeExpr::Scalar(_) => false, - NativeExpr::Instant { input, .. } - | NativeExpr::Window { input, .. } - | NativeExpr::Map { input, .. } => reads_across_series(input), - NativeExpr::Aggregate { .. } | NativeExpr::Binary { .. } | NativeExpr::Vector { .. } => { - true - } - } + matches!(engine, MITO_ENGINE | MITO2_ENGINE) } fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 7eddd3b81db..01af409fca3 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1881,14 +1881,15 @@ async fn native_fused_tag_partitioned_table_still_plans_fused_series() { } /// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. -/// A metric-engine logical table is served by its physical region's scanner, so it is admitted. +/// A metric-engine logical table keeps the existing plan. /// -/// The ordinary metric rewrite puts a filter and an alias between the scan and this node, which the -/// bare-scan check refuses - that is why such a query usually falls back. This test pins the one -/// shape that reaches the node bare (the physical table could not be resolved), so that the engine -/// whitelist does not refuse a table whose region really does serve frames. +/// A logical table reads through its physical region, which is a mito scanner, but the planner +/// only reaches this node with a bare scan when that physical table could not be resolved: the +/// ordinary rewrite puts a `__table_id` filter and an alias in between. Such a table's `TableInfo` +/// also carries no `append_mode`, so the preflight's own check could not be made here. Both are +/// reasons the engine is not on the route's whitelist, and this test pins that it falls back. #[tokio::test] -async fn native_fused_metric_engine_bare_scan_is_admitted() { +async fn native_fused_metric_engine_table_falls_back() { let native = build_native_fused_plan_with_provider( "rate(some_metric[5m])", true, @@ -1905,10 +1906,30 @@ async fn native_fused_metric_engine_bare_scan_is_admitted() { .await, ) .await; + let without_flag = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_options( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + vec![], + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ConcreteDataType::float64_datatype(), + METRIC_ENGINE, + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_none(), + "a metric-engine table must not reach the native route" + ); assert_eq!( - find_fused_series(&native).map(|fused| fused.promql().to_string()), - Some("rate(some_metric[5m])".to_string()), - "a metric-engine table serves frames, so a bare scan over one keeps the native route" + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "a metric-engine table must keep the existing plan" ); } From ca804689eb93e14f5ab68d9ea3e2a2889ca749a5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 14:01:18 +0800 Subject: [PATCH 17/38] chore(query): share the offset conversion and correct a column-order note Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 6dabf979649..14fa4d99734 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -201,14 +201,9 @@ impl PromPlanner { { return Ok(None); } - let offset = selector - .offset - .as_ref() - .map(|offset| match offset { - promql_parser::parser::Offset::Pos(duration) => duration.as_millis() as i64, - promql_parser::parser::Offset::Neg(duration) => -(duration.as_millis() as i64), - }) - .unwrap_or_default(); + // The legacy selector converts an offset the same way; sharing that conversion keeps the + // two scan paths from disagreeing about an extreme duration. + let offset = Self::offset_millis(&selector.offset); let (saved_start, saved_end, saved_range) = (self.ctx.start, self.ctx.end, self.ctx.range); self.ctx.range = selector_range_millis(prom_expr); // Keep the fused scan's bounds identical to the legacy selector scan, including @ and offset. @@ -327,8 +322,9 @@ impl PromPlanner { /// The columns a [`FusedSeries`] node produces. /// -/// The shape is the one the existing unary pipeline consumes — the series fields, the tags, then -/// the time index — and each column is declared the way the node's own conversion reads it back: +/// The shape is the one the existing unary pipeline consumes — the time index, then the series +/// fields, then the tags — and each column is declared the way the node's own conversion reads it +/// back: /// the time column is a non-nullable millisecond timestamp with no timezone, because every /// evaluated chunk is addressed on the statement's grid and carries one; a field is a nullable /// `Float64`, because a grid point inside the statement's interval can have no sample; and a tag is From 51da078dc83ed5c12cf858605d2e2cdca533930e Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 15:23:54 +0800 Subject: [PATCH 18/38] chore(query): clear the two lints clippy reports on this branch A doc comment on the categorizer was followed by a blank line, and the test table builder takes eight arguments. Both are what CI's `-D warnings` turns into errors; neither changes behaviour. Signed-off-by: discord9 --- src/query/src/dist_plan/commutativity.rs | 1 - src/query/src/promql/planner/test.rs | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 17114e5dc55..940f921c7cf 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -393,7 +393,6 @@ impl Categorizer { /// several series at once. Neither has a compiler producer yet, which is why refusing them costs /// nothing today: a tree built some other way is exactly what this check is for, because evaluating /// it one region at a time would fold a partial input and answer a wrong query silently. - #[cfg(test)] mod tests { use std::collections::{BTreeMap, BTreeSet}; diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 01af409fca3..15732c987bb 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -970,6 +970,7 @@ async fn native_engine_provider(partition_key_indices: Vec) -> DfTableSou .await } +#[allow(clippy::too_many_arguments)] async fn build_test_table_provider_with_options( table_name_tuples: &[(String, String)], num_tag: usize, From 29ef60f6e89e46cd51cef77e587562d74e9629d7 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:19:39 +0800 Subject: [PATCH 19/38] test(promql): restore scan assignment coverage and update native guide Signed-off-by: discord9 --- src/promql/AGENTS.md | 80 +++++++++--------------- src/query/src/dist_plan/commutativity.rs | 8 --- src/table/src/table/scan.rs | 56 +++++++++++++++-- 3 files changed, 81 insertions(+), 63 deletions(-) diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md index dc90c5a0eec..3962add5e1b 100644 --- a/src/promql/AGENTS.md +++ b/src/promql/AGENTS.md @@ -16,38 +16,41 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`. | Kernels | `src/promql/src/functions.rs`, `src/promql/src/functions/` | PromQL functions | | Matching | `src/promql/src/vector_match.rs` | Binary vector matching | | Extension operators | `src/promql/src/extension_plan.rs`, `src/promql/src/extension_plan/` | Logical nodes and execution plans | -| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics | +| Native model | `src/promql/src/native/` | Compiled tree, chunks, display, metrics and execution | -## Two evaluation paths, and their status +## Evaluation paths -The easiest thing to get wrong here. Both exist; only one runs. - -- **Extension operators are live.** `PromPlanner` builds `SeriesDivide`, - `SeriesNormalize`, `RangeManipulate`, `InstantManipulate`, `EmptyMetric` and - `UnionDistinctOn`; `PromExtensionPlanner` - (`src/promql/src/extension_plan/planner.rs`) maps each logical node to its - execution plan, one branch per node, `None` for anything else. -- **Nothing plans into the native model yet.** `src/query/src/promql/native_compile.rs` - compiles a statement into a `NativeTree`, but no planner calls it, so a tree - built here reaches no execution. `src/promql/src/native/` holds the tree, the - chunk protocol, display and metrics, and — once a runtime lands — the driver - that runs a tree node by node. Compiling and running are separate concerns - owned by separate files; do not assume a tree is executable because it - compiles. +- **Extension operators:** the existing PromQL path plans operators such as + `SeriesDivide`, `SeriesNormalize`, `RangeManipulate`, `InstantManipulate`, + `EmptyMetric` and `UnionDistinctOn`. `PromExtensionPlanner` + (`src/promql/src/extension_plan/planner.rs`) maps the logical nodes to their + execution plans. +- **Fused native-series route:** the query planner offers a PromQL statement to + `native_compile::compile` only when the engine's + `in_process_region_transport` and `QueryOptions::experimental_native_promql` + options are enabled and the per-request `query.experimental_native_promql` + extension is true. Accepted plans use process-local `FusedSeries`; declined + or unsupported shapes fall back to the existing planner. The supported subset + requires one selector, a direct scan with a bound metric name, and no residual + label matchers or `__field__` matcher that narrows the scan, plus an eligible + raw-series table/field set and supported expression/grid. Metric scans retain + the static `__table_id` predicate. Hybrid CTEs and CREATE VIEW TQL creation + currently fall back to the existing planner. ## Change coupling -- **An extension node reaches a datanode only if three things line up**: a payload - branch in `src/common/substrait/src/extension_serializer.rs`, a physical branch - in `PromExtensionPlanner`, and a `Commutativity::Commutative` arm in - `check_extension_plan` (`src/query/src/dist_plan/commutativity.rs`). Missing the - third is silent: the node falls through to `Unsupported`, the distributed - boundary is placed above it, and it never reaches a region — with no error. +- Distributed extension operators need matching handling in the Substrait + serializer (`src/common/substrait/src/extension_serializer.rs`), the physical + `PromExtensionPlanner`, and `check_extension_plan` + (`src/query/src/dist_plan/commutativity.rs`). Process-local `FusedSeries` is + the exception: it is not serialized to a datanode. Its distributed placement + is restricted to a direct table scan, a bound metric name and trees that read + one series at a time. The single-region placement warning in + `commutativity.rs` remains important: do not assume a series stays in one + region. - A new logical node also needs its physical planner reachable from `DfQueryPlanner::create_physical_plan`, assembled in `src/query/src/query_engine/state.rs`. -- User-visible changes need planner coverage plus a sqlness case under - `tests/cases/standalone/common/promql/`. ## Gotchas @@ -72,9 +75,9 @@ The easiest thing to get wrong here. Both exist; only one runs. the producer ordering the adapter adds. Nothing in that metadata binds the key to a decode — the frame decoder compares those format keys and the field list, and never a label identity. `unpack_raw_series` therefore requires a - caller-supplied `RawLabelSource`: decode with the source that was actually bound - when the stream was produced, because the frame schema does not replace that - binding. The metadata describes format, not label identity. + caller-supplied `RawLabelSource`: decode with the source that was actually + bound when the stream was produced, because the frame schema does not replace + that binding. The metadata describes format, not label identity. - **A wrong label binding is silently wrong, not an error.** Storage validates the binding it is given (metadata identity, tag-ness, name matches, no duplicates) but cannot tell whether that set is the one the query meant, and the metric name @@ -85,26 +88,5 @@ The easiest thing to get wrong here. Both exist; only one runs. ```bash cargo nextest run -p promql cargo nextest run -p query # planner and native-compile changes -cargo sqlness bare -t # PromQL cases live under tests/cases/standalone/common/promql/ +cargo sqlness bare -t # PromQL cases under tests/cases/standalone/common/promql/ ``` - -## Not wired up yet - -The storage side of the raw path exists; what is missing is planner wiring. Both -gaps below are structural, not bugs. - -- **The native tree has no planner entry point.** `native_compile::compile` has no - caller, and a `NativeTree` becomes a physical plan only once something builds it - while it still holds the AST and the selector binding — the point where the - matchers, `__field__` and the lookback are still available. Compiling later, from - a plan, loses them. -- **No planner reaches a raw scan.** mito2 implements - `RegionScanner::raw_series_schema` (`mito2/src/read/series_scan.rs`), so - `RegionScanExec::try_to_series_scan` can build `SeriesScanExec` from a prepared - scanner; the trait's `Ok(None)` default only refuses engines that have not - implemented it. What is missing is the planner side: nothing selects that - conversion, and `SeriesScanExec::children()` is empty, so wrapping a - `RegionScanExec` before partition preparation leaves `ParallelizeScan` unable to - assign partitions. Either give the partitioner a child it can see, or convert to - the native leaf only after partition preparation — the second is the contract - `RegionScanExec::try_to_series_scan` documents. diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index 940f921c7cf..c11e78cf819 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -385,14 +385,6 @@ impl Categorizer { } } -/// Whether any node of `tree` reads more than one series at a time. -/// -/// Everything the native compiler emits today reads one series: an instant selection, a window and -/// a point-wise map each select, fold or map the samples of the series they were handed. An -/// aggregate folds a label set's series together and a binary operation pairs them, so both consume -/// several series at once. Neither has a compiler producer yet, which is why refusing them costs -/// nothing today: a tree built some other way is exactly what this check is for, because evaluating -/// it one region at a time would fold a partial input and answer a wrong query silently. #[cfg(test)] mod tests { use std::collections::{BTreeMap, BTreeSet}; diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 38548fa0853..6e69daa00a2 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -1360,9 +1360,11 @@ mod test { } #[test] - fn assign_partition_ranges_preserves_requested_partition_count() { - let ranges = (0..4) - .map(|identifier| PartitionRange { + fn assign_partition_ranges_balances_rows_deterministically() { + let ranges: Vec<_> = [100, 200, 150, 250] + .into_iter() + .enumerate() + .map(|(identifier, num_rows)| PartitionRange { start: common_time::Timestamp::new( identifier as i64, common_time::timestamp::TimeUnit::Second, @@ -1371,11 +1373,53 @@ mod test { identifier as i64 + 1, common_time::timestamp::TimeUnit::Second, ), - num_rows: 1, - identifier, + num_rows, + identifier: identifier + 1, }) .collect(); - assert_eq!(assign_partition_ranges(ranges, 4).len(), 4); + let assigned = assign_partition_ranges(ranges.clone(), 2); + assert_eq!( + assigned + .iter() + .map(|partition| { + partition + .iter() + .map(|range| range.identifier) + .collect::>() + }) + .collect::>(), + vec![vec![4, 1], vec![2, 3]] + ); + let mut assigned_identifiers = assigned + .iter() + .flatten() + .map(|range| range.identifier) + .collect::>(); + assigned_identifiers.sort_unstable(); + assert_eq!(assigned_identifiers, vec![1, 2, 3, 4]); + + let assigned = assign_partition_ranges(ranges.clone(), 5); + assert_eq!( + assigned + .iter() + .map(|partition| { + partition + .iter() + .map(|range| range.identifier) + .collect::>() + }) + .collect::>(), + vec![vec![4], vec![1], vec![2], vec![], vec![3]] + ); + let mut assigned_identifiers = assigned + .iter() + .flatten() + .map(|range| range.identifier) + .collect::>(); + assigned_identifiers.sort_unstable(); + assert_eq!(assigned_identifiers, vec![1, 2, 3, 4]); + + assert_eq!(assign_partition_ranges(vec![], 5), vec![vec![]; 5]); } #[test] From 340ea2d71e1d787fc5aabf4c690c4432aadc38ee Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:26:19 +0800 Subject: [PATCH 20/38] fix(query): retain metric table restriction in native scans Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 111 ++++++++++++++++++------- src/query/src/promql/planner/test.rs | 91 ++++++++++++++++---- 2 files changed, 155 insertions(+), 47 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 14fa4d99734..fbca18f39d4 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -30,11 +30,12 @@ use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; -use common_catalog::consts::{MITO_ENGINE, MITO2_ENGINE}; +use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion_expr::Expr; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; @@ -45,6 +46,9 @@ use promql_parser::parser::{ }; use session::context::QueryContextRef; use snafu::ResultExt; +use store_api::metric_engine_consts::{ + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, +}; use crate::promql::error::{DataFusionPlanningSnafu, Result}; use crate::promql::native_compile::{compile, statement_grid}; @@ -62,7 +66,7 @@ use crate::promql::planner::PromPlanner; /// [`QueryOptions::experimental_native_promql`]: crate::options::QueryOptions::experimental_native_promql /// [`remote_dyn_filter_pushdown_enabled_from_extensions`]: /// crate::options::remote_dyn_filter_pushdown_enabled_from_extensions -pub(crate) const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql"; +pub const NATIVE_PROMQL_EXTENSION_KEY: &str = "query.experimental_native_promql"; /// Whether `query_ctx` asks for the native fused-series route. /// @@ -160,12 +164,41 @@ impl PromPlanner { // instead would let the two paths disagree about the columns they read. let table_ref = self.table_ref()?; let mut input = self.create_table_scan_plan(table_ref.clone()).await?; - // The node reads the region scan itself, so the plan below it has to *be* that scan: a - // filter, a projection or a `__tsid` rewrite in between produces a different physical - // child, and the physical layer cannot reach through it. Declining here, while the logical - // plan is still this method's to choose, is what keeps such a query on the existing path. - let LogicalPlan::TableScan(scan) = &input else { - return Ok(None); + let logical_scan = metric_scan_and_table_id(&input); + let has_verified_metric_rewrite = logical_scan.is_some(); + let scan = match logical_scan { + Some((mut scan, predicate, table_id)) => { + let logical_provider = self + .table_provider + .resolve_table(table_ref.clone()) + .await + .context(crate::promql::error::CatalogSnafu)?; + let logical_table = self.table_from_source(&logical_provider)?; + let meta = &logical_table.table_info().meta; + if meta.engine != METRIC_ENGINE_NAME + || !meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY) + || logical_table.table_info().ident.table_id != table_id + { + return Ok(None); + } + // The physical scan bypasses logical-region request transformation, so preserve + // the validated logical-table restriction as a pushed-down scan predicate. + scan.filters.push(predicate); + input = LogicalPlan::TableScan(scan.clone()); + scan + } + None => { + if self.ctx.use_tsid { + return Ok(None); + } + match &input { + LogicalPlan::TableScan(scan) => scan.clone(), + _ => return Ok(None), + } + } }; let table = self.table_from_source(&scan.source)?; let table_meta = &table.table_info().meta; @@ -174,7 +207,9 @@ impl PromPlanner { // planning error, not a fallback. So the engine is checked here, where the existing plan is // still the answer. A file-engine table is the concrete case: its scanner is a // `SinglePartitionScanner`, which leaves `raw_series_schema` at the default `None`. - if !supports_native_raw_series(&table_meta.engine) { + if !supports_native_raw_series(&table_meta.engine) + || (table_meta.engine == METRIC_ENGINE_NAME && !has_verified_metric_rewrite) + { return Ok(None); } // The region's raw-series preflight refuses these, and the physical layer turns that refusal into an error, so decline this route here instead. @@ -379,29 +414,45 @@ fn selector_range_millis(expr: &PromExpr) -> Option { } } -/// Whether `engine` can serve raw series frames to the node this route builds. +/// Whether `engine` can serve raw series frames to the fused-series node. /// -/// A whitelist, not a probe: the logical planner cannot ask a region what its scanner supports -/// without reading the table, and a false negative only costs the query its native route. -/// -/// Mito's series scanner implements the raw path. -/// -/// The metric engine is deliberately absent. A logical metric table does read through its physical -/// region, which is a mito scanner, but the planner reaches this node with a bare scan only when it -/// could not resolve that physical table: the ordinary rewrite puts a `__table_id` filter and an -/// alias between the scan and this node, which the bare-scan check above already refuses. A logical -/// table whose physical table is missing is not a table this route can serve - its `TableInfo` -/// carries no `append_mode` either, so the preflight's own check could not be made here - and -/// naming the engine would admit exactly that unverifiable case. -/// -/// Every other engine in the workspace leaves `raw_series_schema` at its default `None`: a -/// file-engine table, for one, is read by a `SinglePartitionScanner` that cannot frame a series. -/// -/// This is deliberately not [`supports_pk_dictionary_encoding`]: that answers a different question. -/// -/// [`supports_pk_dictionary_encoding`]: table::table::scan +/// A whitelist avoids probing a region from logical planning; declining only costs the native route. fn supports_native_raw_series(engine: &str) -> bool { - matches!(engine, MITO_ENGINE | MITO2_ENGINE) + matches!(engine, MITO_ENGINE | MITO2_ENGINE | METRIC_ENGINE) +} + +/// Recognizes only the metric logical-table scan emitted by `create_table_scan_plan`. +fn metric_scan_and_table_id( + plan: &LogicalPlan, +) -> Option<(datafusion::logical_expr::TableScan, Expr, u32)> { + let plan = match plan { + LogicalPlan::Projection(projection) => projection.input.as_ref(), + plan => plan, + }; + let plan = match plan { + LogicalPlan::SubqueryAlias(alias) => alias.input.as_ref(), + plan => plan, + }; + let LogicalPlan::Filter(filter) = plan else { + return None; + }; + let Expr::BinaryExpr(binary) = &filter.predicate else { + return None; + }; + if binary.op != datafusion::logical_expr::Operator::Eq + || !matches!(binary.left.as_ref(), Expr::Column(column) if column.name == DATA_SCHEMA_TABLE_ID_COLUMN_NAME) + { + return None; + } + let Expr::Literal(datafusion::scalar::ScalarValue::UInt32(Some(table_id)), _) = + binary.right.as_ref() + else { + return None; + }; + let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { + return None; + }; + Some((scan.clone(), filter.predicate.clone(), *table_id)) } fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 15732c987bb..b085e3b4f2f 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -38,7 +38,7 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::{MemTable, provider_as_source}; use datafusion::execution::context::SessionContext; -use datafusion::logical_expr::{Extension, col}; +use datafusion::logical_expr::{Extension, UserDefinedLogicalNodeCore, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use promql::extension_plan::{FusedSeries, HistogramFold}; @@ -1748,6 +1748,22 @@ async fn build_native_fused_plan_with_provider( .unwrap() } +fn find_metric_physical_scan(plan: &LogicalPlan) -> Option { + match plan { + LogicalPlan::Extension(Extension { node }) => node + .as_any() + .downcast_ref::() + .and_then(|fused| match fused.inputs().first() { + Some(LogicalPlan::TableScan(scan)) => Some(scan.clone()), + _ => None, + }), + _ => plan + .inputs() + .into_iter() + .find_map(find_metric_physical_scan), + } +} + fn find_fused_series(plan: &LogicalPlan) -> Option<&FusedSeries> { if let LogicalPlan::Extension(Extension { node }) = plan && let Some(fused) = node.as_any().downcast_ref::() @@ -1881,16 +1897,61 @@ async fn native_fused_tag_partitioned_table_still_plans_fused_series() { assert!(plan.display_indent().to_string().contains("FusedSeries")); } -/// The flag on and a shape the compiler refuses: the existing planner plans it unchanged. -/// A metric-engine logical table keeps the existing plan. -/// -/// A logical table reads through its physical region, which is a mito scanner, but the planner -/// only reaches this node with a bare scan when that physical table could not be resolved: the -/// ordinary rewrite puts a `__table_id` filter and an alias in between. Such a table's `TableInfo` -/// also carries no `append_mode`, so the preflight's own check could not be made here. Both are -/// reasons the engine is not on the route's whitelist, and this test pins that it falls back. +/// Metric logical scans retain their table-id restriction on the direct fused scan input. #[tokio::test] -async fn native_fused_metric_engine_table_falls_back() { +async fn native_fused_metric_engine_table_is_accepted() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await, + ) + .await; + let legacy = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + false, + build_test_table_provider_with_tsid( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_some(), + "{}", + native.display_indent() + ); + assert!( + legacy + .display_indent() + .to_string() + .contains("PromRangeManipulate") + ); + let scan = find_metric_physical_scan(&native).expect("fused route keeps physical scan"); + assert!(scan.filters.iter().any(|filter| { + *filter + == datafusion::logical_expr::Expr::Column(datafusion_common::Column::new( + Some(datafusion_common::TableReference::bare("phy")), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + )) + .eq(datafusion::logical_expr::lit(1024u32)) + })); + assert!( + scan.filters + .iter() + .any(|filter| format!("{filter:?}").contains("timestamp")) + ); +} + +/// A metric table without logical-table metadata declines to the legacy plan. +#[tokio::test] +async fn native_fused_metric_table_without_logical_metadata_falls_back() { let native = build_native_fused_plan_with_provider( "rate(some_metric[5m])", true, @@ -1907,7 +1968,7 @@ async fn native_fused_metric_engine_table_falls_back() { .await, ) .await; - let without_flag = build_native_fused_plan_with_provider( + let legacy = build_native_fused_plan_with_provider( "rate(some_metric[5m])", false, build_test_table_provider_with_options( @@ -1923,14 +1984,10 @@ async fn native_fused_metric_engine_table_falls_back() { .await, ) .await; - assert!( - find_fused_series(&native).is_none(), - "a metric-engine table must not reach the native route" - ); + assert!(find_fused_series(&native).is_none()); assert_eq!( native.display_indent().to_string(), - without_flag.display_indent().to_string(), - "a metric-engine table must keep the existing plan" + legacy.display_indent().to_string() ); } From 27cd09ec553b0127881ff9cfa5d569bd4d16bb64 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:27:53 +0800 Subject: [PATCH 21/38] fix(query): keep TQL views serializable and expose native diagnostics Signed-off-by: discord9 --- src/operator/src/statement/ddl.rs | 8 +++++- src/query/src/dist_plan/analyzer.rs | 7 ++++- src/query/src/promql/planner.rs | 18 ++++++++++-- src/sql/src/parsers/create_parser.rs | 43 +++++++++++++++++++++++++++- 4 files changed, 70 insertions(+), 6 deletions(-) diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index fa9da9f322e..22c4a4cf25c 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -75,6 +75,7 @@ use partition::expr::{Operand, PartitionExpr, RestrictedOp}; use partition::multi_dim::MultiDimPartitionRule; use query::parser::QueryStatement; use query::plan::extract_and_rewrite_full_table_names; +use query::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use query::query_engine::DefaultSerializer; use query::sql::create_table_stmt; use session::context::QueryContextRef; @@ -1164,7 +1165,12 @@ impl StatementExecutor { ) .await? } - Statement::Tql(query) => self.plan_tql(query.clone(), &ctx).await?, + Statement::Tql(query) => { + // The serialized view plan must not depend on the native execution extension. + let mut view_ctx = ctx.as_ref().clone(); + view_ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, "false"); + self.plan_tql(query.clone(), &Arc::new(view_ctx)).await? + } _ => { return InvalidViewStmtSnafu {}.fail(); } diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 76020bf262b..92ec8502fab 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -513,7 +513,12 @@ impl PlanRewriter { /// `plan` must be the root of the tree that is about to be rewritten, otherwise /// [`PlanRewriter::should_expand`] may skip a check it has to perform. fn new(plan: &LogicalPlan, in_process_region_transport: bool) -> Self { - let skip_encodability_gate = in_process_region_transport && is_fused_series_plan(plan); + let native_root = match plan { + LogicalPlan::Analyze(analyze) => analyze.input.as_ref(), + _ => plan, + }; + let skip_encodability_gate = + in_process_region_transport && is_fused_series_plan(native_root); Self { // Preserve the ordinary whole-plan fast path unless this is the in-process root // exception: no node in that tree needs an encoding check. diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 07d4e61dd90..05acc1a9de5 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -17,6 +17,7 @@ mod function_plans; mod island; mod matching_filters; pub(crate) mod native; +pub use native::NATIVE_PROMQL_EXTENSION_KEY; mod set_operator; @@ -2421,9 +2422,20 @@ impl PromPlanner { // let promql_parser::parser::ast::Extension { expr } = ext_expr; let expr = &ext_expr.expr; let children = expr.children(); - let plan = self - .prom_root_expr_to_plan(&children[0], query_engine_state) - .await?; + let native_plan = match expr.name() { + ANALYZE_NODE_NAME + | ANALYZE_VERBOSE_NODE_NAME + | EXPLAIN_NODE_NAME + | EXPLAIN_VERBOSE_NODE_NAME => self.try_plan_native_fused(&children[0]).await?, + _ => None, + }; + let plan = match native_plan { + Some(plan) => plan, + None => { + self.prom_root_expr_to_plan(&children[0], query_engine_state) + .await? + } + }; // Wrapper for the explanation/analyze of the existing plan // https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/builder/struct.LogicalPlanBuilder.html#method.explain // if `analyze` is true, runs the actual plan and produces diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index 20eba9f2087..adb0bccf898 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -55,6 +55,7 @@ use crate::statements::create::{ CreateTableLike, CreateView, Partitions, SqlOrTql, TableConstraint, VECTOR_OPT_DIM, }; use crate::statements::statement::Statement; +use crate::statements::tql::Tql; use crate::statements::transform::type_alias::get_data_type_by_alias_name; use crate::statements::{OptionMap, sql_data_type_to_concrete_data_type}; use crate::util::{OptionValue, location_to_index, parse_option_string}; @@ -148,7 +149,18 @@ impl<'a> ParserContext<'a> { .expect_keyword(Keyword::AS) .context(SyntaxSnafu)?; - let query = self.parse_query()?; + let query = match self.parser.peek_token().token { + Token::Word(word) + if word.quote_style.is_none() + && word.value.eq_ignore_ascii_case(tql_parser::TQL) => + { + match self.parse_tql(false)? { + Statement::Tql(tql @ Tql::Eval(_)) => Statement::Tql(tql), + _ => return self.expected("TQL EVAL", self.parser.peek_token()), + } + } + _ => self.parse_query()?, + }; Ok(Statement::CreateView(CreateView { name: view_name, @@ -3315,6 +3327,35 @@ CREATE TABLE spans ( } } + #[test] + fn test_parse_create_view_tql_eval() { + let sql = "CREATE VIEW test AS TQL EVAL (0, 15, '5s') metric"; + let result = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + assert!(matches!( + &result[0], + Statement::CreateView(CreateView { + query, + .. + }) if matches!(query.as_ref(), Statement::Tql(Tql::Eval(_))) + )); + + for sql in [ + "CREATE VIEW test AS TQL EXPLAIN (0, 15, '5s') metric", + "CREATE VIEW test AS TQL ANALYZE (0, 15, '5s') metric", + ] { + assert!( + ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .is_err() + ); + } + } + #[test] fn test_parse_create_view_with_columns() { let sql = "CREATE VIEW test () AS SELECT * FROM NUMBERS"; From cd64176418ea100f48fa118cc26d137486ae14da Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:29:22 +0800 Subject: [PATCH 22/38] fix(sql): preserve lookback when displaying TQL view definitions Signed-off-by: discord9 --- src/sql/src/parsers/create_parser.rs | 40 ++++++++++++++++++++-------- src/sql/src/statements/tql.rs | 2 +- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index adb0bccf898..de9c5657358 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -3329,17 +3329,35 @@ CREATE TABLE spans ( #[test] fn test_parse_create_view_tql_eval() { - let sql = "CREATE VIEW test AS TQL EVAL (0, 15, '5s') metric"; - let result = - ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) - .unwrap(); - assert!(matches!( - &result[0], - Statement::CreateView(CreateView { - query, - .. - }) if matches!(query.as_ref(), Statement::Tql(Tql::Eval(_))) - )); + for (query, expected_lookback) in [ + ("TQL EVAL (0, 15, '5s') metric", None), + ("TQL EVAL (0, 15, '5s', '5m') metric", Some("5m")), + ("TQL EVAL (0, 15, '5s', '300') metric", Some("300")), + ] { + let sql = format!("CREATE VIEW test AS {query}"); + let result = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::CreateView(view) = &result[0] else { + panic!("expected CREATE VIEW, got {:?}", result[0]); + }; + let Statement::Tql(Tql::Eval(eval)) = view.query.as_ref() else { + panic!("expected TQL EVAL, got {}", view.query); + }; + assert_eq!(eval.lookback.as_deref(), expected_lookback); + + let formatted = result[0].to_string(); + let reparsed = ParserContext::create_with_dialect( + &formatted, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + assert_eq!(reparsed[0], result[0], "{formatted}"); + } for sql in [ "CREATE VIEW test AS TQL EXPLAIN (0, 15, '5s') metric", diff --git a/src/sql/src/statements/tql.rs b/src/sql/src/statements/tql.rs index 835a957eada..948c31ddd8d 100644 --- a/src/sql/src/statements/tql.rs +++ b/src/sql/src/statements/tql.rs @@ -46,7 +46,7 @@ fn format_tql( ) -> std::fmt::Result { write!(f, "({start}, {end}, '{step}'")?; if let Some(lookback) = lookback { - write!(f, ", {lookback}")?; + write!(f, ", '{lookback}'")?; } if let Some(alias) = alias { write!(f, ") {query} AS {alias}") From d62051e12a56ed3749192a7b40d339b47d23173d Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:44:03 +0800 Subject: [PATCH 23/38] fix(query): preserve logical metric field order in native scans Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 73 +++++++++++++------------ src/query/src/promql/planner/test.rs | 76 +++++++++++++++++++++++++- 2 files changed, 112 insertions(+), 37 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index fbca18f39d4..d607eddfeef 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -34,7 +34,7 @@ use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; -use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion::logical_expr::{Extension, LogicalPlan, TableScanBuilder}; use datafusion_expr::Expr; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; @@ -46,9 +46,7 @@ use promql_parser::parser::{ }; use session::context::QueryContextRef; use snafu::ResultExt; -use store_api::metric_engine_consts::{ - DATA_SCHEMA_TABLE_ID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, -}; +use store_api::metric_engine_consts::{DATA_SCHEMA_TABLE_ID_COLUMN_NAME, METRIC_ENGINE_NAME}; use crate::promql::error::{DataFusionPlanningSnafu, Result}; use crate::promql::native_compile::{compile, statement_grid}; @@ -167,38 +165,40 @@ impl PromPlanner { let logical_scan = metric_scan_and_table_id(&input); let has_verified_metric_rewrite = logical_scan.is_some(); let scan = match logical_scan { - Some((mut scan, predicate, table_id)) => { - let logical_provider = self - .table_provider - .resolve_table(table_ref.clone()) - .await - .context(crate::promql::error::CatalogSnafu)?; - let logical_table = self.table_from_source(&logical_provider)?; - let meta = &logical_table.table_info().meta; - if meta.engine != METRIC_ENGINE_NAME - || !meta - .options - .extra_options - .contains_key(LOGICAL_TABLE_METADATA_KEY) - || logical_table.table_info().ident.table_id != table_id - { + Some((mut scan, predicate)) => { + // Preserve logical field order and the table restriction on the direct physical scan. + let source_schema = scan.source.schema(); + let Some(mut ordered_projection) = fields + .iter() + .map(|field| source_schema.index_of(field).ok()) + .collect::>>() + else { return Ok(None); + }; + let original_projection = scan.projection.as_ref(); + let selected_indices = original_projection + .cloned() + .unwrap_or_else(|| (0..source_schema.fields().len()).collect()); + for index in selected_indices { + let is_field = fields + .iter() + .any(|field| field == source_schema.field(index).name()); + if !is_field && !ordered_projection.contains(&index) { + ordered_projection.push(index); + } } - // The physical scan bypasses logical-region request transformation, so preserve - // the validated logical-table restriction as a pushed-down scan predicate. scan.filters.push(predicate); + scan = TableScanBuilder::from(scan) + .with_projection(Some(ordered_projection)) + .build() + .context(DataFusionPlanningSnafu)?; input = LogicalPlan::TableScan(scan.clone()); scan } - None => { - if self.ctx.use_tsid { - return Ok(None); - } - match &input { - LogicalPlan::TableScan(scan) => scan.clone(), - _ => return Ok(None), - } - } + None => match &input { + LogicalPlan::TableScan(scan) => scan.clone(), + _ => return Ok(None), + }, }; let table = self.table_from_source(&scan.source)?; let table_meta = &table.table_info().meta; @@ -424,7 +424,7 @@ fn supports_native_raw_series(engine: &str) -> bool { /// Recognizes only the metric logical-table scan emitted by `create_table_scan_plan`. fn metric_scan_and_table_id( plan: &LogicalPlan, -) -> Option<(datafusion::logical_expr::TableScan, Expr, u32)> { +) -> Option<(datafusion::logical_expr::TableScan, Expr)> { let plan = match plan { LogicalPlan::Projection(projection) => projection.input.as_ref(), plan => plan, @@ -444,15 +444,16 @@ fn metric_scan_and_table_id( { return None; } - let Expr::Literal(datafusion::scalar::ScalarValue::UInt32(Some(table_id)), _) = - binary.right.as_ref() - else { + if !matches!( + binary.right.as_ref(), + Expr::Literal(datafusion::scalar::ScalarValue::UInt32(Some(_)), _) + ) { return None; - }; + } let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { return None; }; - Some((scan.clone(), filter.predicate.clone(), *table_id)) + Some((scan.clone(), filter.predicate.clone())) } fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index b085e3b4f2f..1846aab2002 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -1439,6 +1439,13 @@ async fn build_test_table_provider_with_tsid_fields( async fn build_test_table_provider_with_tsid_tag_fields( table_specs: &[((String, String), usize, usize)], +) -> DfTableSourceProvider { + build_test_table_provider_with_tsid_tag_fields_and_field_order(table_specs, None).await +} + +async fn build_test_table_provider_with_tsid_tag_fields_and_field_order( + table_specs: &[((String, String), usize, usize)], + logical_field_order: Option<&[usize]>, ) -> DfTableSourceProvider { let catalog_list = MemoryCatalogManager::with_default_setup(); @@ -1543,7 +1550,16 @@ async fn build_test_table_provider_with_tsid_tag_fields( ) .with_time_index(true), ); - for i in 0..*num_field { + let field_order = logical_field_order + .map(|order| { + order + .iter() + .copied() + .filter(|i| i < num_field) + .collect::>() + }) + .unwrap_or_else(|| (0..*num_field).collect()); + for i in field_order { columns.push(ColumnSchema::new( format!("field_{i}"), ConcreteDataType::float64_datatype(), @@ -1934,6 +1950,22 @@ async fn native_fused_metric_engine_table_is_accepted() { .contains("PromRangeManipulate") ); let scan = find_metric_physical_scan(&native).expect("fused route keeps physical scan"); + let source_schema = scan.source.schema(); + let projected_names = scan + .projection + .as_ref() + .unwrap() + .iter() + .map(|index| source_schema.field(*index).name()) + .collect::>(); + assert_eq!( + projected_names + .iter() + .filter(|name| name.starts_with("field_")) + .copied() + .collect::>(), + vec!["field_0"] + ); assert!(scan.filters.iter().any(|filter| { *filter == datafusion::logical_expr::Expr::Column(datafusion_common::Column::new( @@ -1949,6 +1981,48 @@ async fn native_fused_metric_engine_table_is_accepted() { ); } +/// Metric fields follow logical metadata order even when physical columns use another order. +#[tokio::test] +async fn native_fused_metric_fields_follow_logical_order() { + let native = build_native_fused_plan_with_provider( + "rate(some_metric[5m])", + true, + build_test_table_provider_with_tsid_tag_fields_and_field_order( + &[( + (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()), + 1, + 2, + )], + Some(&[1, 0]), + ) + .await, + ) + .await; + assert!( + find_fused_series(&native).is_some(), + "{}", + native.display_indent() + ); + let scan = find_metric_physical_scan(&native).expect("fused route keeps physical scan"); + let source_schema = scan.source.schema(); + let projected_names = scan + .projection + .as_ref() + .unwrap() + .iter() + .map(|index| source_schema.field(*index).name()) + .collect::>(); + assert_eq!(projected_names[..2], ["field_1", "field_0"]); + assert!(scan.filters.iter().any(|filter| { + *filter + == datafusion::logical_expr::Expr::Column(datafusion_common::Column::new( + Some(datafusion_common::TableReference::bare("phy")), + DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(), + )) + .eq(datafusion::logical_expr::lit(1024u32)) + })); +} + /// A metric table without logical-table metadata declines to the legacy plan. #[tokio::test] async fn native_fused_metric_table_without_logical_metadata_falls_back() { From c48c5ec5e13cc055817d2f7bccc8a9b0adb2f3dc Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:28:40 +0800 Subject: [PATCH 24/38] test(promql): execute native standalone queries and compare results Signed-off-by: discord9 (cherry picked from commit 12ea89c6ca329a91b6b3ee50ff73c93ab6626dc1) --- tests-integration/src/standalone.rs | 11 + tests-integration/src/tests/promql_test.rs | 295 +++++++++++++++++++++ 2 files changed, 306 insertions(+) diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 074be282543..80fb6ed36ba 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -91,6 +91,7 @@ pub struct GreptimeDbStandaloneBuilder { logical_batcher: Option, table_batcher: BatcherOptions, mito_config: Option, + query_options: query::options::QueryOptions, } impl GreptimeDbStandaloneBuilder { @@ -112,9 +113,17 @@ impl GreptimeDbStandaloneBuilder { logical_batcher: None, table_batcher: BatcherOptions::default(), mito_config: None, + query_options: query::options::QueryOptions::default(), } } + /// Overrides query options for this standalone test instance. + #[must_use] + pub fn with_query_options(mut self, options: query::options::QueryOptions) -> Self { + self.query_options = options; + self + } + /// Overrides the Mito configuration for this test instance. #[must_use] pub fn with_mito_config(mut self, config: mito2::config::MitoConfig) -> Self { @@ -316,6 +325,7 @@ impl GreptimeDbStandaloneBuilder { procedure_executor.clone(), Arc::new(ProcessManager::new(server_addr, None)), ) + .with_in_process_region_transport() .with_local_file_access(local_file_access) .with_plugin(plugins.clone()) .try_build() @@ -389,6 +399,7 @@ impl GreptimeDbStandaloneBuilder { slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, + query: self.query_options.clone(), pending_rows_batcher: PendingRowsBatcherOptions { logical_table: self.logical_batcher.clone(), table: self.table_batcher.clone(), diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 7796be4d15b..c26670f87e7 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -29,7 +29,9 @@ use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; use frontend::instance::Instance; use query::datafusion::QUERY_PARALLELISM_HINT; +use query::options::QueryOptions; use query::parser::{PromQuery, QueryLanguageParser, QueryStatement}; +use query::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use rstest::rstest; use rstest_reuse::apply; use servers::query_handler::sql::SqlQueryHandler; @@ -38,6 +40,7 @@ use session::context::QueryContext; use super::test_util::{ both_instances_cases, check_unordered_output_stream, distributed, standalone, }; +use crate::standalone::GreptimeDbStandaloneBuilder; use crate::tests::test_util::MockInstance; #[allow(clippy::too_many_arguments)] @@ -1752,3 +1755,295 @@ async fn f06_scalar_metric_table_lag_order_by_ts_desc(instance: Arc Arc { + let mut ctx = QueryContext::with_db_name(None); + ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); + ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, enabled.to_string()); + Arc::new(ctx) +} + +async fn native_tql_output(instance: &Arc, sql: &str, enabled: bool) -> Output { + instance + .do_query(sql, native_query_context(enabled, 4)) + .await + .remove(0) + .unwrap() +} + +async fn native_output_text(output: Output) -> String { + output.data.pretty_print().await +} + +async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("TQL EVAL returns record batches"), + }; + let mut rows = batches + .iter() + .flat_map(|batch| { + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let hosts = cast( + batch.column_by_name("host").unwrap().as_ref(), + &DataType::Utf8, + ) + .unwrap(); + let hosts = hosts.as_any().downcast_ref::().unwrap(); + let values = batch + .column_by_name("val") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|row| { + ( + timestamps.value(row), + hosts.value(row).to_string(), + values.value(row), + ) + }) + .collect::>() + }) + .collect::>(); + rows.sort_by(|left, right| (left.0, &left.1).cmp(&(right.0, &right.1))); + rows +} + +fn assert_four_native_partitions(plan: &Arc) { + use datafusion::physical_plan::metrics::MetricValue; + + let mut stack = vec![plan.clone()]; + let mut fused_nodes = 0; + let mut scan_partition_counts = Vec::new(); + let mut values_by_partition = [0_u64; 4]; + while let Some(node) = stack.pop() { + if node.name().contains("FusedSeriesExec") { + fused_nodes += 1; + scan_partition_counts.push( + node.children() + .first() + .map(|child| child.output_partitioning().partition_count()) + .unwrap_or(0), + ); + } + if let Some(metrics) = node.metrics() { + for metric in metrics.iter() { + if metric.value().name().contains("native.root.output_values") { + if let MetricValue::Count { count, .. } = metric.value() { + if let Some(partition) = metric.partition() + && let Some(value) = values_by_partition.get_mut(partition) + { + *value += count.value() as u64; + } + } + } + } + } + stack.extend(node.children().into_iter().cloned()); + } + assert_eq!(fused_nodes, 1, "expected one fused execution node"); + assert_eq!( + scan_partition_counts, + [4], + "fused scan must have four partitions" + ); + assert!( + values_by_partition.iter().all(|value| *value > 0), + "every native execution partition must produce values: {values_by_partition:?}" + ); +} + +fn execution_plan_names(plan: &Arc) -> Vec { + let mut names = Vec::new(); + let mut stack = vec![plan.clone()]; + while let Some(node) = stack.pop() { + names.push(node.name().to_string()); + stack.extend(node.children().into_iter().cloned()); + } + names +} + +async fn native_test_instance(name: &str) -> crate::standalone::GreptimeDbStandalone { + let options = QueryOptions { + experimental_native_promql: true, + ..Default::default() + }; + GreptimeDbStandaloneBuilder::new(name) + .with_query_options(options) + .build() + .await +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_executes_fused_node_and_matches_legacy() { + common_telemetry::init_default_ut_logging(); + let standalone = native_test_instance("native_promql_fused_e2e").await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + "CREATE TABLE native_e2e_phy (ts TIMESTAMP TIME INDEX, val DOUBLE) ENGINE=metric WITH (\"physical_metric_table\" = \"true\"); CREATE TABLE native_e2e (host STRING, ts TIMESTAMP TIME INDEX, val DOUBLE, PRIMARY KEY(host)) ENGINE=metric WITH (\"on_physical_table\" = 'native_e2e_phy')", + QueryContext::arc(), + ) + .await; + execute_all(&instance, "INSERT INTO native_e2e (host, ts, val) VALUES ('a', 0, 2), ('b', 0, 4), ('c', 0, 6), ('d', 0, 8), ('e', 0, 10), ('f', 0, 12), ('g', 0, 14), ('h', 0, 16), ('i', 0, 18), ('j', 0, 20), ('k', 0, 22), ('l', 0, 24), ('m', 0, 26), ('n', 0, 28), ('o', 0, 30), ('p', 0, 32)", QueryContext::arc()).await; + let sql = "TQL EVAL (0, 0, '1s') native_e2e"; + let native = native_tql_output(&instance, sql, true).await; + let plan = native + .meta + .plan + .as_ref() + .expect("native query has physical plan") + .clone(); + assert!( + execution_plan_names(&plan) + .iter() + .any(|name| name.contains("FusedSeriesExec")), + "expected fused physical node: {plan:?}" + ); + let native_rows = native_result_rows(native).await; + assert_four_native_partitions(&plan); + let expected = (b'a'..=b'p') + .zip((1..=16).map(|index| index as f64 * 2.0)) + .map(|(host, value)| (0, (host as char).to_string(), value)) + .collect::>(); + let legacy_rows = native_result_rows(native_tql_output(&instance, sql, false).await).await; + assert_eq!(native_rows, expected); + assert_eq!(native_rows, legacy_rows); +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_explain_and_analyze_execute_fused_plan() { + common_telemetry::init_default_ut_logging(); + let standalone = native_test_instance("native_promql_explain_analyze_e2e").await; + let instance = standalone.fe_instance().clone(); + execute_all(&instance, "CREATE TABLE native_explain (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_explain VALUES (0, 'a', 2)", QueryContext::arc()).await; + for (statement, enabled) in [ + ("TQL EXPLAIN (0, 0, '1s') native_explain", true), + ("TQL EXPLAIN VERBOSE (0, 0, '1s') native_explain", true), + ("TQL ANALYZE (0, 0, '1s') native_explain", true), + ("TQL ANALYZE VERBOSE (0, 0, '1s') native_explain", true), + ("TQL EXPLAIN (0, 0, '1s') native_explain", false), + ] { + let output = native_tql_output(&instance, statement, enabled).await; + let text = native_output_text(output).await; + if enabled { + assert!( + text.contains("FusedSeriesExec"), + "expected native physical tree for {statement}: {text}" + ); + } else { + assert!( + !text.contains("FusedSeriesExec"), + "flag-off EXPLAIN unexpectedly has native plan: {text}" + ); + } + if statement.contains("ANALYZE") && enabled { + assert!( + text.contains("output_values"), + "ANALYZE must report executed native metrics: {text}" + ); + } + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_view_fallback_preserves_native_hint() { + common_telemetry::init_default_ut_logging(); + let standalone = native_test_instance("native_promql_view_fallback_e2e").await; + let instance = standalone.fe_instance().clone(); + execute_all(&instance, "CREATE TABLE native_view_source (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_view_source VALUES (0, 'a', 2)", QueryContext::arc()).await; + let ctx = native_query_context(true, 4); + let expected_extensions = ctx.extensions(); + execute_all( + &instance, + "CREATE VIEW native_view AS TQL EVAL (0, 0, '1s') native_view_source", + ctx.clone(), + ) + .await; + let view = instance + .do_query("SELECT * FROM native_view", ctx.clone()) + .await + .remove(0) + .unwrap(); + let view_rows = native_result_rows(view).await; + assert_eq!(ctx.extensions(), expected_extensions); + assert_eq!(view_rows, vec![(0, "a".to_string(), 2.0)]); + let direct = instance + .do_query("TQL EVAL (0, 0, '1s') native_view_source", ctx) + .await + .remove(0) + .unwrap(); + assert!( + execution_plan_names(direct.meta.plan.as_ref().expect("physical plan")) + .iter() + .any(|name| name.contains("FusedSeriesExec")) + ); + assert_eq!( + native_result_rows(direct).await, + vec![(0, "a".to_string(), 2.0)] + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_metric_logical_tables_are_isolated() { + common_telemetry::init_default_ut_logging(); + let standalone = native_test_instance("native_promql_metric_isolation_e2e").await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + r#"CREATE TABLE native_metric_phy ( + ts TIMESTAMP TIME INDEX, + val DOUBLE + ) ENGINE=metric WITH ("physical_metric_table" = "true"); + CREATE TABLE native_metric_a ( + host STRING, + ts TIMESTAMP TIME INDEX, + val DOUBLE, + PRIMARY KEY (host) + ) ENGINE=metric WITH ("on_physical_table" = 'native_metric_phy'); + CREATE TABLE native_metric_b ( + host STRING, + ts TIMESTAMP TIME INDEX, + val DOUBLE, + PRIMARY KEY (host) + ) ENGINE=metric WITH ("on_physical_table" = 'native_metric_phy'); + INSERT INTO native_metric_a (host, ts, val) VALUES ('only_a', 0, 101), ('only_a', 1000, 102); + INSERT INTO native_metric_b (host, ts, val) VALUES ('only_b', 0, 202), ('only_b', 1000, 204);"#, + QueryContext::arc(), + ) + .await; + for (table, host, value, foreign_host) in [ + ("native_metric_a", "only_a", 1.0, "only_b"), + ("native_metric_b", "only_b", 2.0, "only_a"), + ] { + let sql = format!("TQL EVAL (1, 1, '1s') rate({table}[2s])"); + let native = native_tql_output(&instance, &sql, true).await; + let plan = native + .meta + .plan + .as_ref() + .expect("native metric physical plan"); + assert!( + execution_plan_names(plan) + .iter() + .any(|name| name.contains("FusedSeriesExec")) + ); + let native_rows = native_result_rows(native).await; + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(native_rows, vec![(1000, host.to_string(), value)]); + assert_eq!(native_rows, legacy_rows); + assert!( + native_rows + .iter() + .all(|(_, row_host, _)| row_host != foreign_host) + ); + } +} From c1dcac0f2be7acc836918d7f828452fe8a40bc2d Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:33:17 +0800 Subject: [PATCH 25/38] test(promql): verify actual frontend results and view definition Signed-off-by: discord9 (cherry picked from commit 64a52a2a87245f16c4f2c7be20ba275302e586df) --- tests-integration/src/tests/promql_test.rs | 44 +++++++++++++++++++--- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index c26670f87e7..15cd2cd0fae 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -1797,8 +1797,10 @@ async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { .unwrap(); let hosts = hosts.as_any().downcast_ref::().unwrap(); let values = batch - .column_by_name("val") - .unwrap() + .columns() + .iter() + .find(|column| column.data_type() == &DataType::Float64) + .expect("native/legacy result has one Float64 value column") .as_any() .downcast_ref::() .unwrap(); @@ -1833,6 +1835,11 @@ fn assert_four_native_partitions(plan: &Arc>(); - let legacy_rows = native_result_rows(native_tql_output(&instance, sql, false).await).await; + let legacy_rows = native_result_rows( + instance + .do_promql_query(query, native_query_context(false, 4)) + .await + .remove(0) + .unwrap(), + ) + .await; assert_eq!(native_rows, expected); assert_eq!(native_rows, legacy_rows); } @@ -1964,7 +1988,7 @@ async fn native_promql_view_fallback_preserves_native_hint() { let expected_extensions = ctx.extensions(); execute_all( &instance, - "CREATE VIEW native_view AS TQL EVAL (0, 0, '1s') native_view_source", + "CREATE VIEW native_view AS TQL EVAL (0, 0, '1s', '5m') native_view_source", ctx.clone(), ) .await; @@ -1976,6 +2000,14 @@ async fn native_promql_view_fallback_preserves_native_hint() { let view_rows = native_result_rows(view).await; assert_eq!(ctx.extensions(), expected_extensions); assert_eq!(view_rows, vec![(0, "a".to_string(), 2.0)]); + let create_view = instance + .do_query("SHOW CREATE VIEW native_view", ctx.clone()) + .await + .remove(0) + .unwrap(); + let create_view_text = native_output_text(create_view).await; + assert!(create_view_text.contains("TQL EVAL")); + assert!(create_view_text.contains("5m")); let direct = instance .do_query("TQL EVAL (0, 0, '1s') native_view_source", ctx) .await From 980edef08d579206edd067e39f2cf538af438318 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 16:42:49 +0800 Subject: [PATCH 26/38] test(promql): observe actual fused region execution Signed-off-by: discord9 (cherry picked from commit 313c8c181f3331bbaa67539a2739e26d6d0e3dec) --- tests-integration/src/tests/promql_test.rs | 108 +++++++++++++-------- 1 file changed, 67 insertions(+), 41 deletions(-) diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 15cd2cd0fae..f113a7846aa 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -12,16 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use api::v1::value::ValueData; use api::v1::{ ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; +use common_base::Plugins; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::{Output, OutputData}; use common_recordbatch::util::collect_batches; +use datafusion::physical_plan::ExecutionPlanProperties; use datatypes::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, }; @@ -31,6 +33,7 @@ use frontend::instance::Instance; use query::datafusion::QUERY_PARALLELISM_HINT; use query::options::QueryOptions; use query::parser::{PromQuery, QueryLanguageParser, QueryStatement}; +use query::physical_wrapper::{PhysicalPlanWrapper, PhysicalPlanWrapperRef}; use query::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use rstest::rstest; use rstest_reuse::apply; @@ -1756,6 +1759,41 @@ async fn f06_scalar_metric_table_lag_order_by_ts_desc(instance: Arc>>); + +impl NativePlanObserver { + fn clear(&self) { + *self.0.lock().unwrap() = None; + } + + fn take(&self) -> Arc { + self.0 + .lock() + .unwrap() + .take() + .expect("native fused plan was captured") + } +} + +impl PhysicalPlanWrapper for NativePlanObserver { + fn wrap( + &self, + origin: Arc, + _ctx: session::context::QueryContextRef, + ) -> Arc { + let mut stack = vec![origin.clone()]; + while let Some(node) = stack.pop() { + if node.name() == "FusedSeriesExec" { + *self.0.lock().unwrap() = Some(node); + break; + } + stack.extend(node.children().into_iter().cloned()); + } + origin + } +} + fn native_query_context(enabled: bool, parallelism: usize) -> Arc { let mut ctx = QueryContext::with_db_name(None); ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string()); @@ -1878,21 +1916,32 @@ fn execution_plan_names(plan: &Arc names } -async fn native_test_instance(name: &str) -> crate::standalone::GreptimeDbStandalone { +async fn native_test_instance( + name: &str, +) -> ( + crate::standalone::GreptimeDbStandalone, + Arc, +) { + let observer = Arc::new(NativePlanObserver::default()); + let plugins = Plugins::new(); + let wrapper: PhysicalPlanWrapperRef = observer.clone(); + plugins.insert(wrapper); let options = QueryOptions { experimental_native_promql: true, ..Default::default() }; - GreptimeDbStandaloneBuilder::new(name) + let standalone = GreptimeDbStandaloneBuilder::new(name) + .with_plugin(plugins) .with_query_options(options) .build() - .await + .await; + (standalone, observer) } #[tokio::test(flavor = "multi_thread")] async fn native_promql_executes_fused_node_and_matches_legacy() { common_telemetry::init_default_ut_logging(); - let standalone = native_test_instance("native_promql_fused_e2e").await; + let (standalone, observer) = native_test_instance("native_promql_fused_e2e").await; let instance = standalone.fe_instance().clone(); execute_all( &instance, @@ -1908,24 +1957,14 @@ async fn native_promql_executes_fused_node_and_matches_legacy() { step: "1s".to_string(), ..PromQuery::default() }; + observer.clear(); let native = instance - .do_promql_query(query.clone(), native_query_context(true, 4)) + .do_promql_query(&query, native_query_context(true, 4)) .await .remove(0) .unwrap(); - let plan = native - .meta - .plan - .as_ref() - .expect("native query has physical plan") - .clone(); - assert!( - execution_plan_names(&plan) - .iter() - .any(|name| name.contains("FusedSeriesExec")), - "expected fused physical node: {plan:?}" - ); let native_rows = native_result_rows(native).await; + let plan = observer.take(); assert_four_native_partitions(&plan); let expected = (b'a'..=b'p') .zip((1..=16).map(|index| index as f64 * 2.0)) @@ -1933,7 +1972,7 @@ async fn native_promql_executes_fused_node_and_matches_legacy() { .collect::>(); let legacy_rows = native_result_rows( instance - .do_promql_query(query, native_query_context(false, 4)) + .do_promql_query(&query, native_query_context(false, 4)) .await .remove(0) .unwrap(), @@ -1946,7 +1985,7 @@ async fn native_promql_executes_fused_node_and_matches_legacy() { #[tokio::test(flavor = "multi_thread")] async fn native_promql_explain_and_analyze_execute_fused_plan() { common_telemetry::init_default_ut_logging(); - let standalone = native_test_instance("native_promql_explain_analyze_e2e").await; + let (standalone, _observer) = native_test_instance("native_promql_explain_analyze_e2e").await; let instance = standalone.fe_instance().clone(); execute_all(&instance, "CREATE TABLE native_explain (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_explain VALUES (0, 'a', 2)", QueryContext::arc()).await; for (statement, enabled) in [ @@ -1981,7 +2020,7 @@ async fn native_promql_explain_and_analyze_execute_fused_plan() { #[tokio::test(flavor = "multi_thread")] async fn native_promql_view_fallback_preserves_native_hint() { common_telemetry::init_default_ut_logging(); - let standalone = native_test_instance("native_promql_view_fallback_e2e").await; + let (standalone, observer) = native_test_instance("native_promql_view_fallback_e2e").await; let instance = standalone.fe_instance().clone(); execute_all(&instance, "CREATE TABLE native_view_source (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_view_source VALUES (0, 'a', 2)", QueryContext::arc()).await; let ctx = native_query_context(true, 4); @@ -2008,26 +2047,21 @@ async fn native_promql_view_fallback_preserves_native_hint() { let create_view_text = native_output_text(create_view).await; assert!(create_view_text.contains("TQL EVAL")); assert!(create_view_text.contains("5m")); + observer.clear(); let direct = instance .do_query("TQL EVAL (0, 0, '1s') native_view_source", ctx) .await .remove(0) .unwrap(); - assert!( - execution_plan_names(direct.meta.plan.as_ref().expect("physical plan")) - .iter() - .any(|name| name.contains("FusedSeriesExec")) - ); - assert_eq!( - native_result_rows(direct).await, - vec![(0, "a".to_string(), 2.0)] - ); + let direct_rows = native_result_rows(direct).await; + assert_eq!(direct_rows, vec![(0, "a".to_string(), 2.0)]); + assert_eq!(observer.take().name(), "FusedSeriesExec"); } #[tokio::test(flavor = "multi_thread")] async fn native_promql_metric_logical_tables_are_isolated() { common_telemetry::init_default_ut_logging(); - let standalone = native_test_instance("native_promql_metric_isolation_e2e").await; + let (standalone, observer) = native_test_instance("native_promql_metric_isolation_e2e").await; let instance = standalone.fe_instance().clone(); execute_all( &instance, @@ -2057,18 +2091,10 @@ async fn native_promql_metric_logical_tables_are_isolated() { ("native_metric_b", "only_b", 2.0, "only_a"), ] { let sql = format!("TQL EVAL (1, 1, '1s') rate({table}[2s])"); + observer.clear(); let native = native_tql_output(&instance, &sql, true).await; - let plan = native - .meta - .plan - .as_ref() - .expect("native metric physical plan"); - assert!( - execution_plan_names(plan) - .iter() - .any(|name| name.contains("FusedSeriesExec")) - ); let native_rows = native_result_rows(native).await; + assert_eq!(observer.take().name(), "FusedSeriesExec"); let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; assert_eq!(native_rows, vec![(1000, host.to_string(), value)]); assert_eq!(native_rows, legacy_rows); From ea060286cc7995f3181c1140d961b6eedd1e1964 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 17:03:27 +0800 Subject: [PATCH 27/38] test(promql): require real partition output and strict native results Signed-off-by: discord9 (cherry picked from commit a13b1a723e7839f5b4e799a4493a507369de8893) --- tests-integration/src/tests/promql_test.rs | 154 ++++++++++----------- 1 file changed, 74 insertions(+), 80 deletions(-) diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index f113a7846aa..75fd38c387f 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -1810,7 +1810,12 @@ async fn native_tql_output(instance: &Arc, sql: &str, enabled: bool) - } async fn native_output_text(output: Output) -> String { - output.data.pretty_print().await + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("query returns record batches"), + }; + batches.pretty_print().unwrap() } async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { @@ -1834,14 +1839,19 @@ async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { ) .unwrap(); let hosts = hosts.as_any().downcast_ref::().unwrap(); - let values = batch + let value_columns = batch .columns() .iter() - .find(|column| column.data_type() == &DataType::Float64) - .expect("native/legacy result has one Float64 value column") + .filter(|column| column.data_type() == &DataType::Float64) + .collect::>(); + assert_eq!(value_columns.len(), 1, "expected one Float64 value column"); + let values = value_columns[0] .as_any() .downcast_ref::() .unwrap(); + assert_eq!(timestamps.null_count(), 0); + assert_eq!(hosts.null_count(), 0); + assert_eq!(values.null_count(), 0); (0..batch.num_rows()) .map(|row| { ( @@ -1860,64 +1870,31 @@ async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { fn assert_four_native_partitions(plan: &Arc) { use datafusion::physical_plan::metrics::MetricValue; - let mut stack = vec![plan.clone()]; - let mut fused_nodes = 0; - let mut scan_partition_counts = Vec::new(); - let mut values_by_partition = [0_u64; 4]; - while let Some(node) = stack.pop() { - if node.name().contains("FusedSeriesExec") { - fused_nodes += 1; - scan_partition_counts.push( - node.children() - .first() - .map(|child| child.output_partitioning().partition_count()) - .unwrap_or(0), - ); - assert_eq!( - node.output_partitioning().partition_count(), - 4, - "fused output must retain four partitions" - ); - } - if let Some(metrics) = node.metrics() { - for metric in metrics.iter() { - if metric.value().name().contains("native.root.output_values") { - if let MetricValue::Count { count, .. } = metric.value() { - if let Some(partition) = metric.partition() - && let Some(value) = values_by_partition.get_mut(partition) - { - *value += count.value() as u64; - } - } - } - } - } - stack.extend(node.children().into_iter().cloned()); - } - assert_eq!(fused_nodes, 1, "expected one fused execution node"); + assert_eq!(plan.name(), "FusedSeriesExec"); + assert_eq!(plan.output_partitioning().partition_count(), 4); assert_eq!( - scan_partition_counts, - [4], + plan.children()[0].output_partitioning().partition_count(), + 4, "fused scan must have four partitions" ); + let mut values_by_partition = [0_u64; 4]; + for metric in plan.metrics().expect("fused plan metrics").iter() { + if metric.value().name() == "native.root.output_values" + && let MetricValue::Count { count, .. } = metric.value() + && let Some(partition) = metric.partition() + && let Some(value) = values_by_partition.get_mut(partition) + { + *value += count.value() as u64; + } + } assert!( values_by_partition.iter().all(|value| *value > 0), "every native execution partition must produce values: {values_by_partition:?}" ); } -fn execution_plan_names(plan: &Arc) -> Vec { - let mut names = Vec::new(); - let mut stack = vec![plan.clone()]; - while let Some(node) = stack.pop() { - names.push(node.name().to_string()); - stack.extend(node.children().into_iter().cloned()); - } - names -} - async fn native_test_instance( - name: &str, + builder: GreptimeDbStandaloneBuilder, ) -> ( crate::standalone::GreptimeDbStandalone, Arc, @@ -1930,7 +1907,7 @@ async fn native_test_instance( experimental_native_promql: true, ..Default::default() }; - let standalone = GreptimeDbStandaloneBuilder::new(name) + let standalone = builder .with_plugin(plugins) .with_query_options(options) .build() @@ -1941,7 +1918,11 @@ async fn native_test_instance( #[tokio::test(flavor = "multi_thread")] async fn native_promql_executes_fused_node_and_matches_legacy() { common_telemetry::init_default_ut_logging(); - let (standalone, observer) = native_test_instance("native_promql_fused_e2e").await; + let mut mito_config = mito2::config::MitoConfig::default(); + mito_config.experimental_series_scan_v2 = true; + let builder = + GreptimeDbStandaloneBuilder::new("native_promql_fused_e2e").with_mito_config(mito_config); + let (standalone, observer) = native_test_instance(builder).await; let instance = standalone.fe_instance().clone(); execute_all( &instance, @@ -1985,7 +1966,10 @@ async fn native_promql_executes_fused_node_and_matches_legacy() { #[tokio::test(flavor = "multi_thread")] async fn native_promql_explain_and_analyze_execute_fused_plan() { common_telemetry::init_default_ut_logging(); - let (standalone, _observer) = native_test_instance("native_promql_explain_analyze_e2e").await; + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_explain_analyze_e2e", + )) + .await; let instance = standalone.fe_instance().clone(); execute_all(&instance, "CREATE TABLE native_explain (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_explain VALUES (0, 'a', 2)", QueryContext::arc()).await; for (statement, enabled) in [ @@ -1995,24 +1979,33 @@ async fn native_promql_explain_and_analyze_execute_fused_plan() { ("TQL ANALYZE VERBOSE (0, 0, '1s') native_explain", true), ("TQL EXPLAIN (0, 0, '1s') native_explain", false), ] { + observer.clear(); let output = native_tql_output(&instance, statement, enabled).await; let text = native_output_text(output).await; - if enabled { - assert!( - text.contains("FusedSeriesExec"), - "expected native physical tree for {statement}: {text}" - ); - } else { - assert!( - !text.contains("FusedSeriesExec"), - "flag-off EXPLAIN unexpectedly has native plan: {text}" - ); - } if statement.contains("ANALYZE") && enabled { - assert!( - text.contains("output_values"), - "ANALYZE must report executed native metrics: {text}" - ); + assert!(text.contains("FusedSeriesExec")); + assert!(text.contains("output_values")); + let plan = observer.take(); + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 1, "ANALYZE executes the fused series"); + } else if statement.contains("EXPLAIN") && enabled { + assert!(text.contains("FusedSeries: native_explain")); + assert!(text.contains("MergeScanExec")); + assert!(!text.contains("FusedSeriesExec")); + } else if !enabled { + assert!(!text.contains("FusedSeries:")); } } } @@ -2020,7 +2013,10 @@ async fn native_promql_explain_and_analyze_execute_fused_plan() { #[tokio::test(flavor = "multi_thread")] async fn native_promql_view_fallback_preserves_native_hint() { common_telemetry::init_default_ut_logging(); - let (standalone, observer) = native_test_instance("native_promql_view_fallback_e2e").await; + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_view_fallback_e2e", + )) + .await; let instance = standalone.fe_instance().clone(); execute_all(&instance, "CREATE TABLE native_view_source (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_view_source VALUES (0, 'a', 2)", QueryContext::arc()).await; let ctx = native_query_context(true, 4); @@ -2061,7 +2057,10 @@ async fn native_promql_view_fallback_preserves_native_hint() { #[tokio::test(flavor = "multi_thread")] async fn native_promql_metric_logical_tables_are_isolated() { common_telemetry::init_default_ut_logging(); - let (standalone, observer) = native_test_instance("native_promql_metric_isolation_e2e").await; + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_metric_isolation_e2e", + )) + .await; let instance = standalone.fe_instance().clone(); execute_all( &instance, @@ -2086,9 +2085,9 @@ async fn native_promql_metric_logical_tables_are_isolated() { QueryContext::arc(), ) .await; - for (table, host, value, foreign_host) in [ - ("native_metric_a", "only_a", 1.0, "only_b"), - ("native_metric_b", "only_b", 2.0, "only_a"), + for (table, host, value) in [ + ("native_metric_a", "only_a", 1.0), + ("native_metric_b", "only_b", 2.0), ] { let sql = format!("TQL EVAL (1, 1, '1s') rate({table}[2s])"); observer.clear(); @@ -2098,10 +2097,5 @@ async fn native_promql_metric_logical_tables_are_isolated() { let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; assert_eq!(native_rows, vec![(1000, host.to_string(), value)]); assert_eq!(native_rows, legacy_rows); - assert!( - native_rows - .iter() - .all(|(_, row_host, _)| row_host != foreign_host) - ); } } From ada66665f303c33536d068b4d0b7447a424edee5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 17:10:53 +0800 Subject: [PATCH 28/38] test(sql): cover textual TQL view definitions Signed-off-by: discord9 (cherry picked from commit 9df296e65b7e1c45b64c1e48906afa0f02013bce) --- .../cases/standalone/common/view/tql_eval.sql | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 tests/cases/standalone/common/view/tql_eval.sql diff --git a/tests/cases/standalone/common/view/tql_eval.sql b/tests/cases/standalone/common/view/tql_eval.sql new file mode 100644 index 00000000000..367eecb6bde --- /dev/null +++ b/tests/cases/standalone/common/view/tql_eval.sql @@ -0,0 +1,18 @@ +CREATE TABLE tql_view_source( + ts TIMESTAMP TIME INDEX, + host STRING, + val DOUBLE, + PRIMARY KEY (host) +); + +INSERT INTO tql_view_source VALUES (0, 'a', 2); + +CREATE VIEW tql_view AS TQL EVAL (0, 0, '1s', '5m') tql_view_source; + +SELECT * FROM tql_view; + +SHOW CREATE VIEW tql_view; + +DROP VIEW tql_view; + +DROP TABLE tql_view_source; From bb5a7ee540c919ede28d71dbb691a1bd63d37802 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 17:19:50 +0800 Subject: [PATCH 29/38] test(promql): clarify partition fixture and allow richer explain Signed-off-by: discord9 (cherry picked from commit ee5566d201774983e235716a67f2600b267c7c99) --- tests-integration/src/tests/promql_test.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 75fd38c387f..dbce5829110 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -1918,8 +1918,12 @@ async fn native_test_instance( #[tokio::test(flavor = "multi_thread")] async fn native_promql_executes_fused_node_and_matches_legacy() { common_telemetry::init_default_ut_logging(); - let mut mito_config = mito2::config::MitoConfig::default(); - mito_config.experimental_series_scan_v2 = true; + // The fixture opts in so its single bulk insert is distributed across real TSID scan + // partitions; legacy may send its one batch only to partition 0. + let mito_config = mito2::config::MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }; let builder = GreptimeDbStandaloneBuilder::new("native_promql_fused_e2e").with_mito_config(mito_config); let (standalone, observer) = native_test_instance(builder).await; @@ -2003,7 +2007,6 @@ async fn native_promql_explain_and_analyze_execute_fused_plan() { } else if statement.contains("EXPLAIN") && enabled { assert!(text.contains("FusedSeries: native_explain")); assert!(text.contains("MergeScanExec")); - assert!(!text.contains("FusedSeriesExec")); } else if !enabled { assert!(!text.contains("FusedSeries:")); } From afca3328f6a4f87c38cfb5d6b755a6ee1939db9f Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 17:27:07 +0800 Subject: [PATCH 30/38] test(sql): record generated TQL view regression results Signed-off-by: discord9 (cherry picked from commit 93fc2db0e4ee3eb46507865c8869d9426ac7ea2f) --- .../standalone/common/view/tql_eval.result | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 tests/cases/standalone/common/view/tql_eval.result diff --git a/tests/cases/standalone/common/view/tql_eval.result b/tests/cases/standalone/common/view/tql_eval.result new file mode 100644 index 00000000000..704ba763969 --- /dev/null +++ b/tests/cases/standalone/common/view/tql_eval.result @@ -0,0 +1,41 @@ +CREATE TABLE tql_view_source( + ts TIMESTAMP TIME INDEX, + host STRING, + val DOUBLE, + PRIMARY KEY (host) +); + +Affected Rows: 0 + +INSERT INTO tql_view_source VALUES (0, 'a', 2); + +Affected Rows: 1 + +CREATE VIEW tql_view AS TQL EVAL (0, 0, '1s', '5m') tql_view_source; + +Affected Rows: 0 + +SELECT * FROM tql_view; + ++---------------------+------+-----+ +| ts | host | val | ++---------------------+------+-----+ +| 1970-01-01T00:00:00 | a | 2.0 | ++---------------------+------+-----+ + +SHOW CREATE VIEW tql_view; + ++----------+---------------------------------------------------------------------+ +| View | Create View | ++----------+---------------------------------------------------------------------+ +| tql_view | CREATE VIEW tql_view AS TQL EVAL (0, 0, '1s', '5m') tql_view_source | ++----------+---------------------------------------------------------------------+ + +DROP VIEW tql_view; + +Affected Rows: 0 + +DROP TABLE tql_view_source; + +Affected Rows: 0 + From c6881447c1397766befa46a8f0c277c2c5d6eb93 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 18:53:24 +0800 Subject: [PATCH 31/38] fix(promql): bind native input to common scan ordering Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 15 +++++++-------- src/query/src/promql/planner/test.rs | 16 ++++++++++++++++ 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index d607eddfeef..7af48e9d393 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -80,14 +80,13 @@ pub(crate) fn native_promql_requested(query_ctx: &QueryContextRef) -> bool { /// The ordering this route compiles its trees against. /// -/// The logical planner cannot read the ordering a scanner will declare — it is decided when the -/// region scanner is prepared, which happens during physical planning — so the route names the -/// strongest ordering it can assume instead of guessing: `KeyThenTime` orders each series' chunks -/// contiguously *and* orders the series by identity, which `Contiguous` (the weaker declaration a -/// two-phase legacy scan makes) also satisfies. Nothing downstream enforces the tree's declared -/// ordering against the frame's own declaration beyond rejecting `Interleaved`, so a scanner that -/// declares `Contiguous` still feeds a tree compiled here. -const NATIVE_ORDERING: SeriesOrdering = SeriesOrdering::KeyThenTime; +/// Which ordering a raw-series scan declares is decided when the region scanner is prepared, during +/// physical planning, and the logical planner cannot read it: a legacy scan declares `KeyThenTime` +/// and a two-phase scan `Contiguous`. So the route names the weakest ordering both provide — +/// `Contiguous`, every series contiguous and the series themselves unordered — rather than +/// overclaiming the legacy scan's stronger `KeyThenTime`. A native evaluation only refuses +/// `Interleaved`, so a tree compiled here runs on either scanner. +const NATIVE_ORDERING: SeriesOrdering = SeriesOrdering::Contiguous; impl PromPlanner { /// Plans `prom_expr` as one [`FusedSeries`] node, or returns `None` to keep the existing plan. diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 1846aab2002..629712c4483 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -41,7 +41,9 @@ use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{Extension, UserDefinedLogicalNodeCore, col}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::series::SeriesOrdering; use promql::extension_plan::{FusedSeries, HistogramFold}; +use promql::native::plan::NativeExpr; use promql_parser::label::Labels; use promql_parser::parser; use session::context::QueryContext; @@ -1813,6 +1815,20 @@ async fn native_fused_rate_plans_to_a_fused_series_node() { .map(|field| field.name().as_str()) .collect::>() ); + // The scan below is chosen during physical planning, so the route cannot know whether it is + // the legacy scan (`KeyThenTime`) or the two-phase one (`Contiguous`) and must bind the weaker + // guarantee both provide. `rate` compiles to one window over that raw input. + let NativeExpr::Window { input, .. } = &fused.tree().root().expr else { + panic!("the root is not a window: {:?}", fused.tree().root().expr); + }; + let NativeExpr::Input(bound) = &input.expr else { + panic!("the window's input is not an input: {:?}", input.expr); + }; + assert_eq!( + SeriesOrdering::Contiguous, + bound.ordering, + "the tree claimed an ordering a raw-series scan need not guarantee" + ); } #[tokio::test] From 990d932198ab0c5e41b724a54d990e5a8ffc8acb Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 23 Oct 2026 20:06:13 +0800 Subject: [PATCH 32/38] fix(query): keep label rewrites on the existing plan Signed-off-by: discord9 --- src/query/src/promql/planner/native.rs | 6 +++++ src/query/src/promql/planner/test.rs | 35 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 7af48e9d393..d2d23c8a123 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -470,6 +470,12 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>) PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => out.push(vs), PromExpr::Paren(ParenExpr { expr }) => collect_selectors(expr, out), PromExpr::Call(call) => { + // A label rewrite adds a label the node's output schema cannot declare: that schema is + // built from the table's own tags, so folding `label_replace`/`label_join` in would + // silently drop the destination label the existing plan projects. + if matches!(call.func.name, "label_replace" | "label_join") { + return; + } for arg in &call.args.args { collect_selectors(arg, out); } diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 629712c4483..6a77ce48cc8 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -2199,6 +2199,41 @@ async fn native_fused_unsupported_shape_falls_back() { assert_eq!(display, without_flag.display_indent().to_string()); } +/// A label rewrite stays on the existing plan, which is the one that projects its new label. +/// +/// The native node's output schema is the table's own time index, fields and tags, so a tree that +/// rewrote labels would report a label set the existing plan's projection emits and the node cannot. +/// Each query below writes `newlabel`, which the fixture table does not have (its only tag is +/// `tag_0`), so a fused plan would visibly lose it. +#[tokio::test] +async fn native_fused_label_rewrites_keep_legacy_output_labels() { + for query in [ + // The rewrite at the statement root... + r#"label_replace(some_metric, "newlabel", "$1", "tag_0", "(.*)")"#, + r#"label_join(some_metric, "newlabel", "-", "tag_0", "tag_0")"#, + // ... and below another call, which is the same shape one level deeper. + r#"abs(label_replace(some_metric, "newlabel", "$1", "tag_0", "(.*)"))"#, + r#"abs(label_join(some_metric, "newlabel", "-", "tag_0", "tag_0"))"#, + ] { + let native = build_native_fused_plan(query, true).await; + let display = native.display_indent().to_string(); + assert!( + find_fused_series(&native).is_none(), + "`{query}` rewrites labels the node's schema cannot declare: {display}" + ); + assert!( + display.contains("newlabel"), + "`{query}` must keep the projection that emits its new label: {display}" + ); + let without_flag = build_native_fused_plan(query, false).await; + assert_eq!( + display, + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan, not an approximation of it" + ); + } +} + /// The flag off: no statement is offered to the native route at all. #[tokio::test] async fn native_fused_flag_off_keeps_the_existing_plan() { From ebfbd2ae5da1a2e02ce27db98ede723e5a4308d1 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 24 Oct 2026 18:58:51 +0800 Subject: [PATCH 33/38] fix(query): preserve legacy semantics at native admission Signed-off-by: discord9 --- src/promql/AGENTS.md | 23 +- src/query/src/dist_plan/commutativity.rs | 22 +- src/query/src/planner.rs | 121 +++++- src/query/src/promql/planner.rs | 19 - src/query/src/promql/planner/native.rs | 55 ++- src/query/src/promql/planner/test.rs | 197 ++++++++- src/sql/src/parsers/create_parser.rs | 61 +-- src/sql/src/statements/tql.rs | 2 +- tests-integration/src/tests/promql_test.rs | 392 +++++++++++++++++- .../standalone/common/view/tql_eval.result | 41 -- .../cases/standalone/common/view/tql_eval.sql | 18 - 11 files changed, 754 insertions(+), 197 deletions(-) delete mode 100644 tests/cases/standalone/common/view/tql_eval.result delete mode 100644 tests/cases/standalone/common/view/tql_eval.sql diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md index 3962add5e1b..2381461d3c0 100644 --- a/src/promql/AGENTS.md +++ b/src/promql/AGENTS.md @@ -34,8 +34,21 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`. requires one selector, a direct scan with a bound metric name, and no residual label matchers or `__field__` matcher that narrows the scan, plus an eligible raw-series table/field set and supported expression/grid. Metric scans retain - the static `__table_id` predicate. Hybrid CTEs and CREATE VIEW TQL creation - currently fall back to the existing planner. + the static `__table_id` predicate. Hybrid CTEs keep the existing plan. The + operator also disables native planning for a programmatic TQL view so its + persisted Substrait plan remains serializable; `CREATE VIEW ... AS TQL EVAL` + syntax is a separate change. +- **Admission and SQL output:** the guards live in + `src/query/src/promql/planner/native.rs`. Non-`Float64` fields, alternative + (`or`) matchers and `rate`/`increase` on temporality-tagged tables keep the + existing plan; the latter preserves legacy per-series delta summation. + Allowing query fallback through the engine option, session variable or + `query_fallback` hint also disables native admission. Native SQL/TQL output + has column order `time index, fields, tags`, raw field names with nullable + `Float64` values, and nullable `Utf8` tags (including integer table tags). + Unlike a legacy raw selector, it does not preserve tag types; legacy function + expressions can also use different field names and column order. The native + time index is a non-nullable millisecond timestamp without a time zone. ## Change coupling @@ -45,9 +58,9 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`. (`src/query/src/dist_plan/commutativity.rs`). Process-local `FusedSeries` is the exception: it is not serialized to a datanode. Its distributed placement is restricted to a direct table scan, a bound metric name and trees that read - one series at a time. The single-region placement warning in - `commutativity.rs` remains important: do not assume a series stays in one - region. + one series at a time, and it is only correct because native admission refuses a + table whose partition columns are not all tag columns: that is what keeps a + series inside one region. Keep that admission and this placement together. - A new logical node also needs its physical planner reachable from `DfQueryPlanner::create_physical_plan`, assembled in `src/query/src/query_engine/state.rs`. diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index c11e78cf819..4221b009390 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -263,20 +263,14 @@ impl Categorizer { return Commutativity::NonCommutative; } - // TEMPORARY, and this is the one thing here that is not proved: the placement is - // only correct while each series reaches exactly one region. A region evaluates the - // whole tree over the series its scan returns, so a series whose samples are split - // across regions is evaluated per fragment, and a range function over a fragment is - // not the range function over the series. Nothing checks that a series stays in one - // region: the storage partition key is whatever `partition_column_names()` names and - // may be a column that varies within one series, and this categorizer runs before - // region selection, so it has no selected-region vector to consult either. A single - // region deployment satisfies it; a partitioned table over several regions does not, - // and would return a wrong answer rather than an error. - // - // What would make this unconditional is a partition key the label set determines, - // checked against the table's physical partition expressions. Until then the - // placement is a proof of concept over one region. + // A region evaluates the whole tree over the series its scan returns, so this is + // only correct while each series reaches exactly one region: a series split + // across regions would be ranged per fragment, and a range over fragments is not + // the range over the series. Native admission + // (`promql::planner::native::try_plan_native_fused`) only builds the node when + // every partition column of the table is one of the binding's tag columns, so + // the partition key is a function of the label set and a series cannot straddle + // regions. Commutativity::Commutative } name if name == SeriesNormalize::name() diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index de45e55e8d8..a70e0108817 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -50,6 +50,8 @@ use sql::statements::statement::Statement; use sql::statements::tql::Tql; use sqlparser::ast::{AccessExpr, Value, visit_expressions_mut}; +use crate::datafusion::QUERY_FALLBACK_HINT; +use crate::dist_plan::DistPlannerOptions; use crate::error::{ CteColumnSchemaMismatchSnafu, PlanSqlSnafu, QueryPlanSnafu, Result, SqlSnafu, UnimplementedSnafu, @@ -323,6 +325,20 @@ impl DfLogicalPlanner { native_route_allowed: bool, ) -> Result { let mut scheduled_state = self.derive_session_state_with_scheduled_time(&query_ctx)?; + // Whether the distributed planner may fall back to the fallback plan rewriter. The + // execution engine resolves the same three sources the same way: the engine option, which + // it copies into the session state it plans with, and then, per query, the session + // variable or the `query_fallback` hint. Each source only ever turns the fallback on, so + // the value they end up with is the disjunction of them. + let allow_query_fallback = scheduled_state + .config_options() + .extensions + .get::() + .is_some_and(|options| options.allow_query_fallback) + || query_ctx.configuration_parameter().allow_query_fallback() + || query_ctx + .extension(QUERY_FALLBACK_HINT) + .is_some_and(|fallback| fallback.to_lowercase().parse::().unwrap_or(false)); let promql_annotations = query_ctx.remote_query_id().map(promql_annotation_collector); if let Some(collector) = &promql_annotations { scheduled_state @@ -350,7 +366,13 @@ impl DfLogicalPlanner { stmt, &self.engine_state, promql_annotations, + // A query allowed to fall back keeps the existing plan. The native node's physical + // planner needs the bare region scan as its input; the fallback rewriter instead wraps + // every table scan in `MergeScan`, so the node above it would have no physical plan and + // the query would fail after the fallback was already taken. Declining here keeps the + // whole query on the legacy path. native_route_allowed + && !allow_query_fallback && self.engine_state.experimental_native_promql() && self.engine_state.in_process_region_transport() && crate::promql::planner::native::native_promql_requested(&query_ctx), @@ -778,7 +800,9 @@ mod tests { use table::test_util::EmptyTable; use super::*; + use crate::datafusion::QUERY_FALLBACK_HINT; use crate::parser::{PromQuery, QueryLanguageParser}; + use crate::promql::planner::NATIVE_PROMQL_EXTENSION_KEY; use crate::{QueryEngineFactory, QueryEngineRef}; async fn create_test_engine() -> QueryEngineRef { @@ -837,6 +861,14 @@ mod tests { } fn create_promql_test_engine() -> QueryEngineRef { + create_promql_test_engine_with(crate::options::QueryOptions::default(), false) + } + + /// The PromQL fixture, additionally stating the engine options and how its regions are reached. + fn create_promql_test_engine_with( + query_options: crate::options::QueryOptions, + in_process_region_transport: bool, + ) -> QueryEngineRef { let catalog_manager = MemoryCatalogManager::with_default_setup(); let physical_table_name = "phy"; let physical_table_id = 999u32; @@ -927,15 +959,19 @@ mod tests { }) .unwrap(); - QueryEngineFactory::new( + QueryEngineFactory::try_new_with_plugins_and_region_transport( catalog_manager, None, None, None, None, + None, false, - crate::options::QueryOptions::default(), + in_process_region_transport, + Default::default(), + query_options, ) + .unwrap() .query_engine() } @@ -968,6 +1004,87 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `rate(some_metric[5m])` on the fixture above with the fallback sources the caller + /// states — the engine option, the session variable and the `query_fallback` hint — while the + /// native route is enabled and requested, so only the fallback guard decides. + async fn plan_native_promql_with_fallback( + engine_fallback: bool, + session_fallback: bool, + hint: Option<&str>, + ) -> String { + let engine = create_promql_test_engine_with( + crate::options::QueryOptions { + allow_query_fallback: engine_fallback, + experimental_native_promql: true, + ..Default::default() + }, + // The native route plans a region input the region server in this process accepts. + true, + ); + // The native route request and the fallback hint both reach the planner as query context + // extensions, the way a client sends them. + let mut builder = QueryContextBuilder::default() + .set_extension(NATIVE_PROMQL_EXTENSION_KEY.to_string(), "true".to_string()); + if let Some(hint) = hint { + builder = builder.set_extension(QUERY_FALLBACK_HINT.to_string(), hint.to_string()); + } + let query_ctx = builder.build(); + query_ctx + .configuration_parameter() + .set_allow_query_fallback(session_fallback); + let query_ctx = Arc::new(query_ctx); + let stmt = QueryLanguageParser::parse_promql( + &PromQuery { + query: "rate(some_metric[5m])".to_string(), + start: "0".to_string(), + end: "10".to_string(), + step: "5s".to_string(), + lookback: "300s".to_string(), + alias: None, + }, + &query_ctx, + ) + .unwrap(); + + engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent() + .to_string() + } + + /// The native route is admitted only when the query cannot fall back, and every source the + /// execution engine resolves the fallback from decides that at plan time too. + #[tokio::test] + async fn test_plan_pql_declines_native_route_when_query_may_fall_back() { + // (engine option, session variable, hint, native route admitted) + for (engine, session, hint, native) in [ + // Nothing asks for the fallback, so the route is planned as before. + (false, false, None, true), + // A hint that does not parse as `true` leaves the fallback off, as it does at execution. + (false, false, Some("false"), true), + (false, false, Some("yes"), true), + // Either per-query source turns the fallback on, so the route is declined. + (false, true, None, false), + (false, false, Some("true"), false), + (false, false, Some("TRUE"), false), + // Each source only turns the fallback on: neither a `false` hint nor a `false` session + // variable clears the engine option, and a `false` hint does not clear the variable. + (true, false, None, false), + (true, false, Some("false"), false), + (true, true, Some("false"), false), + ] { + let plan = plan_native_promql_with_fallback(engine, session, hint).await; + assert_eq!( + plan.contains("FusedSeries"), + native, + "engine={engine}, session={session}, hint={hint:?}: {plan}" + ); + } + } + /// Plans `sql` and runs the DataFusion analyzer, which is where /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these /// assertions would not see the assignment rewrite at all. diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 05acc1a9de5..18bb5f2c316 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -5517,25 +5517,6 @@ impl PromPlanner { .map(|idx| schema.field(idx).data_type()) } - /// Whether every field column carries a type the region's raw series preflight serves. - /// - /// The preflight decides on `ConcreteDataType::is_numeric`, which excludes the decimals, so - /// the Arrow type is converted rather than asked: Arrow's own `is_numeric` includes - /// `Decimal128`/`Decimal256` (and `Float16`, which has no column type here), and a decimal - /// field would then pass here and fail once the logical plan is fixed. Anything else - a - /// string, a boolean, a timestamp - has no raw framing either. - fn field_columns_are_raw_series_servable( - schema: &DFSchemaRef, - field_columns: &[String], - ) -> bool { - field_columns.iter().all(|field| { - Self::field_column_type(schema, field).is_some_and(|data_type| { - datatypes::prelude::ConcreteDataType::from_arrow_type(data_type).is_numeric() - || data_type == &Self::native_histogram_arrow_type() - }) - }) - } - fn field_column_is_native_histogram(schema: &DFSchemaRef, field_column: &str) -> bool { Self::field_column_type(schema, field_column) .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index d2d23c8a123..9cd23a87cc3 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -31,6 +31,7 @@ use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE, MITO2_ENGINE}; +use common_query::prelude::OTLP_AGGREGATION_TEMPORALITY_LABEL; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as ArrowTimeUnit}; use datafusion::common::{DFSchema, DFSchemaRef}; use datafusion::error::DataFusionError; @@ -102,8 +103,9 @@ impl PromPlanner { /// caller plans it with the existing planner as if this method did not exist. The two ways that /// happens are a shape the compiler refuses (`DataFusionError::NotImplemented`, which names /// what it found) and a query this route cannot bind faithfully — more than one selector, a - /// label matcher the bare scan below the node would not apply, an `__field__` matcher that - /// narrows the field set, a missing table, a native histogram field, or an unusable grid. + /// label matcher the bare scan below the node would not apply, alternative (`or`) matchers, an + /// `__field__` matcher that narrows the field set, a missing table, a field that is not + /// `Float64`, a `rate`/`increase` over a delta-temporality table, or an unusable grid. pub(super) async fn try_plan_native_fused( &mut self, prom_expr: &PromExpr, @@ -125,6 +127,11 @@ impl PromPlanner { return Ok(None); }; + // These matchers are alternatives the one scan below the node cannot represent, so the + // query keeps the plan the existing path builds for it. + if !selector.matchers.or_matchers.is_empty() { + return Ok(None); + } // Resolve the selector the way the existing planner does, and take what is left as the // rows it selects. The scan below the node is the selector's own scan, unfiltered, so a // matcher that survives preprocessing selects fewer rows here than the scan reads. @@ -155,6 +162,18 @@ impl PromPlanner { if fields.is_empty() { return Ok(None); } + // On a delta-temporality table a series' samples are increments, which the existing + // planner sums over the window for `increase` and divides by the range for `rate` + // ([`PromPlanner::select_delta_range_math`]), per series because the tag can differ between + // them. A compiled tree computes both for a cumulative series, so any `rate`/`increase` + // over such a table keeps the existing plan; a bare selector still folds. + if tags + .iter() + .any(|tag| tag == OTLP_AGGREGATION_TEMPORALITY_LABEL) + && calls_rate_or_increase(prom_expr) + { + return Ok(None); + } // The selector's own table scan, built by the same method the existing path uses: the // physical planner turns it into the region scan the node reads, and building a scan here @@ -258,17 +277,13 @@ impl PromPlanner { }; scan.filters.push(filter); } - // A window over a native histogram yields a histogram, which the output schema below - // cannot declare; the existing path evaluates those with its own kernels. - if Self::field_columns_contain_native_histogram(input.schema(), &fields) { - return Ok(None); - } - // The region's raw series preflight serves a field only when it is numeric or a native - // histogram, and the physical layer turns that refusal into a planning error because the - // logical plan is already fixed by then. A table whose selected fields are, say, a string - // and a boolean therefore has to be declined here, while the existing plan is still the - // fallback - `stringsfield` over such a table answered normally before this route existed. - if !Self::field_columns_are_raw_series_servable(input.schema(), &fields) { + // The region's columnar series-index samples, one of the readers this scan may be served + // by, store `Float64` fields only, and the logical plan cannot tell which reader arrives. + // So every field has to be `Float64`: an integer, a `Float32`, a decimal, a string or a + // native histogram field keeps the existing plan, which the physical layer still serves. + if !fields.iter().all(|field| { + Self::field_column_type(input.schema(), field) == Some(&ArrowDataType::Float64) + }) { return Ok(None); } let Some(stmt) = self.native_statement(prom_expr) else { @@ -484,6 +499,20 @@ fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>) } } +/// Whether `expr` calls `rate`/`increase`, the functions the existing planner rewrites per series +/// for a delta-temporality table. Only nested calls and parentheses can hold one in an expression +/// this route accepts; every other shape names no single selector or reads across series. +fn calls_rate_or_increase(expr: &PromExpr) -> bool { + match expr { + PromExpr::Call(call) => { + matches!(call.func.name, "rate" | "increase") + || call.args.args.iter().any(|arg| calls_rate_or_increase(arg)) + } + PromExpr::Paren(ParenExpr { expr }) => calls_rate_or_increase(expr), + _ => false, + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 6a77ce48cc8..2e22ee59003 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -972,6 +972,73 @@ async fn native_engine_provider(partition_key_indices: Vec) -> DfTableSou .await } +/// A delta-temporality table on the engine the native route accepts. +/// +/// The OTLP receiver writes the aggregation temporality as a tag of every series, so a table with +/// that tag can hold delta and cumulative series side by side: the existing planner sums the +/// increments of a delta series instead of rating its counter. The value field is the `Float64` +/// this route folds. +async fn native_engine_temporality_marker_provider() -> DfTableSourceProvider { + let catalog_list = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + OTLP_AGGREGATION_TEMPORALITY_LABEL.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + greptime_timestamp().to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .engine(NATIVE_ENGINE.to_string()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::default() + .name("some_metric") + .meta(table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + assert!( + catalog_list + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "some_metric".to_string(), + table_id: 1024, + table, + }) + .is_ok() + ); + + DfTableSourceProvider::new( + catalog_list, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) +} + #[allow(clippy::too_many_arguments)] async fn build_test_table_provider_with_options( table_name_tuples: &[(String, String)], @@ -2132,19 +2199,23 @@ async fn native_fused_engine_without_raw_series_falls_back() { ); } -/// A value field the region's raw scan cannot serve keeps the existing plan. +/// A value field the raw-series readers cannot both serve keeps the existing plan. /// -/// The preflight serves numeric columns and native histograms only, and the physical layer turns -/// its refusal into a planning error once the logical plan is fixed. A bare selector over a string -/// field answered normally before this route existed, so it must keep doing that. +/// The region's preflight serves numeric columns and native histograms, while the columnar +/// series-index samples the scan may be read through store `Float64` fields only, and the logical +/// plan cannot tell which reader the region prepares. So only `Float64` is admitted; a bare +/// selector over any other field type answered normally before this route existed. #[tokio::test] -async fn native_fused_non_numeric_field_falls_back() { +async fn native_fused_unsupported_field_type_falls_back() { for field_type in [ ConcreteDataType::string_datatype(), ConcreteDataType::boolean_datatype(), // Arrow's `is_numeric` includes the decimals while the region's preflight does not, so // asking Arrow would let a decimal field through to a planning error. ConcreteDataType::decimal128_datatype(10, 2), + // Numeric, but not the `Float64` the columnar samples store. + ConcreteDataType::int64_datatype(), + ConcreteDataType::float32_datatype(), ] { let native = build_native_fused_plan_with_provider( "some_metric", @@ -2170,6 +2241,122 @@ async fn native_fused_non_numeric_field_falls_back() { } } +/// The `Float64` field every raw reader serves controls the field-type decline above. +#[tokio::test] +async fn native_fused_float64_field_is_accepted() { + let plan = build_native_fused_plan_with_provider( + "some_metric", + true, + build_test_table_provider_with_field_type(ConcreteDataType::float64_datatype()).await, + ) + .await; + + assert!( + find_fused_series(&plan).is_some(), + "a `Float64` field must still fold: {}", + plan.display_indent() + ); +} + +/// Alternative (`or`) matchers keep the existing plan. +/// +/// `a{x="1" or x="2"}` names alternative matchers, which the one scan below the node cannot +/// represent. These queries are only asserted to fall back: what the existing plan answers for +/// them is its own question, not one this route claims to have fixed. +#[tokio::test] +async fn native_fused_or_matchers_fall_back() { + for query in [ + r#"some_metric{tag_0="a" or tag_0="b"}"#, + r#"rate(some_metric{tag_0="a" or tag_0="b"}[5m])"#, + ] { + let native = build_native_fused_plan(query, true).await; + let without_flag = build_native_fused_plan(query, false).await; + assert!( + find_fused_series(&native).is_none(), + "`{query}` names alternative matchers: {}", + native.display_indent() + ); + assert_eq!( + native.display_indent().to_string(), + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan" + ); + } +} + +/// A `rate`/`increase` over a delta-temporality table keeps the existing plan, which sums each +/// series' increments. +/// +/// The samples of a series tagged `delta` are increments, which the existing planner sums over the +/// window for `increase` and divides by the range for `rate`, while a cumulative series is rated +/// from its counter; it picks that per series. A compiled tree rates every series the same way, so +/// these queries fall back. The same queries over a table without the marker still fold, which is +/// what makes the decline the marker's doing rather than the shape's. +#[tokio::test] +async fn native_fused_delta_temporality_rates_fall_back() { + for query in [ + "rate(some_metric[5m])", + "increase(some_metric[5m])", + // Below another call and inside parentheses, which are the same decline one level down. + "abs(rate(some_metric[5m]))", + "abs(increase(some_metric[5m]))", + "(rate(some_metric[5m]))", + ] { + let marked = build_native_fused_plan_with_provider( + query, + true, + native_engine_temporality_marker_provider().await, + ) + .await; + let marked_without_flag = build_native_fused_plan_with_provider( + query, + false, + native_engine_temporality_marker_provider().await, + ) + .await; + assert!( + find_fused_series(&marked).is_none(), + "`{query}` needs the per-series delta math: {}", + marked.display_indent() + ); + assert_eq!( + marked.display_indent().to_string(), + marked_without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan, not the counter form for every series" + ); + + let unmarked = build_native_fused_plan_with_provider( + query, + true, + native_engine_provider(vec![]).await, + ) + .await; + assert!( + find_fused_series(&unmarked).is_some(), + "`{query}` over a table without the marker must still fold: {}", + unmarked.display_indent() + ); + } +} + +/// The marker alone does not decline a query: only `rate`/`increase` are refused. +#[tokio::test] +async fn native_fused_delta_temporality_non_counter_math_still_folds() { + for query in ["some_metric", "abs(some_metric)"] { + let plan = build_native_fused_plan_with_provider( + query, + true, + native_engine_temporality_marker_provider().await, + ) + .await; + assert!( + find_fused_series(&plan).is_some(), + "`{query}` needs no per-series delta math and must still fold: {}", + plan.display_indent() + ); + } +} + /// A whole-vector operation is not a per-series tree, so it stays on the existing plan. #[tokio::test] async fn native_fused_whole_vector_operation_falls_back() { diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index de9c5657358..20eba9f2087 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -55,7 +55,6 @@ use crate::statements::create::{ CreateTableLike, CreateView, Partitions, SqlOrTql, TableConstraint, VECTOR_OPT_DIM, }; use crate::statements::statement::Statement; -use crate::statements::tql::Tql; use crate::statements::transform::type_alias::get_data_type_by_alias_name; use crate::statements::{OptionMap, sql_data_type_to_concrete_data_type}; use crate::util::{OptionValue, location_to_index, parse_option_string}; @@ -149,18 +148,7 @@ impl<'a> ParserContext<'a> { .expect_keyword(Keyword::AS) .context(SyntaxSnafu)?; - let query = match self.parser.peek_token().token { - Token::Word(word) - if word.quote_style.is_none() - && word.value.eq_ignore_ascii_case(tql_parser::TQL) => - { - match self.parse_tql(false)? { - Statement::Tql(tql @ Tql::Eval(_)) => Statement::Tql(tql), - _ => return self.expected("TQL EVAL", self.parser.peek_token()), - } - } - _ => self.parse_query()?, - }; + let query = self.parse_query()?; Ok(Statement::CreateView(CreateView { name: view_name, @@ -3327,53 +3315,6 @@ CREATE TABLE spans ( } } - #[test] - fn test_parse_create_view_tql_eval() { - for (query, expected_lookback) in [ - ("TQL EVAL (0, 15, '5s') metric", None), - ("TQL EVAL (0, 15, '5s', '5m') metric", Some("5m")), - ("TQL EVAL (0, 15, '5s', '300') metric", Some("300")), - ] { - let sql = format!("CREATE VIEW test AS {query}"); - let result = ParserContext::create_with_dialect( - &sql, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .unwrap(); - let Statement::CreateView(view) = &result[0] else { - panic!("expected CREATE VIEW, got {:?}", result[0]); - }; - let Statement::Tql(Tql::Eval(eval)) = view.query.as_ref() else { - panic!("expected TQL EVAL, got {}", view.query); - }; - assert_eq!(eval.lookback.as_deref(), expected_lookback); - - let formatted = result[0].to_string(); - let reparsed = ParserContext::create_with_dialect( - &formatted, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .unwrap(); - assert_eq!(reparsed[0], result[0], "{formatted}"); - } - - for sql in [ - "CREATE VIEW test AS TQL EXPLAIN (0, 15, '5s') metric", - "CREATE VIEW test AS TQL ANALYZE (0, 15, '5s') metric", - ] { - assert!( - ParserContext::create_with_dialect( - sql, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .is_err() - ); - } - } - #[test] fn test_parse_create_view_with_columns() { let sql = "CREATE VIEW test () AS SELECT * FROM NUMBERS"; diff --git a/src/sql/src/statements/tql.rs b/src/sql/src/statements/tql.rs index 948c31ddd8d..835a957eada 100644 --- a/src/sql/src/statements/tql.rs +++ b/src/sql/src/statements/tql.rs @@ -46,7 +46,7 @@ fn format_tql( ) -> std::fmt::Result { write!(f, "({start}, {end}, '{step}'")?; if let Some(lookback) = lookback { - write!(f, ", '{lookback}'")?; + write!(f, ", {lookback}")?; } if let Some(alias) = alias { write!(f, ") {query} AS {alias}") diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index dbce5829110..7f6e900bc1d 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -20,6 +20,7 @@ use api::v1::{ ColumnDataType, ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, SemanticType, }; use common_base::Plugins; +use common_query::prelude::{GREPTIME_TEMPORALITY_DELTA, OTLP_AGGREGATION_TEMPORALITY_LABEL}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::{Output, OutputData}; use common_recordbatch::util::collect_batches; @@ -30,7 +31,7 @@ use datatypes::arrow::array::{ use datatypes::arrow::compute::cast; use datatypes::arrow::datatypes::DataType; use frontend::instance::Instance; -use query::datafusion::QUERY_PARALLELISM_HINT; +use query::datafusion::{QUERY_FALLBACK_HINT, QUERY_PARALLELISM_HINT}; use query::options::QueryOptions; use query::parser::{PromQuery, QueryLanguageParser, QueryStatement}; use query::physical_wrapper::{PhysicalPlanWrapper, PhysicalPlanWrapperRef}; @@ -39,6 +40,10 @@ use rstest::rstest; use rstest_reuse::apply; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; +use sql::ast::{Ident, ObjectName, ObjectNamePart}; +use sql::statements::create::CreateView; +use sql::statements::statement::Statement as SqlStatement; +use sql::statements::tql::{Tql, TqlEval}; use super::test_util::{ both_instances_cases, check_unordered_output_stream, distributed, standalone, @@ -1774,6 +1779,11 @@ impl NativePlanObserver { .take() .expect("native fused plan was captured") } + + /// The fused node the last wrapped plan contained, if any; `take` requires one instead. + fn fused(&self) -> Option> { + self.0.lock().unwrap().clone() + } } impl PhysicalPlanWrapper for NativePlanObserver { @@ -1867,6 +1877,90 @@ async fn native_result_rows(output: Output) -> Vec<(i64, String, f64)> { rows } +/// [`native_result_rows`] plus each row's OTLP temporality marker, sorted by label then time. +async fn native_marker_rows(output: Output) -> Vec<(i64, String, String, f64)> { + let batches = match output.data { + OutputData::Stream(stream) => collect_batches(stream).await.unwrap(), + OutputData::RecordBatches(batches) => batches, + _ => unreachable!("TQL EVAL returns record batches"), + }; + let mut rows = Vec::new(); + for batch in &batches { + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let hosts = cast( + batch.column_by_name("host").unwrap().as_ref(), + &DataType::Utf8, + ) + .unwrap(); + let hosts = hosts.as_any().downcast_ref::().unwrap(); + let markers = cast( + batch + .column_by_name(OTLP_AGGREGATION_TEMPORALITY_LABEL) + .unwrap() + .as_ref(), + &DataType::Utf8, + ) + .unwrap(); + let markers = markers.as_any().downcast_ref::().unwrap(); + let values = batch + .columns() + .iter() + .filter(|column| column.data_type() == &DataType::Float64) + .collect::>(); + assert_eq!(values.len(), 1, "expected one Float64 value column"); + let values = values[0].as_any().downcast_ref::().unwrap(); + assert_eq!( + markers.null_count(), + 0, + "every row must keep its temporality label" + ); + assert_eq!(values.null_count(), 0); + for row in 0..batch.num_rows() { + rows.push(( + timestamps.value(row), + hosts.value(row).to_string(), + markers.value(row).to_string(), + values.value(row), + )); + } + } + rows.sort_by(|left, right| (left.1.as_str(), left.0).cmp(&(right.1.as_str(), right.0))); + rows +} + +/// Compares the whole answer against `expected`, in the same label order, key by key and value by +/// value; a duplicated or missing series therefore fails on the keys rather than on the count. +fn assert_delta_rows( + rows: &[(i64, String, String, f64)], + expected: &[(i64, &str, &str, f64)], + context: &str, +) { + let actual = rows + .iter() + .map(|(ts, host, marker, value)| (*ts, host.as_str(), marker.as_str(), *value)) + .collect::>(); + assert_eq!(actual.len(), expected.len(), "{context}: rows {actual:?}"); + for (actual, expected) in actual.iter().zip(expected) { + assert_eq!( + (actual.0, actual.1, actual.2), + (expected.0, expected.1, expected.2), + "{context}: row {actual:?}" + ); + assert!( + (actual.3 - expected.3).abs() <= 1e-9 * expected.3.abs().max(1.0), + "{context}: series {:?} must report {}, found {}", + actual.1, + expected.3, + actual.3 + ); + } +} + fn assert_four_native_partitions(plan: &Arc) { use datafusion::physical_plan::metrics::MetricValue; @@ -1893,20 +1987,35 @@ fn assert_four_native_partitions(plan: &Arc ( crate::standalone::GreptimeDbStandalone, Arc, +) { + native_test_instance_with_options( + builder, + QueryOptions { + experimental_native_promql: true, + ..Default::default() + }, + ) + .await +} + +/// [`native_test_instance`] with the engine-wide query options spelled out by the caller. +async fn native_test_instance_with_options( + builder: GreptimeDbStandaloneBuilder, + options: QueryOptions, +) -> ( + crate::standalone::GreptimeDbStandalone, + Arc, ) { let observer = Arc::new(NativePlanObserver::default()); let plugins = Plugins::new(); let wrapper: PhysicalPlanWrapperRef = observer.clone(); plugins.insert(wrapper); - let options = QueryOptions { - experimental_native_promql: true, - ..Default::default() - }; let standalone = builder .with_plugin(plugins) .with_query_options(options) @@ -2024,12 +2133,30 @@ async fn native_promql_view_fallback_preserves_native_hint() { execute_all(&instance, "CREATE TABLE native_view_source (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_view_source VALUES (0, 'a', 2)", QueryContext::arc()).await; let ctx = native_query_context(true, 4); let expected_extensions = ctx.extensions(); - execute_all( - &instance, - "CREATE VIEW native_view AS TQL EVAL (0, 0, '1s', '5m') native_view_source", - ctx.clone(), - ) - .await; + // The AST the parser would build for `CREATE VIEW native_view AS TQL EVAL (0, 0, '1s', '5m') + // native_view_source`; the statement's spelling is covered by the separate view case. + let create_view = SqlStatement::CreateView(CreateView { + name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new("native_view"))]), + columns: vec![], + query: Box::new(SqlStatement::Tql(Tql::Eval(TqlEval { + start: "0".to_string(), + end: "0".to_string(), + step: "1s".to_string(), + lookback: Some("5m".to_string()), + query: "native_view_source".to_string(), + alias: None, + }))), + or_replace: false, + if_not_exists: false, + }); + instance + .statement_executor() + .execute_stmt(QueryStatement::Sql(create_view), ctx.clone()) + .await + .unwrap(); + // The view's stored plan must not depend on the caller's native hint, so reading it runs the + // plan the view was created with rather than one that needs the hint to be readable. + observer.clear(); let view = instance .do_query("SELECT * FROM native_view", ctx.clone()) .await @@ -2038,14 +2165,11 @@ async fn native_promql_view_fallback_preserves_native_hint() { let view_rows = native_result_rows(view).await; assert_eq!(ctx.extensions(), expected_extensions); assert_eq!(view_rows, vec![(0, "a".to_string(), 2.0)]); - let create_view = instance - .do_query("SHOW CREATE VIEW native_view", ctx.clone()) - .await - .remove(0) - .unwrap(); - let create_view_text = native_output_text(create_view).await; - assert!(create_view_text.contains("TQL EVAL")); - assert!(create_view_text.contains("5m")); + assert!( + observer.fused().is_none(), + "a view plan must be readable without the native hint" + ); + // The hint is still the caller's: the same context keeps folding a direct query. observer.clear(); let direct = instance .do_query("TQL EVAL (0, 0, '1s') native_view_source", ctx) @@ -2102,3 +2226,233 @@ async fn native_promql_metric_logical_tables_are_isolated() { assert_eq!(native_rows, legacy_rows); } } + +/// OTLP delta temporality on the public TQL route: a `delta`-marked series answers the raw sample +/// sum (10 + 10 = 20) while a cumulative series with the same samples answers 0. The fused node +/// computes one cumulative window for every series, so it must decline a marker-bearing table +/// instead of reporting 0 for the delta series. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_delta_temporality_rate_and_increase_match_legacy() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_delta_temporality"; + const CONTROL: &str = "native_delta_control"; + // Samples at 60s and 120s, evaluated at 120s over a 2m window: the `delta` series' samples are + // raw deltas, so its increase is their sum 10 + 10; the `cumulative` one did not rise, so 0. + const DELTA_INCREASE: f64 = 20.0; + const CUMULATIVE_INCREASE: f64 = 0.0; + const DELTA_RATE: f64 = 20.0 / 120.0; + const CUMULATIVE_RATE: f64 = 0.0; + + let (standalone, observer) = native_test_instance(GreptimeDbStandaloneBuilder::new( + "native_promql_delta_temporality", + )) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + r#"CREATE TABLE {TABLE} ( + host STRING, + ts TIMESTAMP(3) TIME INDEX, + val DOUBLE, + otlp_aggregation_temporality STRING, + PRIMARY KEY (host, otlp_aggregation_temporality) + ) ENGINE=mito; + INSERT INTO {TABLE} (host, ts, val, otlp_aggregation_temporality) VALUES + ('delta', 60000, 10, '{GREPTIME_TEMPORALITY_DELTA}'), + ('delta', 120000, 10, '{GREPTIME_TEMPORALITY_DELTA}'), + ('cumulative', 60000, 10, 'cumulative'), + ('cumulative', 120000, 10, 'cumulative'); + CREATE TABLE {CONTROL} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=mito; + INSERT INTO {CONTROL} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30);"# + ), + QueryContext::arc(), + ) + .await; + + for (query, delta, cumulative) in [ + ( + format!("increase({TABLE}[2m])"), + DELTA_INCREASE, + CUMULATIVE_INCREASE, + ), + (format!("rate({TABLE}[2m])"), DELTA_RATE, CUMULATIVE_RATE), + ( + format!("abs(increase({TABLE}[2m]))"), + DELTA_INCREASE, + CUMULATIVE_INCREASE, + ), + ( + format!("abs(rate({TABLE}[2m]))"), + DELTA_RATE, + CUMULATIVE_RATE, + ), + ] { + let sql = format!("TQL EVAL (120, 120, '1m') {query}"); + // Both series are in the window, in label order, with their own marker and value. + let expected = [ + (120_000, "cumulative", "cumulative", cumulative), + (120_000, "delta", GREPTIME_TEMPORALITY_DELTA, delta), + ]; + // Values before routes: a wrong number is the regression, a declined plan is not. + let legacy_rows = native_marker_rows(native_tql_output(&instance, &sql, false).await).await; + assert_delta_rows(&legacy_rows, &expected, &format!("native-off `{query}`")); + observer.clear(); + let native_rows = native_marker_rows(native_tql_output(&instance, &sql, true).await).await; + assert_delta_rows(&native_rows, &expected, &format!("native-on `{query}`")); + assert_eq!( + native_rows, legacy_rows, + "`{query}`: same rows on both routes" + ); + assert!( + observer.fused().is_none(), + "`{query}`: the native route must decline a table that carries the temporality marker" + ); + } + + // The control: a marker-free table (samples 10 and 30 here, unlike the table above) still fuses + // and agrees with the existing plan, so the decline is about the marker, not the table shape. + let control_sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({CONTROL}[2m])"); + observer.clear(); + let control_native = + native_result_rows(native_tql_output(&instance, &control_sql, true).await).await; + assert_eq!( + control_native, + vec![(120_000, "a".to_string(), 40.0)], + "the fused control reports the sampled sum" + ); + assert_eq!(observer.take().name(), "FusedSeriesExec"); + let control_legacy = + native_result_rows(native_tql_output(&instance, &control_sql, false).await).await; + assert_eq!( + control_native, control_legacy, + "the control agrees with the existing plan" + ); +} + +/// Enabling query fallback — by the engine option or by the per-request `query_fallback` hint — must +/// keep an otherwise fusible query on the existing plan without changing its answer. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_fallback_options_keep_expected_answer() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_fallback_control"; + // One eligible selector, so only the fallback settings decide the route. `sum_over_time` is + // exact over the two samples in the window: 10 + 30 for `a`, 10 + 70 for `b`. + let query = format!("sum_over_time({TABLE}[2m])"); + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + for (tag, allow_query_fallback) in [("engine_off", false), ("engine_on", true)] { + let (standalone, observer) = native_test_instance_with_options( + GreptimeDbStandaloneBuilder::new(&format!("native_promql_fallback_{tag}")), + QueryOptions { + experimental_native_promql: true, + allow_query_fallback, + ..Default::default() + }, + ) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + "CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=mito; \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30), \ + ('b', 60000, 10), ('b', 120000, 70);" + ), + QueryContext::arc(), + ) + .await; + let mut hinted = native_query_context(true, 4).as_ref().clone(); + hinted.set_extension(QUERY_FALLBACK_HINT, "true"); + for (route, ctx) in [ + ("plain", native_query_context(true, 4)), + ("fallback-hint", Arc::new(hinted)), + ] { + let sql = format!("TQL EVAL (120, 120, '1m') {query}"); + observer.clear(); + let rows = + native_result_rows(instance.do_query(&sql, ctx).await.remove(0).unwrap()).await; + assert_eq!( + rows, expected, + "{tag} {route}: the answer must not depend on the fallback settings" + ); + if allow_query_fallback || route == "fallback-hint" { + // A fallback rewrites every table scan into `MergeScan`, which the fused node + // cannot read, so a query allowed to fall back keeps the existing plan. + assert!( + observer.fused().is_none(), + "{tag} {route}: enabling the fallback must keep the existing plan" + ); + } else { + assert_eq!( + observer.take().name(), + "FusedSeriesExec", + "{tag} {route}: with no fallback enabled this query must be fused" + ); + } + } + } +} + +/// A store configured for the columnar series format must not change the answer of a query over a +/// field that format cannot serve: the columnar samples carry `Float64` series only, so an integer +/// field has no raw frame to fold and the route has to decline it while the existing plan answers. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_config_falls_back_for_integer_field() { + common_telemetry::init_default_ut_logging(); + const TABLE: &str = "native_columnar_int"; + // Both functions are exact over the two samples inside the window: 10 + 30 and their mean. + const SUM: f64 = 40.0; + const AVG: f64 = 20.0; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new("native_promql_columnar_fallback") + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + "CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val BIGINT, PRIMARY KEY (host)) ENGINE=mito; \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30);" + ), + QueryContext::arc(), + ) + .await; + for (query, expected) in [ + (format!("sum_over_time({TABLE}[2m])"), SUM), + (format!("avg_over_time({TABLE}[2m])"), AVG), + ] { + let sql = format!("TQL EVAL (120, 120, '1m') {query}"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + assert_eq!( + native_rows.len(), + 1, + "native-on `{query}`: one row for the one series" + ); + assert_eq!(native_rows[0].0, 120_000, "native-on `{query}`: ts"); + assert_eq!(native_rows[0].1, "a", "native-on `{query}`: label"); + assert!( + (native_rows[0].2 - expected).abs() <= 1e-9 * expected.abs().max(1.0), + "native-on `{query}`: expected {expected}, found {}", + native_rows[0].2 + ); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!( + native_rows, legacy_rows, + "`{query}`: same rows on both routes" + ); + assert!( + observer.fused().is_none(), + "`{query}`: the native route must decline a field the columnar samples cannot serve" + ); + } +} diff --git a/tests/cases/standalone/common/view/tql_eval.result b/tests/cases/standalone/common/view/tql_eval.result deleted file mode 100644 index 704ba763969..00000000000 --- a/tests/cases/standalone/common/view/tql_eval.result +++ /dev/null @@ -1,41 +0,0 @@ -CREATE TABLE tql_view_source( - ts TIMESTAMP TIME INDEX, - host STRING, - val DOUBLE, - PRIMARY KEY (host) -); - -Affected Rows: 0 - -INSERT INTO tql_view_source VALUES (0, 'a', 2); - -Affected Rows: 1 - -CREATE VIEW tql_view AS TQL EVAL (0, 0, '1s', '5m') tql_view_source; - -Affected Rows: 0 - -SELECT * FROM tql_view; - -+---------------------+------+-----+ -| ts | host | val | -+---------------------+------+-----+ -| 1970-01-01T00:00:00 | a | 2.0 | -+---------------------+------+-----+ - -SHOW CREATE VIEW tql_view; - -+----------+---------------------------------------------------------------------+ -| View | Create View | -+----------+---------------------------------------------------------------------+ -| tql_view | CREATE VIEW tql_view AS TQL EVAL (0, 0, '1s', '5m') tql_view_source | -+----------+---------------------------------------------------------------------+ - -DROP VIEW tql_view; - -Affected Rows: 0 - -DROP TABLE tql_view_source; - -Affected Rows: 0 - diff --git a/tests/cases/standalone/common/view/tql_eval.sql b/tests/cases/standalone/common/view/tql_eval.sql deleted file mode 100644 index 367eecb6bde..00000000000 --- a/tests/cases/standalone/common/view/tql_eval.sql +++ /dev/null @@ -1,18 +0,0 @@ -CREATE TABLE tql_view_source( - ts TIMESTAMP TIME INDEX, - host STRING, - val DOUBLE, - PRIMARY KEY (host) -); - -INSERT INTO tql_view_source VALUES (0, 'a', 2); - -CREATE VIEW tql_view AS TQL EVAL (0, 0, '1s', '5m') tql_view_source; - -SELECT * FROM tql_view; - -SHOW CREATE VIEW tql_view; - -DROP VIEW tql_view; - -DROP TABLE tql_view_source; From 51f754356b16cacc0a5d722ae1989b2c9d76be74 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 24 Oct 2026 19:05:28 +0800 Subject: [PATCH 34/38] test(query): iterate captured delta result batches Signed-off-by: discord9 --- tests-integration/src/tests/promql_test.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 7f6e900bc1d..f953078f183 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -1885,7 +1885,7 @@ async fn native_marker_rows(output: Output) -> Vec<(i64, String, String, f64)> { _ => unreachable!("TQL EVAL returns record batches"), }; let mut rows = Vec::new(); - for batch in &batches { + for batch in batches.iter() { let timestamps = batch .column_by_name("ts") .unwrap() @@ -2133,8 +2133,8 @@ async fn native_promql_view_fallback_preserves_native_hint() { execute_all(&instance, "CREATE TABLE native_view_source (ts TIMESTAMP TIME INDEX, host STRING, val DOUBLE, PRIMARY KEY(host)); INSERT INTO native_view_source VALUES (0, 'a', 2)", QueryContext::arc()).await; let ctx = native_query_context(true, 4); let expected_extensions = ctx.extensions(); - // The AST the parser would build for `CREATE VIEW native_view AS TQL EVAL (0, 0, '1s', '5m') - // native_view_source`; the statement's spelling is covered by the separate view case. + // Build the TQL view programmatically; the SQL syntax and SHOW CREATE round trip belong to the + // independent view-persistence PR. let create_view = SqlStatement::CreateView(CreateView { name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new("native_view"))]), columns: vec![], From 5ba2d76f5a9bc05168927380fe5dc85aaf05434a Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 24 Oct 2026 21:05:18 +0800 Subject: [PATCH 35/38] test(query): cover fused Columnar Float64 execution Signed-off-by: discord9 --- tests-integration/src/tests/promql_test.rs | 130 +++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index f953078f183..32539b31dcf 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -25,6 +25,7 @@ use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::{Output, OutputData}; use common_recordbatch::util::collect_batches; use datafusion::physical_plan::ExecutionPlanProperties; +use datafusion::physical_plan::coop::CooperativeExec; use datatypes::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, }; @@ -44,6 +45,9 @@ use sql::ast::{Ident, ObjectName, ObjectNamePart}; use sql::statements::create::CreateView; use sql::statements::statement::Statement as SqlStatement; use sql::statements::tql::{Tql, TqlEval}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::region_engine::RawLabelSource; +use table::table::scan::SeriesScanExec; use super::test_util::{ both_instances_cases, check_unordered_output_stream, distributed, standalone, @@ -2456,3 +2460,129 @@ async fn native_promql_columnar_config_falls_back_for_integer_field() { ); } } + +/// The positive counterpart of the fallback above: a sparse physical metric table whose series +/// index is published in the columnar format fuses a `Float64` selector and matches the old plan. +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_fused_series_matches_legacy() { + common_telemetry::init_default_ut_logging(); + const PHYSICAL: &str = "native_columnar_fused_phy"; + const TABLE: &str = "native_columnar_fused"; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + // Slow the periodic pass, so the manual build below is what publishes the catalog. + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let builder = GreptimeDbStandaloneBuilder::new("native_promql_columnar_fused") + .with_mito_config(mito_config); + let (standalone, observer) = native_test_instance(builder).await; + let instance = standalone.fe_instance().clone(); + // One flushed SST is enough for a manual build, which skips the periodic four-file trigger. + execute_all( + &instance, + &format!( + "CREATE TABLE {PHYSICAL} (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric \ + WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); \ + CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) \ + ENGINE=metric WITH ('on_physical_table' = '{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30), \ + ('b', 60000, 20), ('b', 120000, 60); \ + ADMIN FLUSH_TABLE('{PHYSICAL}'); \ + ADMIN BUILD_SERIES_INDEX('{PHYSICAL}');" + ), + QueryContext::arc(), + ) + .await; + // Both samples of each series fall in the 2m window: 10 + 30 and 20 + 60. + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + let sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}[2m])"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + // The answer is drained first, so the counters below describe execution, not planning. + let plan = observer.take(); + assert_eq!(native_rows, expected, "the fused route reports the sums"); + // `EnsureCooperative` wraps a leaf that does not ask for cooperative scheduling, and the fused + // node reads the same series scan through it. + let child = plan.children()[0]; + let child = child + .downcast_ref::() + .map_or(child, |cooperative| cooperative.input()); + let series_scan = child + .downcast_ref::() + .expect("the fused node reads a prepared series scan"); + let RawLabelSource::PrimaryKey { + metric_name, + physical_metadata, + logical_tags, + } = series_scan.label_source() + else { + panic!("the fused scan must materialize labels from the primary key"); + }; + assert_eq!(metric_name.as_ref(), TABLE); + assert_eq!( + logical_tags + .iter() + .map(|tag| tag.name.as_ref()) + .collect::>(), + vec!["host"], + "the logical tags are the primary-key columns of the scanned region" + ); + for tag in logical_tags.iter() { + let column = physical_metadata.column_by_id(tag.column_id).unwrap(); + assert_eq!(column.column_schema.name, "host"); + assert_eq!(column.semantic_type, SemanticType::Tag); + } + assert_eq!( + physical_metadata.primary_key_encoding, + PrimaryKeyEncoding::Sparse + ); + let series_schema = series_scan.series_schema(); + assert_eq!(series_schema.fields.len(), 1); + assert_eq!(series_schema.fields[0].name.as_ref(), "val"); + // Only the columnar writer publishes these formats, so the scanned region's catalog rules out + // a configuration-only run. + let samples = std::path::Path::new(&standalone.opts.storage.data_home) + .join("series_index") + .join(physical_metadata.region_id.as_u64().to_string()) + .join("samples.json"); + let catalog: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(samples).unwrap()).unwrap(); + let entries = catalog["entries"].as_array().unwrap(); + let series = catalog["series"].as_array().unwrap(); + assert!(!entries.is_empty() && !series.is_empty(), "published"); + assert!( + entries + .iter() + .all(|entry| entry["artifact"]["format"] == "columnar-samples") + ); + assert!( + series + .iter() + .all(|entry| entry["artifact"]["format"] == "columnar-series") + ); + // One value per series: the node ran, rather than being merely planned. + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 2, "both series must report one value"); + // Native off: the same rows must come from the existing plan. + observer.clear(); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(native_rows, legacy_rows, "same rows on both routes"); + assert!(observer.fused().is_none(), "the legacy route must not fuse"); +} From 3b0587aa8d5e4475485d6535ca2d7163fb890875 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 24 Oct 2026 23:01:49 +0800 Subject: [PATCH 36/38] fix(mito2): read raw series when columnar indexes are not ready Keep native queries available for unflushed memtables and SSTs with incomplete columnar series or samples coverage by selecting the ordinary raw-series reader. Preserve structural validation and propagate artifact read errors instead of retrying them through legacy execution. Cover all four incomplete-readiness states in public PromQL queries, both storage scan modes and ordering, complete prepared-source execution, and artifact error propagation. Signed-off-by: discord9 --- src/mito2/src/read/series_scan.rs | 91 ++++-- src/mito2/src/read/series_scan/prepared.rs | 47 +-- .../src/read/series_scan/prepared/tests.rs | 280 ++++++++++++++---- tests-integration/src/tests/promql_test.rs | 173 +++++++++++ 4 files changed, 482 insertions(+), 109 deletions(-) diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 433fac8c7f1..8fb06782240 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -256,10 +256,41 @@ impl SeriesScan { .is_some_and(|context| context.format == crate::config::SeriesIndexFormat::Columnar) } + /// Returns whether raw series come from the columnar prepared source. + /// + /// Derived only from the pinned series index version and every configured partition range, so + /// preflight, the ordering guarantee, the partition reader and scanner-wide preparation agree. + /// `false` means the stored sources are read by the ordinary scan. + fn uses_prepared_columnar(&self) -> Result { + if !self.uses_columnar() { + return Ok(false); + } + let ranges = self + .properties + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new) + } + + /// Returns the raw ordering guarantee for a scan whose columnar readiness is already known. + fn raw_ordering(&self, prepared_columnar: bool) -> SeriesOrdering { + if prepared_columnar { + return SeriesOrdering::KeyThenTime; + } + match self.mode { + SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime, + SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous, + } + } + + /// Prepares the raw scan and reports whether its data comes from the columnar prepared source. fn prepare_raw_scan( &self, label_source: &RawLabelSource, - ) -> Result, BoxedError> { + ) -> Result, BoxedError> { let input = &self.stream_ctx.input; let metadata = input.mapper.metadata(); match label_source { @@ -386,25 +417,19 @@ impl SeriesScan { ); return Ok(None); } - if self.uses_columnar() { - let ranges = self - .properties - .partitions - .iter() - .flatten() - .copied() - .collect::>(); - prepared::validate_inputs(&self.stream_ctx, &ranges).map_err(BoxedError::new)?; - } + // Structural eligibility errors still reject the raw path; ordinary scans cover sources + // without columnar artifacts, so missing coverage is not an error here. + let prepared_columnar = self.uses_prepared_columnar()?; let schema = Arc::new(SeriesSchema { fields: specs.into(), eval_grid: None, }); let typed_schema = input.mapper.input_arrow_schema(false); let mut predicates = Vec::new(); - let predicate = if self.uses_columnar() { + let predicate = if prepared_columnar { input.predicate_group().predicate() } else { + // The ordinary scan applies region predicates before this residual filtering. input.predicate_group().predicate_without_region() }; if let Some(predicate) = predicate { @@ -452,11 +477,14 @@ impl SeriesScan { predicates.push(physical); } } - Ok(Some(RawScanPreparation { - schema, - fields, - predicates, - })) + Ok(Some(( + RawScanPreparation { + schema, + fields, + predicates, + }, + prepared_columnar, + ))) } fn supports_two_phase(input: &ScanInput) -> bool { @@ -844,12 +872,11 @@ impl RegionScanner for SeriesScan { } fn raw_series_ordering(&self) -> SeriesOrdering { - if self.uses_columnar() { - return SeriesOrdering::KeyThenTime; - } - match self.mode { - SeriesScanMode::Legacy => SeriesOrdering::KeyThenTime, - SeriesScanMode::TwoPhase => SeriesOrdering::Contiguous, + // Readiness is fallible here, but every fallible entry point rejects the same structural + // errors before execution, so the `false` fallback below is never observed in a query. + match self.uses_prepared_columnar() { + Ok(prepared_columnar) => self.raw_ordering(prepared_columnar), + Err(_) => self.raw_ordering(false), } } @@ -859,7 +886,7 @@ impl RegionScanner for SeriesScan { ) -> Result>, BoxedError> { Ok(self .prepare_raw_scan(label_source)? - .map(|preparation| preparation.schema)) + .map(|(preparation, _)| preparation.schema)) } fn scan_series_partition( @@ -869,10 +896,10 @@ impl RegionScanner for SeriesScan { partition: usize, label_source: &RawLabelSource, ) -> Result, BoxedError> { - let Some(preparation) = self.prepare_raw_scan(label_source)? else { + let Some((preparation, prepared_columnar)) = self.prepare_raw_scan(label_source)? else { return Ok(None); }; - if self.uses_columnar() { + if prepared_columnar { let partitions = self.properties.num_partitions().max(1); if partition >= partitions { return Err(BoxedError::new( @@ -897,10 +924,12 @@ impl RegionScanner for SeriesScan { return Ok(Some(raw::from_stream( schema, labels, - SeriesOrdering::KeyThenTime, + self.raw_ordering(prepared_columnar), Box::pin(stream), ))); } + // Stored sources without columnar artifacts (memtables, unindexed SSTs) are read by the + // ordinary scan, so the raw path stays available. let metrics = new_partition_metrics( &self.stream_ctx, ctx.explain_verbose, @@ -911,7 +940,7 @@ impl RegionScanner for SeriesScan { let batches = self .scan_batch_in_partition(ctx, partition, metrics.clone(), metrics_set) .map_err(BoxedError::new)?; - let ordering = self.raw_series_ordering(); + let ordering = self.raw_ordering(prepared_columnar); Ok(Some(raw::make_stream( batches, preparation.schema, @@ -957,9 +986,13 @@ impl RegionScanner for SeriesScan { &self.metrics_list, ); Box::pin(async move { - let Some(preparation) = preparation? else { + let Some((preparation, prepared_columnar)) = preparation? else { return Ok(None); }; + if !prepared_columnar { + // Known unsupported coverage: the caller keeps the ordinary scan instead. + return Ok(None); + } state .prepare(stream_ctx, ranges, labels, preparation, metrics) .await diff --git a/src/mito2/src/read/series_scan/prepared.rs b/src/mito2/src/read/series_scan/prepared.rs index 8dc002f9b23..0c87d6f18d4 100644 --- a/src/mito2/src/read/series_scan/prepared.rs +++ b/src/mito2/src/read/series_scan/prepared.rs @@ -137,20 +137,30 @@ fn same_labels(left: &RawLabelSource, right: &RawLabelSource) -> bool { } /// Validates the pinned assignments and coverage without opening artifacts. -pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result<()> { - collect_inputs(context, ranges).map(|_| ()) +/// +/// Returns whether every assigned source can be read from columnar artifacts. `false` is not an +/// error: the caller scans those sources with the ordinary reader instead. +pub(crate) fn validate_inputs(context: &StreamContext, ranges: &[PartitionRange]) -> Result { + collect_inputs(context, ranges).map(|inputs| inputs.is_some()) } +/// Collects the assigned SSTs and their covering columnar series indexes. +/// +/// `None` means the assignments are structurally valid but not ready for columnar reads: an +/// assigned memtable, an SST without columnar series coverage, or an SST without columnar +/// samples. Structural corruption is still an error, and every assignment is validated before +/// either answer is returned, so a later invalid assignment is never masked by missing artifacts. fn collect_inputs( context: &StreamContext, ranges: &[PartitionRange], -) -> Result<(Vec, HashSet)> { +) -> Result, HashSet)>> { let input = &context.input; let metadata = input.region_metadata(); let index = input .series_index .as_ref() .ok_or_else(|| invalid("missing columnar index context"))?; + let mut ready = true; let mut files = BTreeMap::>::new(); for range in ranges { let range = context @@ -159,7 +169,9 @@ fn collect_inputs( .ok_or_else(|| invalid("invalid prepared scan range"))?; for row_group in &range.row_group_indices { if context.is_mem_range_index(*row_group) { - return Err(invalid("columnar prepared reads require flushed memtables")); + // Artifacts exist for flushed SSTs only. + ready = false; + continue; } if !context.is_file_range_index(*row_group) { return Err(invalid( @@ -203,24 +215,18 @@ fn collect_inputs( } } if !covered { - return Err(invalid(format!( - "SST {} has no columnar series coverage; run build_series_index()", - file.file_id() - ))); + ready = false; } - index + if index .version .samples .get(&file.file_id().file_id()) - .filter(|handle| handle.entry().artifact.format == ArtifactFormat::ColumnarSamples) - .ok_or_else(|| { - invalid(format!( - "SST {} has no columnar samples; run build_series_index()", - file.file_id() - )) - })?; + .is_none_or(|handle| handle.entry().artifact.format != ArtifactFormat::ColumnarSamples) + { + ready = false; + } } - Ok((file_indices, indexes)) + Ok(ready.then_some((file_indices, indexes))) } struct Candidate { @@ -254,7 +260,12 @@ impl ColumnarPreparedSource { .series_index .as_ref() .ok_or_else(|| invalid("missing columnar index context"))?; - let (files, indexes) = collect_inputs(&context, &ranges)?; + let Some((files, indexes)) = collect_inputs(&context, &ranges)? else { + // Readiness was checked before preparation, so this is a real capability change. + return Err(invalid( + "columnar series coverage changed after series preparation", + )); + }; let mut samples = Vec::>::new(); for file_idx in files { let file = &input.files[file_idx]; diff --git a/src/mito2/src/read/series_scan/prepared/tests.rs b/src/mito2/src/read/series_scan/prepared/tests.rs index a0fbee4c14c..2d47eb36676 100644 --- a/src/mito2/src/read/series_scan/prepared/tests.rs +++ b/src/mito2/src/read/series_scan/prepared/tests.rs @@ -25,13 +25,14 @@ use datafusion::execution::TaskContext; use datafusion::physical_plan::{ExecutionPlan, metrics::ExecutionPlanMetricsSet}; use datafusion_expr::{col, lit}; use datatypes::prelude::ConcreteDataType; +use datatypes::series::SeriesOrdering; use datatypes::series::raw::RawSeriesChunk; use datatypes::series::samples::SampleRef; use object_store::ObjectStore; use object_store::layers::mock::{self, MockLayerBuilder, oio}; use store_api::codec::PrimaryKeyEncoding; use store_api::region_engine::{ - BoundTag, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, + BoundTag, PartitionRange, PrepareRequest, QueryScanContext, RegionEngine, RegionScanner, }; use store_api::region_request::{ RegionBuildIndexRequest, RegionDeleteRequest, RegionPutRequest, RegionRequest, @@ -51,6 +52,10 @@ use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_wit use crate::test_util::{TestEnv, flush_region}; async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) { + fixture_with(false).await +} + +async fn fixture_with(experimental_series_scan_v2: bool) -> (TestEnv, MitoEngine, MitoRegionRef) { let mut env = TestEnv::with_prefix("prepared-columnar").await; let mut metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); metadata @@ -66,6 +71,7 @@ async fn fixture() -> (TestEnv, MitoEngine, MitoRegionRef) { MitoConfig { experimental_enable_series_index: true, experimental_series_index_format: SelectedFormat::Columnar, + experimental_series_scan_v2, experimental_series_index_maintenance_interval: Duration::from_secs(3600), min_compaction_interval: Duration::from_secs(3600), ..Default::default() @@ -192,6 +198,41 @@ impl mock::Read for ReadRecorder { } } +/// Fails reads of the selected artifact paths so tests can prove that object-store failures +/// surface as errors instead of being classified as not-ready fallbacks. +struct FailingReader { + inner: oio::Reader, + fail: bool, +} +impl mock::Read for FailingReader { + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + if self.fail { + return Err(injected_failure()); + } + self.inner.read(range).await + } + + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> { + if self.fail { + return Err(injected_failure()); + } + self.inner.open(range).await + } +} + +fn injected_failure() -> mock::Error { + mock::Error::new( + mock::ErrorKind::Unexpected, + "injected series manifest read failure", + ) +} + fn scanner(region: &MitoRegionRef, store: ObjectStore) -> SeriesScan { scanner_with_format(region, store, SelectedFormat::Columnar) } @@ -269,6 +310,35 @@ fn rows(chunks: &[RawSeriesChunk]) -> Vec<(Vec, i64, f64)> { .collect() } +fn samples(chunks: &[RawSeriesChunk]) -> Vec<(i64, f64)> { + rows(chunks) + .into_iter() + .map(|(_, time, value)| (time, value)) + .collect() +} + +/// Reads the raw path of every partition, checks each stream declares `ordering`, and returns the +/// samples ordered by time so a check does not depend on how range grouping spread a memtable. +async fn all_raw_samples( + scanner: &dyn RegionScanner, + ctx: &QueryScanContext, + metrics: &ExecutionPlanMetricsSet, + binding: &RawLabelSource, + ordering: SeriesOrdering, +) -> Vec<(i64, f64)> { + let mut collected = Vec::new(); + for partition in 0..scanner.properties().num_partitions().max(1) { + let stream = scanner + .scan_series_partition(ctx, metrics, partition, binding) + .unwrap() + .unwrap(); + assert_eq!(ordering, stream.ordering()); + collected.extend(samples(&stream.try_collect::>().await.unwrap())); + } + collected.sort_unstable_by_key(|(time, _)| *time); + collected +} + #[tokio::test] async fn prepared_sources_share_discovery_and_read_independently() { let (_env, engine, region) = fixture().await; @@ -386,89 +456,133 @@ async fn prepared_sources_share_discovery_and_read_independently() { engine.stop().await.unwrap(); } +#[rstest::rstest] +#[case::legacy(false)] +#[case::two_phase(true)] #[tokio::test] -async fn native_scan_reports_incomplete_coverage_and_uses_prepared_source() { - let (_env, engine, region) = fixture().await; +async fn native_scan_falls_back_to_raw_sources_until_columnar_coverage(#[case] v2: bool) { + let (_env, engine, region) = fixture_with(v2).await; let request = ScanRequest { distribution: Some(TimeSeriesDistribution::PerSeries), ..Default::default() }; - // Three files are below the automatic index threshold. + let ctx = QueryScanContext::default(); + let metrics = ExecutionPlanMetricsSet::new(); + // The three flushed SSTs only have columnar artifacts after `build`. + let fallback_ordering = if v2 { + SeriesOrdering::Contiguous + } else { + SeriesOrdering::KeyThenTime + }; let raw_scanner = engine .handle_query(region.region_id, request.clone()) .await .unwrap(); let binding = labels(raw_scanner.metadata()); - assert!(raw_scanner.raw_series_schema(&binding).is_err()); - assert!( - raw_scanner - .scan_series_partition( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - 0, - &binding, - ) - .is_err() + assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering()); + assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some()); + assert_eq!( + vec![(1000, 1.0), (2000, 1.0)], + all_raw_samples( + raw_scanner.as_ref(), + &ctx, + &metrics, + &binding, + fallback_ordering + ) + .await ); + // Known unready coverage selects the ordinary scan instead of a prepared source. assert!( raw_scanner - .prepare_series( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - &binding - ) + .prepare_series(&ctx, &metrics, &binding) .await - .is_err() + .unwrap() + .is_none() ); + build(&engine, ®ion).await; - let raw_scanner = engine + let mut raw_scanner = engine .handle_query(region.region_id, request.clone()) .await .unwrap(); + // Move every assigned range into a new last partition: readiness, the ordering guarantee and + // the reader must all follow the new assignment, not the pre-prepare partition list. + let assigned = raw_scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + raw_scanner + .prepare(PrepareRequest::default().with_ranges(vec![vec![], assigned])) + .unwrap(); + assert_eq!(2, raw_scanner.properties().num_partitions()); let binding = labels(raw_scanner.metadata()); + assert_eq!( + SeriesOrdering::KeyThenTime, + raw_scanner.raw_series_ordering() + ); + let source = raw_scanner + .prepare_series(&ctx, &metrics, &binding) + .await + .unwrap() + .unwrap(); + assert_eq!(1, source.num_series()); + // The first partition owns no range and the second owns all of them. + assert_eq!( + vec![(1000, 1.0), (2000, 1.0)], + all_raw_samples( + raw_scanner.as_ref(), + &ctx, + &metrics, + &binding, + SeriesOrdering::KeyThenTime + ) + .await + ); let plan = Arc::new(RegionScanExec::new(raw_scanner, request.clone(), None).unwrap()) .try_to_series_scan(binding.clone()) .unwrap() .unwrap(); - let frame = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); - let chunks = unpack_raw_series(frame, Some(plan.series_schema().clone())) - .unwrap() - .try_collect::>() - .await - .unwrap(); + assert_eq!(SeriesOrdering::KeyThenTime, plan.raw_series_ordering()); + let frame = plan.execute(1, Arc::new(TaskContext::default())).unwrap(); + let unpacked = unpack_raw_series(frame, Some(plan.series_schema().clone())).unwrap(); + // The packed frame declares the same guarantee the plan advertises. + assert_eq!(plan.raw_series_ordering(), unpacked.ordering()); assert_eq!( - vec![1000, 2000], - rows(&chunks).iter().map(|(_, t, _)| *t).collect::>() + vec![(1000, 1.0), (2000, 1.0)], + samples(&unpacked.try_collect::>().await.unwrap()) ); + + // An assigned memtable is known unready, so the raw path keeps the ordinary scan and must + // still return the new row. write(&engine, ®ion, 0, 3000, Some(4.0)).await; let raw_scanner = engine .handle_query(region.region_id, request) .await .unwrap(); let binding = labels(raw_scanner.metadata()); - assert!(raw_scanner.raw_series_schema(&binding).is_err()); + assert_eq!(fallback_ordering, raw_scanner.raw_series_ordering()); + assert!(raw_scanner.raw_series_schema(&binding).unwrap().is_some()); assert!( raw_scanner - .scan_series_partition( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - 0, - &binding, - ) - .is_err() - ); - assert!( - raw_scanner - .prepare_series( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - &binding - ) + .prepare_series(&ctx, &metrics, &binding) .await - .err() .unwrap() - .output_msg() - .contains("memtables") + .is_none() + ); + assert_eq!( + vec![(1000, 1.0), (2000, 1.0), (3000, 4.0)], + all_raw_samples( + raw_scanner.as_ref(), + &ctx, + &metrics, + &binding, + fallback_ordering + ) + .await ); engine.stop().await.unwrap(); } @@ -1046,7 +1160,8 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() { }; let mut context = make_context(); let ranges = context.partition_ranges(); - validate_inputs(&context, &ranges).unwrap(); + // Complete columnar coverage is the only ready answer. + assert!(validate_inputs(&context, &ranges).unwrap()); context.ranges[ranges[0].identifier].row_group_indices = vec![RowGroupIndex { index: context.input.num_files(), row_group_index: -1, @@ -1083,30 +1198,30 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() { .contains("foreign SST") ); + // Missing columnar samples are a not-ready answer, not an error. let mut context = make_context(); let index = context.input.series_index.as_mut().unwrap(); index.version = Arc::new(SeriesIndexVersion { series_indexes: index.version.series_indexes.clone(), ..Default::default() }); + assert!(!validate_inputs(&context, &ranges).unwrap()); + // Invalid assignments remain errors even when sample coverage is incomplete. + let mut invalid = ranges.clone(); + invalid.push(PartitionRange { + identifier: ranges[0].identifier + ranges.len(), + ..ranges[0] + }); assert!( - validate_inputs(&context, &ranges) + validate_inputs(&context, &invalid) .unwrap_err() .to_string() - .contains("no columnar samples") + .contains("invalid prepared scan range") ); let scan = SeriesScan::new(context.input, false); let binding = labels(scan.metadata()); - assert!(scan.raw_series_schema(&binding).is_err()); - assert!( - scan.scan_series_partition( - &QueryScanContext::default(), - &ExecutionPlanMetricsSet::new(), - 0, - &binding - ) - .is_err() - ); + // The raw path stays available for the ordinary scan and preparation is skipped. + assert!(scan.raw_series_schema(&binding).unwrap().is_some()); assert!( scan.prepare_series( &QueryScanContext::default(), @@ -1114,8 +1229,49 @@ async fn preflight_rejects_unsupported_assignments_without_artifact_io() { &binding ) .await - .is_err() + .unwrap() + .is_none() ); assert!(reads.lock().unwrap().is_empty()); engine.stop().await.unwrap(); } + +#[tokio::test] +async fn series_artifact_read_failure_propagates() { + let (_env, engine, region) = fixture().await; + build(&engine, ®ion).await; + // The published catalog is complete; the injected manifest read failure must still propagate. + let store = region.series_index_store.clone().unwrap().layer( + MockLayerBuilder::default() + .reader_factory(Arc::new(|path, _, inner| { + Box::new(FailingReader { + inner, + fail: path.contains("/series/") && path.ends_with("/manifest.json"), + }) + })) + .build() + .unwrap(), + ); + let context = StreamContext::seq_scan_ctx(scan_input( + ®ion, + store, + SelectedFormat::Columnar, + (0..region.version().metadata.column_metadatas.len()).collect(), + )); + let ranges = context.partition_ranges(); + assert!(validate_inputs(&context, &ranges).unwrap()); + let scan = SeriesScan::new(context.input, false); + let binding = labels(scan.metadata()); + assert!(scan.raw_series_schema(&binding).unwrap().is_some()); + let error = scan + .prepare_series( + &QueryScanContext::default(), + &ExecutionPlanMetricsSet::new(), + &binding, + ) + .await + .err() + .expect("a failed artifact read must not report a not-ready source"); + assert!(error.output_msg().contains("injected"), "{error}"); + engine.stop().await.unwrap(); +} diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 32539b31dcf..8e4aea3a85a 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -2586,3 +2586,176 @@ async fn native_promql_columnar_fused_series_matches_legacy() { assert_eq!(native_rows, legacy_rows, "same rows on both routes"); assert!(observer.fused().is_none(), "the legacy route must not fuse"); } + +/// No published series index: the columnar `Float64` selector must answer from the raw memtable +/// and from an unindexed SST instead of failing to plan. +#[rstest] +#[case::memtable("memtable", false)] +#[case::unindexed_sst("unindexed_sst", true)] +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_unindexed_float64_matches_legacy( + #[case] stage: &str, + #[case] flush: bool, +) { + common_telemetry::init_default_ut_logging(); + const PHYSICAL: &str = "native_columnar_unindexed_phy"; + const TABLE: &str = "native_columnar_unindexed"; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + // No periodic build, so neither case ever gets an index. + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new(&format!("native_promql_columnar_{stage}")) + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + execute_all( + &instance, + &format!( + "CREATE TABLE {PHYSICAL} (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric \ + WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); \ + CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) \ + ENGINE=metric WITH ('on_physical_table' = '{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('a', 120000, 30), \ + ('b', 60000, 20), ('b', 120000, 60);" + ), + QueryContext::arc(), + ) + .await; + if flush { + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{PHYSICAL}');"), + QueryContext::arc(), + ) + .await; + } + // Both samples of each series fall in the 2m window: 10 + 30 and 20 + 60. + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + let sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}[2m])"); + // Legacy off first, so a red native query is about the fused route. + observer.clear(); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(legacy_rows, expected, "{stage}: existing plan must answer"); + assert!(observer.fused().is_none(), "{stage}: legacy must not fuse"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + // Drained first, so the counter below describes execution, not planning; `take` requires the + // fused node to have run. + let plan = observer.take(); + assert_eq!(native_rows, expected, "{stage}: fused route must answer"); + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 2, "{stage}: two values, one per series"); + assert_eq!( + native_rows, legacy_rows, + "{stage}: same rows on both routes" + ); +} + +/// An index over the first SST only: the newer rows still in the memtable, then flushed into a +/// second unindexed SST, must answer the same sums. +#[rstest] +#[case::partial_memtable("partial_memtable", false)] +#[case::partial_sst("partial_sst", true)] +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_columnar_partial_index_float64_matches_legacy( + #[case] stage: &str, + #[case] flush: bool, +) { + common_telemetry::init_default_ut_logging(); + const PHYSICAL: &str = "native_columnar_partial_phy"; + const TABLE: &str = "native_columnar_partial"; + let mito_config = mito2::config::MitoConfig { + experimental_enable_series_index: true, + experimental_series_index_format: mito2::config::SeriesIndexFormat::Columnar, + // Slow the periodic pass, so only the manual build below publishes an index. + experimental_series_index_maintenance_interval: Duration::from_secs(3600), + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new(&format!("native_promql_columnar_{stage}")) + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + // The first samples are flushed and indexed before the second ones are inserted, so the index + // covers one SST while the later rows start out unflushed. + execute_all( + &instance, + &format!( + "CREATE TABLE {PHYSICAL} (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric \ + WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); \ + CREATE TABLE {TABLE} (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) \ + ENGINE=metric WITH ('on_physical_table' = '{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 60000, 10), ('b', 60000, 20); \ + ADMIN FLUSH_TABLE('{PHYSICAL}'); \ + ADMIN BUILD_SERIES_INDEX('{PHYSICAL}'); \ + INSERT INTO {TABLE} (host, ts, val) VALUES ('a', 120000, 30), ('b', 120000, 60);" + ), + QueryContext::arc(), + ) + .await; + if flush { + // The new SST stays unindexed: the build above is not repeated. + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{PHYSICAL}');"), + QueryContext::arc(), + ) + .await; + } + // Both samples of each series fall in the 2m window: 10 + 30 and 20 + 60. + let expected = vec![ + (120_000, "a".to_string(), 40.0), + (120_000, "b".to_string(), 80.0), + ]; + let sql = format!("TQL EVAL (120, 120, '1m') sum_over_time({TABLE}[2m])"); + // Legacy off first, so a red native query is about the fused route. + observer.clear(); + let legacy_rows = native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!(legacy_rows, expected, "{stage}: existing plan must answer"); + assert!(observer.fused().is_none(), "{stage}: legacy must not fuse"); + observer.clear(); + let native_rows = native_result_rows(native_tql_output(&instance, &sql, true).await).await; + // Drained first, so the counter below describes execution, not planning; `take` requires the + // fused node to have run. + let plan = observer.take(); + assert_eq!(native_rows, expected, "{stage}: fused route must answer"); + let output_values = plan + .metrics() + .expect("fused plan metrics") + .iter() + .filter_map(|metric| match metric.value() { + datafusion::physical_plan::metrics::MetricValue::Count { name, count } + if name.as_ref() == "native.root.output_values" => + { + Some(count.value()) + } + _ => None, + }) + .sum::(); + assert_eq!(output_values, 2, "{stage}: two values, one per series"); + assert_eq!( + native_rows, legacy_rows, + "{stage}: same rows on both routes" + ); +} From aa4a381f7113128ff48c33184b8c2530dc6d97d2 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sun, 25 Oct 2026 03:33:49 +0800 Subject: [PATCH 37/38] refactor(promql): keep fused POC execution partition-local Remove plan-global label-set occupancy and the speculative split-series tree guard. Keep the native evaluator, finality, memory budgets, errors and ordinary RecordBatch output local to each scan partition. Narrow native admission to selectors and the small unary POC subset. Document that cross-partition final-label conflicts and semantic merging are outside this POC, and cover real rate/abs(rate) execution across dense and sparse storage spanning SST and memtable data. Signed-off-by: discord9 --- src/promql/AGENTS.md | 35 +- src/promql/src/extension_plan/fused_series.rs | 394 +---- .../src/extension_plan/fused_series/tests.rs | 1268 +---------------- src/promql/src/native/runs.rs | 9 +- src/query/src/promql/planner/native.rs | 82 +- src/query/src/promql/planner/test.rs | 28 +- tests-integration/src/tests/promql_test.rs | 119 ++ 7 files changed, 262 insertions(+), 1673 deletions(-) diff --git a/src/promql/AGENTS.md b/src/promql/AGENTS.md index 2381461d3c0..53ad3c61169 100644 --- a/src/promql/AGENTS.md +++ b/src/promql/AGENTS.md @@ -25,19 +25,21 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`. `EmptyMetric` and `UnionDistinctOn`. `PromExtensionPlanner` (`src/promql/src/extension_plan/planner.rs`) maps the logical nodes to their execution plans. -- **Fused native-series route:** the query planner offers a PromQL statement to - `native_compile::compile` only when the engine's - `in_process_region_transport` and `QueryOptions::experimental_native_promql` - options are enabled and the per-request `query.experimental_native_promql` - extension is true. Accepted plans use process-local `FusedSeries`; declined - or unsupported shapes fall back to the existing planner. The supported subset - requires one selector, a direct scan with a bound metric name, and no residual - label matchers or `__field__` matcher that narrows the scan, plus an eligible - raw-series table/field set and supported expression/grid. Metric scans retain - the static `__table_id` predicate. Hybrid CTEs keep the existing plan. The - operator also disables native planning for a programmatic TQL view so its - persisted Substrait plan remains serializable; `CREATE VIEW ... AS TQL EVAL` - syntax is a separate change. +- **Experimental fused native-series POC (default off):** the query planner offers + a PromQL statement only when the engine's `in_process_region_transport` and + `QueryOptions::experimental_native_promql` options are enabled and the per-request + `query.experimental_native_promql` extension is true. Admission is limited to a + bare vector selector (parentheses allowed) and unary `rate`, `increase`, + `sum_over_time`, and `abs` calls over a selector; range selectors occur only + inside those calls. It is a standalone route: unsupported expressions and all + other cases use the existing planner. Accepted plans use process-local + `FusedSeries`, with ordinary `RecordBatch` output collected through DataFusion. + Scan N→Fused N retains region/partition parallelism, but this POC does not merge + one logical series across regions or scan partitions, guarantee canonical final + labels across them, or use global label-set occupancy; `NativeExecutor` retains + its partition-local `Emitted` state. Repartitioning redistributes rows but does not perform + semantic series merging. A future gather/repartition plus merge/finalize is + outside this POC. Metric scans retain the static `__table_id` predicate. - **Admission and SQL output:** the guards live in `src/query/src/promql/planner/native.rs`. Non-`Float64` fields, alternative (`or`) matchers and `rate`/`increase` on temporality-tagged tables keep the @@ -58,9 +60,10 @@ operators. Statement-to-plan translation lives in `src/query/src/promql/`. (`src/query/src/dist_plan/commutativity.rs`). Process-local `FusedSeries` is the exception: it is not serialized to a datanode. Its distributed placement is restricted to a direct table scan, a bound metric name and trees that read - one series at a time, and it is only correct because native admission refuses a - table whose partition columns are not all tag columns: that is what keeps a - series inside one region. Keep that admission and this placement together. + one series at a time. Native admission requires partition columns to be tags, + but region/partition-local execution does not merge a logical series across + regions or scan partitions; do not treat that placement guard as a global + series-merge or final-label guarantee. - A new logical node also needs its physical planner reachable from `DfQueryPlanner::create_physical_plan`, assembled in `src/query/src/query_engine/state.rs`. diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index d957851902a..fae78c06390 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -37,18 +37,15 @@ //! 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::collections::VecDeque; use std::fmt::{self, Formatter}; use std::hash::{Hash, Hasher}; use std::pin::Pin; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; 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::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::buffer::ScalarBuffer; use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; use datafusion::arrow::record_batch::RecordBatch; @@ -56,9 +53,7 @@ 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::execution::memory_pool::UnboundedMemoryPool; use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; @@ -69,17 +64,15 @@ use datafusion::physical_plan::{ }; use datatypes::series::SeriesOrdering; use datatypes::series::raw::RawSeriesChunk; -use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, Series, SeriesSchema}; +use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, 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)] @@ -456,89 +449,6 @@ fn tree_input(tree: &NativeTree) -> DataFusionResult { 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 @@ -762,11 +672,6 @@ pub struct FusedSeriesExec { 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. @@ -807,8 +712,6 @@ impl FusedSeriesExec { )?; 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 { @@ -816,7 +719,6 @@ impl FusedSeriesExec { "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)?; @@ -857,10 +759,6 @@ impl FusedSeriesExec { 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, @@ -873,7 +771,6 @@ impl FusedSeriesExec { tags: tags.into(), input, output_schema, - labels, metric: ExecutionPlanMetricsSet::new(), properties, }) @@ -967,14 +864,6 @@ impl ExecutionPlan for FusedSeriesExec { 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 @@ -990,8 +879,6 @@ impl ExecutionPlan for FusedSeriesExec { &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)), @@ -999,7 +886,6 @@ impl ExecutionPlan for FusedSeriesExec { 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, })) @@ -1100,263 +986,6 @@ fn field_column(samples: &Samples) -> DataFusionResult { } } -/// 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. @@ -1375,10 +1004,6 @@ pub struct FusedSeriesStream { 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. @@ -1447,9 +1072,6 @@ impl FusedSeriesStream { 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(); } @@ -1457,12 +1079,6 @@ impl FusedSeriesStream { /// 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, diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index c3837187474..2d98ad0fb38 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -47,8 +47,7 @@ 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::execution::memory_pool::UnboundedMemoryPool; use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_optimizer::PhysicalOptimizerRule; @@ -68,17 +67,13 @@ 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 super::{FusedSeries, FusedSeriesBinding, OutputColumn, POLL_BUDGET, 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, + AggregateKind, AggregateStage, BoundAggregate, BoundInput, 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. @@ -111,8 +106,6 @@ 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 { @@ -129,7 +122,7 @@ fn raw_schema() -> Arc { }) } -/// A raw schema of two float fields, so one series can fill one of them and another the other. +/// A raw schema of two float fields, used to verify a frame with different fields is rejected. fn two_field_schema() -> Arc { Arc::new(SeriesSchema { fields: Arc::from([ @@ -182,102 +175,6 @@ 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, - modifier: 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, - modifier: 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, - modifier: 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 same tree asking for `quantile`, whose state the runtime does not keep yet. fn quantile_aggregate_tree() -> NativeTree { let complete = aggregate_tree(); @@ -329,37 +226,6 @@ fn aggregate_tree() -> NativeTree { .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, @@ -411,6 +277,66 @@ fn output_schema() -> SchemaRef { ])) } +/// 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 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, + ), + ] +} + fn input_plan(schema: SchemaRef) -> LogicalPlan { LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, @@ -447,15 +373,6 @@ fn memory_exec_partitioned(partitions: Vec>) -> 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, @@ -664,11 +581,6 @@ 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, @@ -829,87 +741,6 @@ fn rows(batches: &[RecordBatch]) -> Vec { 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\")]"; @@ -984,7 +815,7 @@ async fn a_runtime_error_after_rows_releases_the_input() { ); } -/// Two series with *different* names collapse once `rate` drops the name. +/// A collision produced inside one native execution is still refused by its local runtime state. /// /// 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 @@ -1489,85 +1320,6 @@ fn an_input_with_no_partitions_is_rejected_when_the_plan_is_built() { ); } -/// 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 @@ -1753,165 +1505,7 @@ fn error_chain(error: &DataFusionError) -> String { 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. +/// Each partition runs its own series and answers it independently. #[tokio::test] async fn every_partition_evaluates_its_own_series() { let partitions = [1.0, 2.0, 3.0] @@ -1998,7 +1592,7 @@ fn the_optimizer_inserts_no_exchange_between_the_input_and_this_node() { 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 + // The count is the input's, which is what the node reports and what it // runs one evaluator per. assert_eq!( 2, @@ -2118,654 +1712,6 @@ async fn a_partition_past_the_inputs_count_is_refused() { ); } -/// 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() { @@ -2806,81 +1752,3 @@ fn unequal_binding_and_output_schema_are_not_ordered_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/runs.rs b/src/promql/src/native/runs.rs index 8de325a5386..5322221ffab 100644 --- a/src/promql/src/native/runs.rs +++ b/src/promql/src/native/runs.rs @@ -21,12 +21,9 @@ //! 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. +//! The native runtime's partition-local [`crate::native::execute::Emitted`] table keeps this state +//! per label set and field as runs, so a chunk of many positions costs one run rather than one +//! entry per position. /// One run of grid positions, half-open, that hold a value in one field. #[derive(Clone, Copy, Debug)] diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 9cd23a87cc3..9cb5bdacba0 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -14,11 +14,11 @@ //! The native fused-series route of the PromQL planner. //! -//! [`PromPlanner::try_plan_native_fused`] offers a whole statement expression to the native -//! compiler and, when the compiler accepts it, replaces everything the existing path would plan -//! below the root with one [`FusedSeries`] node over the selector's own table scan. Declining is -//! the normal answer: every shape the compiler does not accept, and every query this route cannot -//! bind to a scan faithfully, keeps the plan the existing planner builds. +//! [`PromPlanner::try_plan_native_fused`] offers a small experimental subset to the native +//! compiler: a vector selector (optionally parenthesized), or nested unary `rate`, `increase`, +//! `sum_over_time`, and `abs` calls over a selector. It replaces that expression with one +//! [`FusedSeries`] node over the selector's table scan. Other expression shapes keep the existing +//! planner's route. //! //! What the node needs from here is three things the existing planner already resolves: the raw //! series schema the tree's input was compiled against (the selector's selected field columns), the @@ -40,7 +40,6 @@ use datafusion_expr::Expr; use datatypes::series::SeriesOrdering; use datatypes::series::samples::{FieldSpec, SeriesSchema}; use promql::extension_plan::{FusedSeries, FusedSeriesBinding}; -use promql::native::plan::reads_across_series; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector, @@ -92,20 +91,14 @@ const NATIVE_ORDERING: SeriesOrdering = SeriesOrdering::Contiguous; impl PromPlanner { /// Plans `prom_expr` as one [`FusedSeries`] node, or returns `None` to keep the existing plan. /// - /// The route is offered a whole expression because the compiler folds a range function and the - /// math functions over it into one tree: offering a sub-expression instead would compile a tree - /// that computes something the reader did not write. Everything else about the node comes from - /// the planner's resolved selector context: the table name, the time index column, the tag - /// columns and (after the `__field__` matcher) the field columns, which are exactly the raw - /// series schema the tree's input has to be compiled against. + /// The route is offered the whole expression so its tree computes the requested expression + /// over the planner's resolved selector context: table name, time index, tag columns, and + /// field columns (after any `__field__` matcher). These are the raw series schema the input + /// tree is compiled against. /// - /// `None` is not a failure. It says this query cannot take the route *unchanged*, and the - /// caller plans it with the existing planner as if this method did not exist. The two ways that - /// happens are a shape the compiler refuses (`DataFusionError::NotImplemented`, which names - /// what it found) and a query this route cannot bind faithfully — more than one selector, a - /// label matcher the bare scan below the node would not apply, alternative (`or`) matchers, an - /// `__field__` matcher that narrows the field set, a missing table, a field that is not - /// `Float64`, a `rate`/`increase` over a delta-temporality table, or an unusable grid. + /// `None` is not a failure: unsupported shapes and queries this route cannot bind faithfully + /// continue through the existing planner. Admission is limited to selectors and unary calls + /// named above, plus the scan, schema, matcher, and grid constraints checked below. pub(super) async fn try_plan_native_fused( &mut self, prom_expr: &PromExpr, @@ -121,9 +114,8 @@ impl PromPlanner { return Ok(None); } - // A tree reads exactly one raw input, and the scan below the node is that one selector's, - // so a shape that names no selector or more than one is not offered at all. - let Some(selector) = single_selector(prom_expr) else { + // Admit only a direct selector or the selected unary calls applied to one selector. + let Some(selector) = admitted_selector(prom_expr) else { return Ok(None); }; @@ -312,15 +304,6 @@ impl PromPlanner { Err(DataFusionError::NotImplemented(_)) => return Ok(None), Err(error) => return Err(error).context(DataFusionPlanningSnafu), }; - // A vector operation (`sort`, `absent`, `topk`, `histogram_quantile`, ...) reads the whole - // vector at each grid point rather than one series at a time, so it cannot be evaluated by - // one executor per scan partition. Declining here is the only place that still works: once - // the `FusedSeries` node exists the logical plan is fixed, and the categorizer merely - // declining to push it down would leave a node the physical layer must build anyway. - if reads_across_series(tree.root()) { - return Ok(None); - } - let binding = FusedSeriesBinding { series_schema, // The tree's own output is evaluated on the statement's grid, and the binding names @@ -470,32 +453,21 @@ fn metric_scan_and_table_id( Some((scan.clone(), filter.predicate.clone())) } -fn single_selector(expr: &PromExpr) -> Option<&VectorSelector> { - let mut selectors = Vec::new(); - collect_selectors(expr, &mut selectors); - match selectors.as_slice() { - [only] => Some(only), - _ => None, - } -} - -fn collect_selectors<'a>(expr: &'a PromExpr, out: &mut Vec<&'a VectorSelector>) { +fn admitted_selector(expr: &PromExpr) -> Option<&VectorSelector> { match expr { - PromExpr::VectorSelector(selector) => out.push(selector), - PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => out.push(vs), - PromExpr::Paren(ParenExpr { expr }) => collect_selectors(expr, out), - PromExpr::Call(call) => { - // A label rewrite adds a label the node's output schema cannot declare: that schema is - // built from the table's own tags, so folding `label_replace`/`label_join` in would - // silently drop the destination label the existing plan projects. - if matches!(call.func.name, "label_replace" | "label_join") { - return; - } - for arg in &call.args.args { - collect_selectors(arg, out); - } + PromExpr::VectorSelector(selector) => Some(selector), + // Range selectors are admitted only as arguments within one of these supported calls. + PromExpr::MatrixSelector(MatrixSelector { vs, .. }) => Some(vs), + PromExpr::Paren(ParenExpr { expr }) => admitted_selector(expr), + PromExpr::Call(call) + if matches!( + call.func.name, + "rate" | "increase" | "sum_over_time" | "abs" + ) && call.args.args.len() == 1 => + { + admitted_selector(&call.args.args[0]) } - _ => {} + _ => None, } } diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 2e22ee59003..40b342dc2ea 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -2377,13 +2377,27 @@ async fn native_fused_whole_vector_operation_falls_back() { #[tokio::test] async fn native_fused_unsupported_shape_falls_back() { - let plan = build_native_fused_plan("sum(rate(some_metric[5m]))", true).await; - - let display = plan.display_indent().to_string(); - assert!(!display.contains("FusedSeries"), "{display}"); - // Falling back is not an approximation of the existing plan: it is the existing plan. - let without_flag = build_native_fused_plan("sum(rate(some_metric[5m]))", false).await; - assert_eq!(display, without_flag.display_indent().to_string()); + for query in [ + "sum(rate(some_metric[5m]))", + "some_metric + some_metric", + "rate(some_metric[5m])[1h:]", + "histogram_count(some_metric)", + "ceil(some_metric)", + ] { + let plan = build_native_fused_plan(query, true).await; + let display = plan.display_indent().to_string(); + assert!( + !display.contains("FusedSeries"), + "`{query}` must not enter the limited native POC: {display}" + ); + // Falling back is the existing plan, not an approximation of it. + let without_flag = build_native_fused_plan(query, false).await; + assert_eq!( + display, + without_flag.display_indent().to_string(), + "`{query}` must keep the existing plan" + ); + } } /// A label rewrite stays on the existing plan, which is the one that projects its new label. diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 8e4aea3a85a..65a49b83a28 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -2080,6 +2080,125 @@ async fn native_promql_executes_fused_node_and_matches_legacy() { assert_eq!(native_rows, legacy_rows); } +#[tokio::test(flavor = "multi_thread")] +async fn native_promql_rate_and_abs_rate_match_legacy_across_storage_encodings() { + common_telemetry::init_default_ut_logging(); + let mito_config = mito2::config::MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }; + let (standalone, observer) = native_test_instance( + GreptimeDbStandaloneBuilder::new("native_promql_rate_storage_encodings") + .with_mito_config(mito_config), + ) + .await; + let instance = standalone.fe_instance().clone(); + + for (table, ddl, flush_table, encoding, sparse) in [ + ( + "native_rate_dense", + "CREATE TABLE native_rate_dense (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=mito WITH ('primary_key_encoding' = 'dense')", + "native_rate_dense", + PrimaryKeyEncoding::Dense, + false, + ), + ( + "native_rate_sparse", + "CREATE TABLE native_rate_sparse_phy (ts TIMESTAMP(3) TIME INDEX, val DOUBLE) ENGINE=metric WITH ('physical_metric_table' = 'true', 'primary_key_encoding' = 'sparse'); CREATE TABLE native_rate_sparse (host STRING, ts TIMESTAMP(3) TIME INDEX, val DOUBLE, PRIMARY KEY (host)) ENGINE=metric WITH ('on_physical_table' = 'native_rate_sparse_phy')", + "native_rate_sparse_phy", + PrimaryKeyEncoding::Sparse, + true, + ), + ] { + execute_all(&instance, ddl, QueryContext::arc()).await; + for (start, end) in [(0, 60), (120, 300)] { + let values = (b'a'..=b'p') + .flat_map(|host| { + (start..=end).step_by(60).map(move |seconds| { + format!( + "('{}', {}, {})", + host as char, + seconds * 1000, + seconds as f64 * (host - b'a' + 1) as f64 + ) + }) + }) + .collect::>() + .join(", "); + execute_all( + &instance, + &format!("INSERT INTO {table} (host, ts, val) VALUES {values}"), + QueryContext::arc(), + ) + .await; + if start == 0 { + execute_all( + &instance, + &format!("ADMIN FLUSH_TABLE('{flush_table}')"), + QueryContext::arc(), + ) + .await; + } + } + + let expected = (b'a'..=b'p') + .map(|host| { + ( + 300_000, + (host as char).to_string(), + (host - b'a' + 1) as f64, + ) + }) + .collect::>(); + for expression in [ + format!("rate({table}[5m])"), + format!("abs(rate({table}[5m]))"), + ] { + let sql = format!("TQL EVAL (300, 300, '1m') {expression}"); + let legacy_rows = + native_result_rows(native_tql_output(&instance, &sql, false).await).await; + assert_eq!( + legacy_rows, expected, + "native-off `{expression}` on {table}" + ); + + observer.clear(); + let native_rows = + native_result_rows(native_tql_output(&instance, &sql, true).await).await; + assert_eq!(native_rows, expected, "native-on `{expression}` on {table}"); + assert_eq!(native_rows, legacy_rows, "`{expression}` on {table}"); + let plan = observer.take(); + assert_eq!(plan.name(), "FusedSeriesExec"); + assert_eq!( + plan.children()[0].output_partitioning().partition_count(), + plan.output_partitioning().partition_count(), + "fused child and output partitions agree for {table}" + ); + assert!( + plan.output_partitioning().partition_count() > 1, + "fused execution must retain partitioning for {table}" + ); + if sparse { + assert_four_native_partitions(&plan); + } + let child = plan.children()[0]; + let child = child + .downcast_ref::() + .map_or(child, |cooperative| cooperative.input()); + let series_scan = child + .downcast_ref::() + .expect("fused node reads a prepared series scan"); + let RawLabelSource::PrimaryKey { + physical_metadata, .. + } = series_scan.label_source() + else { + panic!("fused scan must materialize labels from the primary key"); + }; + assert_eq!(physical_metadata.primary_key_encoding, encoding); + } + } +} + #[tokio::test(flavor = "multi_thread")] async fn native_promql_explain_and_analyze_execute_fused_plan() { common_telemetry::init_default_ut_logging(); From 9f1bf087ebe13464187e4bc1d51bf38f36561240 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sun, 25 Oct 2026 03:39:52 +0800 Subject: [PATCH 38/38] test(promql): use an instant-vector subquery boundary case Signed-off-by: discord9 --- src/query/src/promql/planner/test.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 40b342dc2ea..13229b11474 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -2380,7 +2380,7 @@ async fn native_fused_unsupported_shape_falls_back() { for query in [ "sum(rate(some_metric[5m]))", "some_metric + some_metric", - "rate(some_metric[5m])[1h:]", + "sum_over_time(rate(some_metric[5m])[1h:])", "histogram_count(some_metric)", "ceil(some_metric)", ] {