From c25b6a9fbdc0bf6ad7bbf194d315ebf1372c5985 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:05:42 +0800 Subject: [PATCH 1/3] feat(promql): fuse independent operands of hybrid binary queries Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/analyzer.rs | 55 ++- src/query/src/dist_plan/analyzer/test.rs | 112 ++++- src/query/src/promql/planner.rs | 75 +++- src/query/src/promql/planner/native.rs | 130 +++++- .../src/tests/native_aggregate_test.rs | 386 ++++++++++++++++++ 5 files changed, 732 insertions(+), 26 deletions(-) diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index f58a9ee3ff7..5300cb0b002 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -287,8 +287,9 @@ impl DistPlannerAnalyzer { 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; + let plan = rewrite_hybrid_binary_operands(plan.data, in_process_region_transport)?; + let mut rewriter = PlanRewriter::new(&plan, in_process_region_transport); + let result = plan.rewrite(&mut rewriter)?.data; Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } @@ -508,6 +509,56 @@ struct PlanRewriter { new_child_plan: Option, } +fn native_operand_wrapper_child(plan: &LogicalPlan) -> Option<&LogicalPlan> { + match plan { + LogicalPlan::SubqueryAlias(alias) => Some(alias.input.as_ref()), + LogicalPlan::Projection(projection) => Some(projection.input.as_ref()), + LogicalPlan::Filter(filter) => Some(filter.input.as_ref()), + LogicalPlan::Window(window) => Some(window.input.as_ref()), + _ => None, + } +} + +fn native_operand_fused(plan: &LogicalPlan) -> Option<&LogicalPlan> { + let mut node = plan; + while let Some(child) = native_operand_wrapper_child(node) { + node = child; + } + (is_fused_series_plan(node) && matches!(node.inputs().as_slice(), [LogicalPlan::TableScan(_)])) + .then_some(node) +} + +fn is_native_operand_subtree(plan: &LogicalPlan) -> bool { + native_operand_fused(plan).is_some() +} + +fn rewrite_hybrid_binary_operands(plan: LogicalPlan, transport: bool) -> DfResult { + plan.transform_up(|plan| { + let LogicalPlan::Join(join) = &plan else { + return Ok(Transformed::no(plan)); + }; + if !is_native_operand_subtree(&join.left) || !is_native_operand_subtree(&join.right) { + return Ok(Transformed::no(plan)); + } + fn rewrite_side(side: &LogicalPlan, transport: bool) -> DfResult { + if let Some(child) = native_operand_wrapper_child(side) { + let rewritten_child = rewrite_side(child, transport)?; + return side + .clone() + .with_new_exprs(side.expressions(), vec![rewritten_child]); + } + let mut rw = PlanRewriter::new(side, transport); + Ok(side.clone().rewrite(&mut rw)?.data) + } + let left = rewrite_side(&join.left, transport)?; + let right = rewrite_side(&join.right, transport)?; + let mut rebuilt = LogicalPlan::Join(join.clone()); + rebuilt = rebuilt.with_new_exprs(rebuilt.expressions(), vec![left, right])?; + Ok(Transformed::yes(rebuilt)) + }) + .map(|transformed| transformed.data) +} + fn is_admitted_pair_plan(plan: &LogicalPlan) -> bool { let LogicalPlan::Extension(extension) = plan else { return false; diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index b512655dd2c..b363c27ff6b 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -16,6 +16,7 @@ use std::collections::{BTreeSet, HashSet}; use std::pin::Pin; use std::sync::Arc; +use crate::promql::native_compile::{compile, statement_grid}; use arrow::datatypes::{DataType, IntervalDayTime, TimeUnit}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, MITO_ENGINE}; use common_error::ext::BoxedError; @@ -43,10 +44,14 @@ use datafusion_functions::datetime::expr_fn::now; use datafusion_functions::unicode::expr_fn::substring; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; +use datatypes::series::SeriesOrdering; +use datatypes::series::samples::{FieldSpec, SeriesSchema}; use futures::Stream; use futures::task::{Context, Poll}; use pretty_assertions::assert_eq; -use promql::extension_plan::{HistogramFold, HistogramFoldOperation}; +use promql::extension_plan::{ + FusedSeries, FusedSeriesBinding, HistogramFold, HistogramFoldOperation, +}; use regex::Regex; use store_api::data_source::DataSource; use store_api::storage::ScanRequest; @@ -322,6 +327,111 @@ impl Stream for EmptyStream { } } +fn hybrid_fused(scan: LogicalPlan, metric: &str) -> LogicalPlan { + let stmt = promql_parser::parser::EvalStmt { + expr: promql_parser::parser::parse(metric).unwrap(), + start: std::time::UNIX_EPOCH, + end: std::time::UNIX_EPOCH + std::time::Duration::from_secs(10), + interval: std::time::Duration::from_secs(1), + lookback_delta: std::time::Duration::from_secs(1), + }; + let series = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { name: "v0".into() }]), + eval_grid: None, + }); + let tree = compile(&stmt, series, SeriesOrdering::KeyThenTime).unwrap(); + let fused = FusedSeries::try_new( + Arc::from(metric), + tree, + FusedSeriesBinding { + grid: statement_grid(&stmt).unwrap(), + time_index: "ts".into(), + tags: vec![], + metric_name: Some(Arc::from(metric)), + }, + vec![scan], + Arc::new( + datafusion_common::DFSchema::try_from(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("v0", DataType::Float64, true), + ])) + .unwrap(), + ), + ) + .unwrap(); + LogicalPlan::Extension(Extension { + node: Arc::new(fused), + }) +} + +fn count_merge_scans(plan: &LogicalPlan) -> usize { + usize::from( + matches!(plan, LogicalPlan::Extension(extension) if extension.node.as_any().is::()), + ) + plan + .inputs() + .iter() + .map(|input| count_merge_scans(input)) + .sum::() +} + +#[test] +fn hybrid_binary_join_rewrites_each_native_operand_boundary() { + let left_table = TestTable::table_with_name(0, "left_metric".into()); + let right_table = TestTable::table_with_name(1, "right_metric".into()); + let scan = |table: TableRef| { + LogicalPlanBuilder::scan( + table.table_info().name.clone(), + Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))), + None, + ) + .unwrap() + .build() + .unwrap() + }; + let left = LogicalPlanBuilder::from(hybrid_fused(scan(left_table), "left_metric")) + .alias(TableReference::bare("left_metric")) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(hybrid_fused(scan(right_table), "right_metric")) + .alias(TableReference::bare("right_metric")) + .unwrap() + .build() + .unwrap(); + let plan = LogicalPlanBuilder::from(left) + .join( + LogicalPlanBuilder::from(right).build().unwrap(), + JoinType::Inner, + (vec![Column::from_name("v0")], vec![Column::from_name("v0")]), + None, + ) + .unwrap() + .build() + .unwrap(); + let rewritten = DistPlannerAnalyzer::new(true).try_push_down(plan).unwrap(); + let LogicalPlan::Join(join) = &rewritten else { + panic!("Join must remain frontend: {rewritten}"); + }; + for side in [&join.left, &join.right] { + assert_eq!(1, count_merge_scans(side), "{rewritten}"); + let text = side.to_string(); + assert!( + text.contains("FusedSeries") && text.contains("MergeScan"), + "{text}" + ); + } + assert_eq!( + 2, + count_merge_scans(&rewritten), + "Join must be outside both MergeScans: {rewritten}" + ); + assert!(rewritten.to_string().contains("FusedSeries")); + assert!( + join.left.to_string().contains("TableScan") && join.right.to_string().contains("TableScan") + ); +} + fn find_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { if let LogicalPlan::Extension(extension) = plan && let Some(merge_scan) = extension diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index c6c96a33156..09397fbbc7e 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -586,7 +586,11 @@ impl PromPlanner { // schema. self.ctx.start = end; self.ctx.end = end; - let plan = self.prom_expr_to_plan(expr, query_engine_state).await?; + let native_fused = self.native_fused; + self.native_fused = false; + let plan_result = self.prom_expr_to_plan(expr, query_engine_state).await; + self.native_fused = native_fused; + let plan = plan_result?; return Ok(LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, schema: plan.schema().clone(), @@ -594,7 +598,11 @@ impl PromPlanner { }; self.ctx.start = start; self.ctx.end = end; - self.prom_expr_to_plan(expr, query_engine_state).await + let native_fused = self.native_fused; + self.native_fused = false; + let plan_result = self.prom_expr_to_plan(expr, query_engine_state).await; + self.native_fused = native_fused; + plan_result } async fn prom_subquery_expr_to_plan( @@ -646,10 +654,14 @@ impl PromPlanner { self.ctx.start = child_start; self.ctx.end = child_end; } - let input = self.prom_expr_to_plan(expr, query_engine_state).await?; + let native_fused = self.native_fused; + self.native_fused = false; + let input_result = self.prom_expr_to_plan(expr, query_engine_state).await; + self.native_fused = native_fused; self.ctx.interval = current_interval; self.ctx.start = current_start; self.ctx.end = current_end; + let input = input_result?; // Return no rows with the child schema so enclosing fallbacks still work. let input = match child_window { Some(_) => input, @@ -1365,32 +1377,57 @@ impl PromPlanner { } // both are columns. join them on time index (None, None) => { - let mut left_input = self.prom_expr_to_plan(lhs, query_engine_state).await?; - let left_field_columns = self.ctx.field_columns.clone(); - let left_time_index_column = self.ctx.time_index_column.clone(); + let native_pair = + self.native_fused && native::hybrid_binary_admitted(*op, modifier, lhs, rhs); + let native = if native_pair { + match ( + self.try_plan_native_operand(lhs, query_engine_state) + .await?, + self.try_plan_native_operand(rhs, query_engine_state) + .await?, + ) { + (Some(left), Some(right)) => Some((left, right)), + _ => None, + } + } else { + None + }; + let used_native = native.is_some(); + let (mut left_input, mut left_context, mut right_input, mut right_context) = + if let Some(((left, left_ctx), (right, right_ctx))) = native { + (left, left_ctx, right, right_ctx) + } else { + let left = self.prom_expr_to_plan(lhs, query_engine_state).await?; + let left_ctx = self.ctx.clone(); + let right = self.prom_expr_to_plan(rhs, query_engine_state).await?; + let right_ctx = self.ctx.clone(); + (left, left_ctx, right, right_ctx) + }; + let left_field_columns = left_context.field_columns.clone(); + let left_time_index_column = left_context.time_index_column.clone(); + self.ctx = left_context.clone(); let mut left_table_ref = self .table_ref() .unwrap_or_else(|_| TableReference::bare("")); - let mut left_context = self.ctx.clone(); - - let mut right_input = self.prom_expr_to_plan(rhs, query_engine_state).await?; - let right_field_columns = self.ctx.field_columns.clone(); - let right_time_index_column = self.ctx.time_index_column.clone(); + let right_field_columns = right_context.field_columns.clone(); + let right_time_index_column = right_context.time_index_column.clone(); + self.ctx = right_context.clone(); let mut right_table_ref = self .table_ref() .unwrap_or_else(|_| TableReference::bare("")); - let mut right_context = self.ctx.clone(); // Both operands are planned first because only the planned contexts tell a tag // from a value field. The rewrite merely adds matchers to a selector, so the // table reference, time index and field columns captured above stay valid. - if let Some(rewritten) = matching_filters::propagate( - binary_expr, - &left_context.tag_columns, - &right_context.tag_columns, - &left_context.aggregation_field_labels, - &right_context.aggregation_field_labels, - ) { + if !used_native + && let Some(rewritten) = matching_filters::propagate( + binary_expr, + &left_context.tag_columns, + &right_context.tag_columns, + &left_context.aggregation_field_labels, + &right_context.aggregation_field_labels, + ) + { // A copied matcher belongs to the scan, not to the operand's identity: // `absent()` turns `selector_matcher` into the labels it reports, so the // re-planned context keeps the matchers the operand was written with. diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index 05083212150..f8e503bc4c4 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -40,7 +40,7 @@ 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; -use datafusion::logical_expr::{Extension, LogicalPlan, TableScanBuilder}; +use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder, TableScanBuilder}; use datafusion::prelude::Expr as DfExpr; use datafusion_expr::Expr; use datatypes::series::SeriesOrdering; @@ -51,7 +51,8 @@ use promql::vector_match::{Cardinality, MatchingLabels}; use promql_parser::label::Matchers; use promql_parser::parser::value::ValueType; use promql_parser::parser::{ - EvalStmt, Expr as PromExpr, MatrixSelector, ParenExpr, VectorSelector, + BinModifier, EvalStmt, Expr as PromExpr, LabelModifier, MatrixSelector, ParenExpr, + VectorMatchCardinality, VectorSelector, token, }; use session::context::QueryContextRef; use snafu::ResultExt; @@ -109,7 +110,102 @@ type PreparedInput = ( Vec, ); +pub(crate) fn hybrid_binary_admitted( + op: token::TokenType, + modifier: &Option, + lhs: &PromExpr, + rhs: &PromExpr, +) -> bool { + let supported_op = matches!( + op.id(), + token::T_ADD + | token::T_SUB + | token::T_MUL + | token::T_DIV + | token::T_MOD + | token::T_POW + | token::T_ATAN2 + ); + let supported_matching = match modifier { + None => true, + Some(modifier) => { + matches!(&modifier.card, VectorMatchCardinality::OneToOne) + && (matches!(&modifier.matching, Some(LabelModifier::Include(labels)) if !labels.labels.is_empty()) + || matches!(&modifier.matching, Some(LabelModifier::Exclude(labels)) if labels.labels.is_empty())) + } + }; + let has_aggregate = |expr: &PromExpr| { + fn visit(expr: &PromExpr) -> bool { + match expr { + PromExpr::Aggregate(_) => true, + PromExpr::Paren(x) => visit(&x.expr), + PromExpr::Unary(x) => visit(&x.expr), + PromExpr::Binary(x) => visit(&x.lhs) || visit(&x.rhs), + PromExpr::Call(x) => x.args.args.iter().any(|arg| visit(arg)), + _ => false, + } + } + visit(expr) + }; + let one_selector = |expr: &PromExpr| { + let mut selectors = Vec::new(); + collect_selectors(expr, &mut selectors) && selectors.len() == 1 + }; + supported_op + && supported_matching + && lhs.value_type() == ValueType::Vector + && rhs.value_type() == ValueType::Vector + && admitted_single_selector(lhs) + && admitted_single_selector(rhs) + && !has_aggregate(lhs) + && !has_aggregate(rhs) + && one_selector(lhs) + && one_selector(rhs) +} + impl PromPlanner { + pub(super) async fn try_plan_native_operand( + &mut self, + expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result> { + let mut selectors = Vec::new(); + if !collect_selectors(expr, &mut selectors) + || selectors.len() != 1 + || !admitted_single_selector(expr) + { + return Ok(None); + } + let selector = selectors[0]; + let saved = self.ctx.clone(); + let plan = match self.try_plan_native_fused(expr, query_engine_state).await? { + Some(plan) => plan, + None => return Ok(None), + }; + self.ctx = saved.clone(); + if self + .preprocess_label_matchers(&selector.matchers, &selector.name) + .is_err() + { + self.ctx = saved; + return Ok(None); + } + if self.setup_context().await?.is_some() { + self.ctx = saved; + return Ok(None); + } + let mut operand_ctx = self.ctx.clone(); + operand_ctx.use_tsid = false; + let table_ref = self.table_ref()?; + self.ctx = saved; + let plan = LogicalPlanBuilder::from(plan) + .alias(table_ref) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + Ok(Some((plan, operand_ctx))) + } + /// Plans `prom_expr` as one native node, or returns `None` to keep the existing plan. /// /// `query_engine_state` is only read to build the *existing* plan of an aggregate query, whose @@ -895,8 +991,8 @@ mod tests { use super::{ NATIVE_PROMQL_EXTENSION_KEY, admitted_selector, admitted_single_selector, - collect_selectors, fused_output_tags, independent_scan_pair, native_promql_requested, - scan_time_filter, + collect_selectors, fused_output_tags, hybrid_binary_admitted, independent_scan_pair, + native_promql_requested, scan_time_filter, }; fn compiled_test_tree(query: &str) -> NativeTree { @@ -1026,6 +1122,32 @@ mod tests { } } + fn hybrid_admission(query: &str) -> bool { + let expression = promql_parser::parser::parse(query).unwrap(); + let promql_parser::parser::Expr::Binary(binary) = expression else { + panic!("expected binary expression: {query}"); + }; + hybrid_binary_admitted(binary.op, &binary.modifier, &binary.lhs, &binary.rhs) + } + + #[test] + fn hybrid_binary_admission_is_narrow() { + for query in ["A / B", "rate(A[1m]) / rate(B[1m])"] { + assert!(hybrid_admission(query), "{query}"); + } + assert!(hybrid_admission("A + on(job) B")); + for query in [ + "A and B", + "A + on() B", + "A + on(job) group_left B", + "sum(A) / B", + "100 / B", + "A / (B + C)", + ] { + assert!(!hybrid_admission(query), "{query}"); + } + } + #[test] fn millisecond_scan_bounds_preserve_open_closed_and_maximum_edges() { let inclusive = scan_time_filter(0, 10, true, "ts").to_string(); diff --git a/tests-integration/src/tests/native_aggregate_test.rs b/tests-integration/src/tests/native_aggregate_test.rs index ab8617dde98..0f4721c8bf1 100644 --- a/tests-integration/src/tests/native_aggregate_test.rs +++ b/tests-integration/src/tests/native_aggregate_test.rs @@ -116,6 +116,7 @@ const EVAL_MILLIS: i64 = EVAL_SECS * 1000; const GAUGE_TABLE: &str = "native_agg_gauge"; /// The table the `rate` queries read: the same five series as [`GAUGE_TABLE`], holding a counter. const COUNTER_TABLE: &str = "native_agg_counter"; +const COUNTER_RHS_TABLE: &str = "native_agg_counter_rhs"; /// The table the range-window smoke test reads. const RANGE_TABLE: &str = "native_agg_range"; /// The tolerance of the value comparisons: relative, with an absolute floor. @@ -434,6 +435,382 @@ async fn run_query( ) } +#[tokio::test(flavor = "multi_thread")] +async fn native_hybrid_binary_05_06_08() { + common_telemetry::init_default_ut_logging(); + let observer = Arc::new(NativePlanObserver::default()); + let plugins = Plugins::new(); + plugins.insert(Arc::clone(&observer) as PhysicalPlanWrapperRef); + let standalone = GreptimeDbStandaloneBuilder::new("native_hybrid_binary") + .with_query_options(QueryOptions { + experimental_native_promql: true, + ..Default::default() + }) + .with_plugin(plugins) + .build() + .await; + let instance = Arc::clone(standalone.fe_instance()); + setup_fixture(&instance).await; + check_hybrid_binary_cases(&instance, &observer).await; + let subquery = format!( + "max_over_time((rate({COUNTER_TABLE}[5m]) / on(job, instance) rate({COUNTER_RHS_TABLE}[5m]))[2m:1m])" + ); + let (legacy, legacy_capture) = run_query(&instance, &subquery, false, &observer).await; + let (native, native_capture) = run_query(&instance, &subquery, true, &observer).await; + assert!( + legacy_capture.fused.is_empty() && native_capture.fused.is_empty(), + "hybrid must be disabled inside subquery: {subquery}" + ); + assert_eq!( + legacy.rows.len(), + native.rows.len(), + "subquery legacy/native row count" + ); + assert_rows( + &native.rows, + &legacy + .rows + .iter() + .map(|row| ResultRow { + ts_millis: row.ts_millis, + labels: row.labels.clone(), + value: row.value, + }) + .collect::>(), + &format!("subquery fallback `{subquery}`"), + ); + check_reduced_label_matching(&instance, &observer).await; +} + +async fn check_reduced_label_matching(instance: &Arc, observer: &NativePlanObserver) { + let unique_query = format!( + "{GAUGE_TABLE}{{instance=~\"h1|h4\"}} / on(job) {COUNTER_RHS_TABLE}{{instance=~\"h1|h4\"}}" + ); + let (legacy, legacy_capture) = run_query(instance, &unique_query, false, observer).await; + assert!(legacy_capture.fused.is_empty()); + let mark = observer.mark(); + let (native, capture) = run_query(instance, &unique_query, true, observer).await; + assert_eq!( + legacy.rows.len(), + native.rows.len(), + "{unique_query}: row count" + ); + assert_rows( + &native.rows, + &legacy + .rows + .iter() + .map(|row| ResultRow { + ts_millis: row.ts_millis, + labels: row.labels.clone(), + value: row.value, + }) + .collect::>(), + &format!("on(job) native == legacy `{unique_query}`"), + ); + assert_eq!( + 1, + native.rows.len(), + "the matching instance filters leave one matching job" + ); + assert_eq!(2, capture.fused.len(), "{unique_query}"); + let plans = observer.since(mark); + let frontend = plans + .iter() + .find(|plan| find_binary_join(plan).is_some()) + .expect("frontend Join plan"); + let join = find_binary_join(frontend).unwrap(); + let uniqueness = join + .children() + .iter() + .map(|branch| uniqueness_operators_above_merge(branch)) + .fold((false, false), |(window, assertion_filter), found| { + (window || found.0, assertion_filter || found.1) + }); + assert!( + uniqueness.0, + "on(job) requires a frontend uniqueness WindowAggExec above MergeScan: {}", + displayable(frontend.as_ref()).indent(false) + ); + assert!( + uniqueness.1, + "on(job) requires a frontend prom_assert_unique_match_group FilterExec above MergeScan: {}", + displayable(frontend.as_ref()).indent(false) + ); + for branch in join.children() { + assert!( + contains_execution_node(branch, "MergeScanExec"), + "on(job) FE branch must retain remote boundary: {}", + displayable(frontend.as_ref()).indent(false) + ); + assert!( + !contains_execution_node(branch, FUSED_EXEC_NAME), + "on(job) fused node must stay in region" + ); + } + for fused in &capture.fused { + assert_fused_reads_scan_directly(fused, &unique_query); + } + + let duplicate_query = format!("{GAUGE_TABLE} / on(job) {COUNTER_RHS_TABLE}"); + for enabled in [false, true] { + let sql = tql_eval(&duplicate_query); + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, SCAN_PARTITIONS.to_string()); + query_ctx.set_extension(NATIVE_PROMQL_EXTENSION_KEY, enabled.to_string()); + let output = instance + .do_query(&sql, Arc::new(query_ctx)) + .await + .remove(0) + .expect("duplicate query should produce an execution stream"); + let error = match output.data { + OutputData::Stream(stream) => collect_batches(stream) + .await + .expect_err("duplicate on(job) matches must fail"), + OutputData::RecordBatches(_) => panic!("duplicate on(job) matches must fail"), + other => panic!("unexpected duplicate-query output: {other:?}"), + }; + let message = error.to_string().to_ascii_lowercase(); + assert!( + message.contains("many-to-many") + || message.contains("multiple matches") + || message.contains("duplicate"), + "`{sql}` should fail for duplicate matching groups, got: {error}" + ); + } +} + +async fn check_hybrid_binary_cases(instance: &Arc, observer: &NativePlanObserver) { + let cases = [ + ( + format!("{GAUGE_TABLE} / on(job, instance) {COUNTER_RHS_TABLE}{{instance!=\"h5\"}}"), + 4usize, + ), + ( + format!( + "100 * (1 - ({GAUGE_TABLE} / on(job, instance) {COUNTER_RHS_TABLE}{{instance!=\"h5\"}}))" + ), + 4usize, + ), + ( + format!("rate({COUNTER_TABLE}[5m]) / on(job, instance) rate({COUNTER_RHS_TABLE}[5m])"), + 5usize, + ), + ]; + for (index, (promql, expected_rows)) in cases.into_iter().enumerate() { + let (legacy, legacy_capture) = run_query(instance, &promql, false, observer).await; + assert!(legacy_capture.fused.is_empty()); + let mark = observer.mark(); + let (native, capture) = run_query(instance, &promql, true, observer).await; + assert_eq!( + legacy.rows.len(), + native.rows.len(), + "{promql}: full row count" + ); + assert_rows( + &native.rows, + &legacy + .rows + .iter() + .map(|row| ResultRow { + ts_millis: row.ts_millis, + labels: row.labels.clone(), + value: row.value, + }) + .collect::>(), + &format!("native == legacy `{promql}`"), + ); + assert_eq!(expected_rows, native.rows.len(), "{promql}"); + assert_eq!( + legacy + .columns + .iter() + .filter(|column| column.ends_with("ts")) + .count(), + 1, + "{promql}: legacy timestamp column" + ); + assert_eq!( + native + .columns + .iter() + .filter(|column| column.ends_with("ts")) + .count(), + 1, + "{promql}: native timestamp column" + ); + if index == 2 { + assert!( + native + .rows + .iter() + .all(|row| row.value.is_some_and(|value| close_enough(value, 1.0))), + "{promql}: independent rates must have ratio 1: {:#?}", + native.rows + ); + } + if index < 2 { + assert!( + native + .rows + .iter() + .all(|row| row.labels.get("instance").is_none_or(|v| v != "h5")), + "one-to-one match must drop unmatched series" + ); + } + assert_eq!(2, capture.fused.len(), "{promql}: plans since {mark}"); + let region_plans = observer + .since(mark) + .into_iter() + .filter(|plan| contains_execution_node(plan, "SeriesScanExec")) + .collect::>(); + let scan_tables = region_plans + .iter() + .map(|plan| { + let text = displayable(plan.as_ref()).indent(false).to_string(); + if text.contains("native_agg_gauge") { + GAUGE_TABLE + } else if text.contains(COUNTER_RHS_TABLE) { + COUNTER_RHS_TABLE + } else if text.contains(COUNTER_TABLE) { + COUNTER_TABLE + } else { + panic!("{promql}: unable to identify scan table: {text}") + } + }) + .collect::>(); + assert!( + scan_tables.contains(&GAUGE_TABLE) || scan_tables.contains(&COUNTER_TABLE), + "{promql}: gauge/counter source missing: {scan_tables:?}" + ); + assert!( + scan_tables.contains(&COUNTER_RHS_TABLE), + "{promql}: independent RHS scan missing: {scan_tables:?}" + ); + for fused in &capture.fused { + assert_fused_reads_scan_directly(fused, &format!("hybrid `{promql}`")); + } + for fused in &capture.fused { + let per_fused = NativeCapture { + fused: vec![Arc::clone(fused)], + finalize: Vec::new(), + }; + let counts = + fused_output_values_per_partition(&per_fused, &format!("hybrid `{promql}`")); + assert!( + counts.iter().all(|count| count.is_some_and(|n| n > 0)), + "{promql}: each scan partition must emit values, got {counts:?}" + ); + } + let plans = observer.since(mark); + let frontend = plans + .iter() + .find(|plan| find_binary_join(plan).is_some()) + .unwrap_or_else(|| panic!("{promql}: FE JoinExec missing; plans: {plans:#?}")); + let join = find_binary_join(frontend).unwrap(); + assert_eq!( + 2, + join.children().len(), + "{promql}: FE join branches: {}", + displayable(frontend.as_ref()).indent(false) + ); + for branch in join.children() { + assert!( + contains_execution_node(branch, "MergeScanExec"), + "{promql}: FE Join branch lacks MergeScanExec: {}", + displayable(frontend.as_ref()).indent(false) + ); + assert!( + !contains_execution_node(branch, FUSED_EXEC_NAME), + "{promql}: FusedSeriesExec leaked onto frontend: {}", + displayable(frontend.as_ref()).indent(false) + ); + } + assert!( + !is_nested_under_merge(frontend, "HashJoinExec"), + "{promql}: FE Join is under MergeScan: {}", + displayable(frontend.as_ref()).indent(false) + ); + if index == 1 { + assert!( + has_projection_parent_of_join(frontend), + "06 outer scalar arithmetic must project above JoinExec: {}", + displayable(frontend.as_ref()).indent(false) + ); + } + } +} + +fn find_binary_join(plan: &Arc) -> Option<&Arc> { + if matches!( + plan.name(), + "HashJoinExec" | "NestedLoopJoinExec" | "CrossJoinExec" + ) { + return Some(plan); + } + plan.children() + .iter() + .find_map(|child| find_binary_join(child)) +} + +fn find_execution_node<'a>( + plan: &'a Arc, + name: &str, +) -> Option<&'a Arc> { + if plan.name() == name { + return Some(plan); + } + plan.children() + .iter() + .find_map(|child| find_execution_node(child, name)) +} + +fn contains_execution_node(plan: &Arc, name: &str) -> bool { + find_execution_node(plan, name).is_some() +} + +/// Finds the PromQL uniqueness validation operators on the frontend side of a remote scan. +/// Traversal stops at each MergeScan so it cannot accidentally match region-side execution. +fn uniqueness_operators_above_merge(plan: &Arc) -> (bool, bool) { + if plan.name() == "MergeScanExec" { + return (false, false); + } + + let mut has_window = matches!(plan.name(), "WindowAggExec" | "BoundedWindowAggExec"); + let mut has_assertion_filter = plan.name() == "FilterExec" + && format!("{plan:?}").contains("prom_assert_unique_match_group"); + for child in plan.children() { + let (child_window, child_filter) = uniqueness_operators_above_merge(child); + has_window |= child_window; + has_assertion_filter |= child_filter; + } + (has_window, has_assertion_filter) +} + +fn is_nested_under_merge(plan: &Arc, name: &str) -> bool { + fn walk(plan: &Arc, under_merge: bool, name: &str) -> bool { + if under_merge && plan.name() == name { + return true; + } + plan.children() + .iter() + .any(|child| walk(child, under_merge || plan.name() == "MergeScanExec", name)) + } + walk(plan, false, name) +} + +fn has_projection_parent_of_join(plan: &Arc) -> bool { + plan.name() == "ProjectionExec" + && plan + .children() + .iter() + .any(|child| find_binary_join(child).is_some()) + || plan + .children() + .iter() + .any(|child| has_projection_parent_of_join(child)) +} + /// Executes one statement and drops its result. async fn execute_sql(instance: &Arc, sql: &str) { for result in instance.do_query(sql, QueryContext::arc()).await { @@ -499,6 +876,7 @@ async fn setup_fixture(instance: &Arc) { } execute_sql(instance, &counter_table_sql()).await; + execute_sql(instance, &format!("CREATE TABLE {COUNTER_RHS_TABLE} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance)) WITH ('compaction.type' = 'twcs', 'compaction.twcs.trigger_file_num' = '100')")).await; let (counter_first, counter_second) = ( COUNTER_FIRST_SECS * 1000, (COUNTER_FIRST_SECS + COUNTER_STEP_SECS) * 1000, @@ -521,8 +899,14 @@ async fn setup_fixture(instance: &Arc) { ), ) .await; + execute_sql(instance, &format!("INSERT INTO {COUNTER_RHS_TABLE} (ts, job, instance, v) VALUES ({counter_first}, '{job}', '{host}', {COUNTER_FIRST_VALUE:?}), ({counter_second}, '{job}', '{host}', {second_value:?})")).await; if index + 1 < COUNTER_SERIES.len() { execute_sql(instance, &format!("ADMIN FLUSH_TABLE('{COUNTER_TABLE}')")).await; + execute_sql( + instance, + &format!("ADMIN FLUSH_TABLE('{COUNTER_RHS_TABLE}')"), + ) + .await; } } @@ -1186,6 +1570,8 @@ async fn native_aggregate_complete_over_four_nonempty_mito_partitions() { check_case(&instance, &observer, case).await; } + check_hybrid_binary_cases(&instance, &observer).await; + // The shapes the native compiler declines, each declined for its own reason, and each checked on // the plans that really ran: none may hold a fused producer or a frontend merge, and both routes // must answer what the fixture derives by hand. From f882c1182d0ed35b32765fcb9fb77edbaa625f93 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:02:38 +0800 Subject: [PATCH 2/3] refactor(promql): reuse prepared native operand context Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/planner/native.rs | 108 ++++++++++++------------- 1 file changed, 50 insertions(+), 58 deletions(-) diff --git a/src/query/src/promql/planner/native.rs b/src/query/src/promql/planner/native.rs index f8e503bc4c4..5fcbb7ca0d3 100644 --- a/src/query/src/promql/planner/native.rs +++ b/src/query/src/promql/planner/native.rs @@ -134,33 +134,12 @@ pub(crate) fn hybrid_binary_admitted( || matches!(&modifier.matching, Some(LabelModifier::Exclude(labels)) if labels.labels.is_empty())) } }; - let has_aggregate = |expr: &PromExpr| { - fn visit(expr: &PromExpr) -> bool { - match expr { - PromExpr::Aggregate(_) => true, - PromExpr::Paren(x) => visit(&x.expr), - PromExpr::Unary(x) => visit(&x.expr), - PromExpr::Binary(x) => visit(&x.lhs) || visit(&x.rhs), - PromExpr::Call(x) => x.args.args.iter().any(|arg| visit(arg)), - _ => false, - } - } - visit(expr) - }; - let one_selector = |expr: &PromExpr| { - let mut selectors = Vec::new(); - collect_selectors(expr, &mut selectors) && selectors.len() == 1 - }; supported_op && supported_matching && lhs.value_type() == ValueType::Vector && rhs.value_type() == ValueType::Vector - && admitted_single_selector(lhs) - && admitted_single_selector(rhs) - && !has_aggregate(lhs) - && !has_aggregate(rhs) - && one_selector(lhs) - && one_selector(rhs) + && admitted_selector(lhs).is_some() + && admitted_selector(rhs).is_some() } impl PromPlanner { @@ -169,35 +148,24 @@ impl PromPlanner { expr: &PromExpr, query_engine_state: &QueryEngineState, ) -> Result> { - let mut selectors = Vec::new(); - if !collect_selectors(expr, &mut selectors) - || selectors.len() != 1 - || !admitted_single_selector(expr) - { + if admitted_selector(expr).is_none() { return Ok(None); } - let selector = selectors[0]; let saved = self.ctx.clone(); - let plan = match self.try_plan_native_fused(expr, query_engine_state).await? { - Some(plan) => plan, - None => return Ok(None), - }; - self.ctx = saved.clone(); - if self - .preprocess_label_matchers(&selector.matchers, &selector.name) - .is_err() + let (plan, operand_ctx) = match self + .try_plan_native_fused_with_context(expr, query_engine_state) + .await? { - self.ctx = saved; - return Ok(None); - } - if self.setup_context().await?.is_some() { - self.ctx = saved; - return Ok(None); - } - let mut operand_ctx = self.ctx.clone(); - operand_ctx.use_tsid = false; - let table_ref = self.table_ref()?; + Some((plan, Some(operand_ctx))) => (plan, operand_ctx), + _ => { + self.ctx = saved; + return Ok(None); + } + }; + self.ctx = operand_ctx.clone(); + let table_ref = self.table_ref(); self.ctx = saved; + let table_ref = table_ref?; let plan = LogicalPlanBuilder::from(plan) .alias(table_ref) .context(DataFusionPlanningSnafu)? @@ -229,6 +197,17 @@ impl PromPlanner { prom_expr: &PromExpr, query_engine_state: &QueryEngineState, ) -> Result> { + Ok(self + .try_plan_native_fused_with_context(prom_expr, query_engine_state) + .await? + .map(|(plan, _)| plan)) + } + + async fn try_plan_native_fused_with_context( + &mut self, + prom_expr: &PromExpr, + query_engine_state: &QueryEngineState, + ) -> Result)>> { if !self.native_fused { return Ok(None); } @@ -284,13 +263,17 @@ impl PromPlanner { output_schema, ) .context(DataFusionPlanningSnafu)?; - return Ok(Some(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))); + return Ok(Some(( + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + None, + ))); } let original_context = self.ctx.clone(); let preparation = async { let mut prepared: Vec = Vec::new(); + let mut single_input_context = None; for selector in &selectors { self.ctx = original_context.clone(); // These matchers are alternatives the one scan below the node cannot represent, so the @@ -306,6 +289,9 @@ impl PromPlanner { if self.setup_context().await?.is_some() { return Ok(None); } + if selectors.len() == 1 { + single_input_context = Some(self.ctx.clone()); + } let (Some(metric), Some(time)) = ( self.ctx.table_name.clone(), self.ctx.time_index_column.clone(), @@ -450,11 +436,11 @@ impl PromPlanner { } prepared.push((input, series_schema, Some(Arc::from(metric)), time, tags)); } - Ok(Some(prepared)) + Ok(Some((prepared, single_input_context))) } .await; self.ctx = original_context; - let Some(mut prepared) = preparation? else { + let Some((mut prepared, single_input_context)) = preparation? else { return Ok(None); }; // The time index, the raw tag list and metric name are shared binding properties. The @@ -612,9 +598,12 @@ impl PromPlanner { Err(DataFusionError::NotImplemented(_)) => return Ok(None), Err(error) => return Err(error).context(DataFusionPlanningSnafu), }; - return Ok(Some(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))); + return Ok(Some(( + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + single_input_context, + ))); } let tree = compiled.tree; let root_fields = tree @@ -650,9 +639,12 @@ impl PromPlanner { output_schema, ) .context(DataFusionPlanningSnafu)?; - Ok(Some(LogicalPlan::Extension(Extension { - node: Arc::new(node), - }))) + Ok(Some(( + LogicalPlan::Extension(Extension { + node: Arc::new(node), + }), + single_input_context, + ))) } /// The statement being planned, rebuilt from the evaluation numbers the planner holds. From 6b7d1bca4c33b795d0d386caf2de23eb11b0d82a Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:03:31 +0800 Subject: [PATCH 3/3] test(promql): strengthen hybrid binary review regressions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/analyzer/test.rs | 43 +++++++++--- .../src/tests/native_aggregate_test.rs | 69 ++++++++++++++++--- 2 files changed, 94 insertions(+), 18 deletions(-) diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index b363c27ff6b..a899c824679 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -413,23 +413,46 @@ fn hybrid_binary_join_rewrites_each_native_operand_boundary() { let LogicalPlan::Join(join) = &rewritten else { panic!("Join must remain frontend: {rewritten}"); }; - for side in [&join.left, &join.right] { + for (side, expected_table) in [(&join.left, "left_metric"), (&join.right, "right_metric")] { assert_eq!(1, count_merge_scans(side), "{rewritten}"); - let text = side.to_string(); - assert!( - text.contains("FusedSeries") && text.contains("MergeScan"), - "{text}" - ); + assert_no_fused_series_outside_merge_scan(side); + + let merge_scan = find_merge_scan(side).expect("each join operand has a MergeScan"); + let LogicalPlan::Extension(extension) = merge_scan.input() else { + panic!("MergeScan remote input must be a FusedSeries root"); + }; + let fused = extension + .node + .as_any() + .downcast_ref::() + .expect("MergeScan remote input must be a FusedSeries root"); + assert_eq!(expected_table, fused.promql()); + assert_eq!(1, fused.children().len()); + let [LogicalPlan::TableScan(table_scan)] = fused.children() else { + panic!("FusedSeries operand must directly consume a TableScan: {fused:?}"); + }; + assert_eq!(expected_table, table_scan.table_name.to_string()); } assert_eq!( 2, count_merge_scans(&rewritten), "Join must be outside both MergeScans: {rewritten}" ); - assert!(rewritten.to_string().contains("FusedSeries")); - assert!( - join.left.to_string().contains("TableScan") && join.right.to_string().contains("TableScan") - ); +} + +fn assert_no_fused_series_outside_merge_scan(plan: &LogicalPlan) { + if let LogicalPlan::Extension(extension) = plan { + if extension.node.as_any().is::() { + return; + } + assert!( + !extension.node.as_any().is::(), + "FusedSeries must be contained in a MergeScan remote input: {plan}" + ); + } + for input in plan.inputs() { + assert_no_fused_series_outside_merge_scan(input); + } } fn find_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { diff --git a/tests-integration/src/tests/native_aggregate_test.rs b/tests-integration/src/tests/native_aggregate_test.rs index 0f4721c8bf1..febfd6d5eff 100644 --- a/tests-integration/src/tests/native_aggregate_test.rs +++ b/tests-integration/src/tests/native_aggregate_test.rs @@ -119,6 +119,8 @@ const COUNTER_TABLE: &str = "native_agg_counter"; const COUNTER_RHS_TABLE: &str = "native_agg_counter_rhs"; /// The table the range-window smoke test reads. const RANGE_TABLE: &str = "native_agg_range"; +const HYBRID_MATCH_LEFT: &str = "native_hybrid_match_left"; +const HYBRID_MATCH_RIGHT: &str = "native_hybrid_match_right"; /// The tolerance of the value comparisons: relative, with an absolute floor. const FLOAT_TOLERANCE: f64 = 1e-10; @@ -451,6 +453,58 @@ async fn native_hybrid_binary_05_06_08() { .await; let instance = Arc::clone(standalone.fe_instance()); setup_fixture(&instance).await; + execute_sql(&instance, &format!("CREATE TABLE {HYBRID_MATCH_LEFT} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance))")).await; + execute_sql(&instance, &format!("CREATE TABLE {HYBRID_MATCH_RIGHT} (ts TIMESTAMP(3) TIME INDEX, job STRING, instance STRING, v DOUBLE, PRIMARY KEY (job, instance))")).await; + execute_sql( + &instance, + &format!("INSERT INTO {HYBRID_MATCH_LEFT} VALUES ({EVAL_MILLIS}, 'api', 'a1', 10)"), + ) + .await; + execute_sql(&instance, &format!("INSERT INTO {HYBRID_MATCH_RIGHT} VALUES ({EVAL_MILLIS}, 'api', 'b1', 2), ({EVAL_MILLIS}, 'other', 'b2', 3), ({EVAL_MILLIS}, 'other', 'b3', 4)" )).await; + let propagated = format!("{HYBRID_MATCH_LEFT}{{job=\"api\"}} / on(job) {HYBRID_MATCH_RIGHT}"); + let (legacy, _) = run_query(&instance, &propagated, false, &observer).await; + assert_eq!( + vec![5.0], + legacy + .rows + .iter() + .map(|row| row.value.unwrap()) + .collect::>() + ); + let mark = observer.mark(); + let (hybrid, capture) = run_query(&instance, &propagated, true, &observer).await; + assert_eq!( + vec![5.0], + hybrid + .rows + .iter() + .map(|row| row.value.unwrap()) + .collect::>() + ); + assert_eq!( + 2, + capture.fused.len(), + "both binary operands use native execution" + ); + let plans = observer.since(mark); + let frontend = plans + .iter() + .find(|plan| find_binary_join(plan).is_some()) + .expect("frontend Join plan"); + let join = find_binary_join(frontend).unwrap(); + let uniqueness = join + .children() + .iter() + .map(|branch| uniqueness_operators_above_merge(branch)) + .fold((false, false), |(window, filter), found| { + (window || found.0, filter || found.1) + }); + assert!( + uniqueness.0 && uniqueness.1, + "on(job) FE uniqueness operators: {}", + displayable(frontend.as_ref()).indent(false) + ); + check_hybrid_binary_cases(&instance, &observer).await; let subquery = format!( "max_over_time((rate({COUNTER_TABLE}[5m]) / on(job, instance) rate({COUNTER_RHS_TABLE}[5m]))[2m:1m])" @@ -640,12 +694,12 @@ async fn check_hybrid_binary_cases(instance: &Arc, observer: &NativePl "{promql}: native timestamp column" ); if index == 2 { + let expected = [0.5; 5]; assert!( - native - .rows - .iter() - .all(|row| row.value.is_some_and(|value| close_enough(value, 1.0))), - "{promql}: independent rates must have ratio 1: {:#?}", + native.rows.iter().zip(expected).all(|(row, expected)| row + .value + .is_some_and(|value| close_enough(value, expected))), + "{promql}: independently scaled RHS rates must have expected ratios: {:#?}", native.rows ); } @@ -899,7 +953,8 @@ async fn setup_fixture(instance: &Arc) { ), ) .await; - execute_sql(instance, &format!("INSERT INTO {COUNTER_RHS_TABLE} (ts, job, instance, v) VALUES ({counter_first}, '{job}', '{host}', {COUNTER_FIRST_VALUE:?}), ({counter_second}, '{job}', '{host}', {second_value:?})")).await; + let rhs_second_value = COUNTER_FIRST_VALUE + 2.0 * increase; + execute_sql(instance, &format!("INSERT INTO {COUNTER_RHS_TABLE} (ts, job, instance, v) VALUES ({counter_first}, '{job}', '{host}', {COUNTER_FIRST_VALUE:?}), ({counter_second}, '{job}', '{host}', {rhs_second_value:?})")).await; if index + 1 < COUNTER_SERIES.len() { execute_sql(instance, &format!("ADMIN FLUSH_TABLE('{COUNTER_TABLE}')")).await; execute_sql( @@ -1570,8 +1625,6 @@ async fn native_aggregate_complete_over_four_nonempty_mito_partitions() { check_case(&instance, &observer, case).await; } - check_hybrid_binary_cases(&instance, &observer).await; - // The shapes the native compiler declines, each declined for its own reason, and each checked on // the plans that really ran: none may hold a fused producer or a frontend merge, and both routes // must answer what the fixture derives by hand.