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] 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.