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 53bb3a187a0..b46a3c0374e 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -1041,7 +1041,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)