From 7941800c9f796e87deaeb01f9dbe02e70b3c615c Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:39:23 +0800 Subject: [PATCH] fix(promql): retain projection dependencies and dropped histogram annotations Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/optimizer/count_nest_aggr.rs | 259 +++++++++++++++++++++ src/query/src/promql/planner.rs | 57 ++++- src/query/src/promql/planner/test.rs | 13 ++ 3 files changed, 327 insertions(+), 2 deletions(-) diff --git a/src/query/src/optimizer/count_nest_aggr.rs b/src/query/src/optimizer/count_nest_aggr.rs index 89ba4260742..b355dd399b7 100644 --- a/src/query/src/optimizer/count_nest_aggr.rs +++ b/src/query/src/optimizer/count_nest_aggr.rs @@ -205,6 +205,13 @@ impl CountNestAggrRule { fn collect_required_instant_columns_into(plan: &LogicalPlan, required: &mut HashSet) { match plan { LogicalPlan::Projection(projection) => { + for expr in &projection.expr { + required.extend( + expr.column_refs() + .into_iter() + .map(|column| column.name.clone()), + ); + } Self::collect_required_instant_columns_into(projection.input.as_ref(), required); } LogicalPlan::Extension(extension) => { @@ -344,3 +351,255 @@ impl CountNestAggrRule { } } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use datafusion::functions_aggregate::sum::sum_udaf; + use datafusion::logical_expr::EmptyRelation; + use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_common::{Column, DFSchema}; + use datafusion_expr::col; + use datatypes::arrow::datatypes::{DataType, Field, TimeUnit}; + + use super::*; + + /// Unmarked source: plain arrow fields, no region metadata. It carries the + /// columns the fixtures reference plus `unused`, which nothing reads. + fn source_plan() -> LogicalPlan { + let fields = vec![ + ( + None, + Arc::new(Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )), + ), + (None, Arc::new(Field::new("tag", DataType::Utf8, true))), + ( + None, + Arc::new(Field::new("value", DataType::Float64, false)), + ), + (None, Arc::new(Field::new("aux", DataType::Float64, true))), + ( + None, + Arc::new(Field::new("unused", DataType::Float64, true)), + ), + ]; + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::new_with_metadata(fields, HashMap::new()).unwrap()), + }) + } + + /// `SeriesDivide(tag, ts) -> InstantManipulate(ts, value)`. + fn instant_chain() -> LogicalPlan { + let divide = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + vec!["tag".to_string()], + "ts".to_string(), + source_plan(), + )), + }); + LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + 0, + 1000, + 1000, + 1000, + 0, + "ts".to_string(), + vec!["tag".to_string()], + Some("value".to_string()), + divide, + )), + }) + } + + /// Builds the oracle shape `Sort <- count <- Sort <- sum <- ` and runs + /// [`CountNestAggrRule::try_rewrite_sort`] on it. + /// + /// The rule only matches bare aggregate expressions, so both aggregates are + /// built without aliases, and the outer `count` argument is a column named + /// exactly like the inner aggregate's output field. That name is read back + /// from the inner schema the same way the rule's own name check does. + fn try_rewrite(input: LogicalPlan, inner_value: Expr) -> Result> { + let inner_agg_plan = LogicalPlanBuilder::from(input) + .aggregate( + vec![col("ts"), col("tag")], + vec![sum_udaf().call(vec![inner_value])], + ) + .unwrap() + .build() + .unwrap(); + let inner_output_column = { + let LogicalPlan::Aggregate(inner_agg) = &inner_agg_plan else { + unreachable!("plan builder must produce an Aggregate node"); + }; + Column::new_unqualified( + inner_agg + .schema + .field(inner_agg.group_expr.len()) + .name() + .clone(), + ) + }; + let inner_sort = LogicalPlanBuilder::from(inner_agg_plan) + .sort(vec![col("ts").sort(true, true)]) + .unwrap() + .build() + .unwrap(); + let outer_agg = LogicalPlanBuilder::from(inner_sort) + .aggregate( + vec![col("ts")], + vec![count_udaf().call(vec![Expr::Column(inner_output_column)])], + ) + .unwrap() + .build() + .unwrap(); + let outer_sort = LogicalPlanBuilder::from(outer_agg) + .sort(vec![col("ts").sort(true, true)]) + .unwrap() + .build() + .unwrap(); + + let LogicalPlan::Sort(sort) = outer_sort else { + unreachable!("plan builder must produce a Sort node"); + }; + CountNestAggrRule::try_rewrite_sort(&sort) + } + + /// Strings of every projection expression in `plan`, in plan order. + fn projection_exprs(plan: &LogicalPlan) -> Vec { + let mut exprs = Vec::new(); + plan.apply(|node| { + if let LogicalPlan::Projection(projection) = node { + exprs.extend(projection.expr.iter().map(|expr| expr.to_string())); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + exprs + } + + /// Column names the `SeriesDivide` node reads after pruning. + fn series_divide_input_columns(plan: &LogicalPlan) -> Vec { + let mut columns = Vec::new(); + plan.apply(|node| { + if let LogicalPlan::Extension(extension) = node + && extension.node.as_any().is::() + { + columns = extension.node.inputs()[0] + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + columns + } + + /// Asserts the shared post-rewrite invariants for one fixture case: the + /// rewrite really happened, pruning below `SeriesDivide` kept `aux` (the + /// retained projection still reads it) and dropped `unused`, and every + /// `kept` expression fragment is still printed in the projection chain. + fn assert_rewritten(rewritten: &LogicalPlan, case: &str, kept: &[&str]) { + assert!( + has_distinct(rewritten), + "{case}: rewrite must occur: {rewritten}" + ); + assert_eq!( + series_divide_input_columns(rewritten), + ["ts", "tag", "value", "aux"], + "{case}: expected aux retained and unused pruned: {rewritten}" + ); + let exprs = projection_exprs(rewritten); + for &fragment in kept { + assert!( + exprs.iter().any(|expr| expr.contains(fragment)), + "{case}: expected a projection expression containing {fragment:?}: {rewritten}" + ); + } + } + + fn has_distinct(plan: &LogicalPlan) -> bool { + let mut found = false; + plan.apply(|node| { + if matches!(node, LogicalPlan::Distinct(_)) { + found = true; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + found + } + + #[test] + fn rewrite_retains_plain_aux_column_for_retained_projection() { + let projection = LogicalPlanBuilder::from(instant_chain()) + .project(vec![col("ts"), col("tag"), col("value"), col("aux")]) + .unwrap() + .build() + .unwrap(); + + let rewritten = try_rewrite(projection, col("value")) + .expect("rewrite must not fail") + .expect("rule must fire"); + assert_rewritten(&rewritten, "plain aux", &["aux"]); + } + + #[test] + fn rewrite_retains_computed_aux_alias_for_retained_projection() { + let projection = LogicalPlanBuilder::from(instant_chain()) + .project(vec![ + col("ts"), + col("tag"), + col("value"), + (col("aux") + lit(1.0)).alias("aux_plus_one"), + ]) + .unwrap() + .build() + .unwrap(); + + let rewritten = try_rewrite(projection, col("value")) + .expect("rewrite must not fail") + .expect("rule must fire"); + assert_rewritten(&rewritten, "computed aux alias", &["aux +", "aux_plus_one"]); + } + + #[test] + fn rewrite_retains_aux_through_two_alias_layers() { + let layer = LogicalPlanBuilder::from(instant_chain()) + .project(vec![ + col("ts"), + col("tag"), + col("value").alias("v1"), + col("aux").alias("aux1"), + ]) + .unwrap() + .build() + .unwrap(); + let projection = LogicalPlanBuilder::from(layer) + .project(vec![ + col("ts"), + col("tag"), + col("v1").alias("v2"), + col("aux1"), + ]) + .unwrap() + .build() + .unwrap(); + + let rewritten = try_rewrite(projection, col("v2")) + .expect("rewrite must not fail") + .expect("rule must fire"); + assert_rewritten(&rewritten, "two alias layers", &["aux AS aux1", "aux1"]); + } +} diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5867301f972..64e193718d1 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -1407,6 +1407,25 @@ impl PromPlanner { &right_field_columns, ); + // `+`/`-` over two alternative float/histogram operands records its discarded + // mixed pairs from the float lane's `coalesce` fallbacks instead of a standalone + // annotation filter. `align_binary_field_columns` emits that lane first, named + // after the left operand's float lane. + let arithmetic_float_lane_output = if matches!(op.id(), token::T_ADD | token::T_SUB) + && Self::field_columns_are_alternative_samples( + left_input.schema(), + &left_field_columns, + ) + && Self::field_columns_are_alternative_samples( + right_input.schema(), + &right_field_columns, + ) { + Self::alternative_sample_columns(left_input.schema(), &left_field_columns) + .map(|(float, _)| float.to_string()) + } else { + None + }; + // normal join if left_table_ref == right_table_ref { // rename table references to avoid ambiguity @@ -1441,7 +1460,7 @@ impl PromPlanner { // Computed scalars reach this join path instead of the literal projection paths. // Broadcast them for arithmetic in the same way as literal scalars. let broadcast_scalar = !is_comparison_op; - let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns( + let (mut field_groups, mut invalid_field_pairs) = Self::align_binary_field_columns( left_input.schema(), right_input.schema(), &left_field_columns, @@ -1472,6 +1491,37 @@ impl PromPlanner { .iter() .map(|(output, _)| output.clone()) .collect(); + // `+`/`-` over two alternative float/histogram operands records the discarded + // mixed pairs from the float lane: the valid float expression stays first and + // each fallback records its pair while returning NULL. A standalone always-true + // annotation filter cannot be used for them, because the outer non-null filter + // can be pushed into the same physical filter and short-circuit those UDFs + // before they run. + if self.promql_annotations.is_some() + && let Some(float_output) = arithmetic_float_lane_output.as_deref() + && let Some((_, float_pairs)) = field_groups + .iter_mut() + .find(|(output, _)| output.as_str() == float_output) + { + float_pairs.extend(std::mem::take(&mut invalid_field_pairs)); + } + // `%` over a mixed float/histogram vector and a native-histogram vector supports + // none of its pairs, so `align_binary_field_columns` falls back to a single lane + // that holds nothing but the discarded pairs. Record the dropped samples from + // that lane instead of a standalone annotation filter: `%` is not a comparison, + // so each lane expression is a `Float64(NULL)` drop UDF, and the outer non-null + // filter still discards those rows. + else if op.id() == token::T_MOD + && self.promql_annotations.is_some() + && !invalid_field_pairs.is_empty() + && field_groups.len() == 1 + && field_groups[0] + .1 + .iter() + .all(|pair| invalid_field_pairs.contains(pair)) + { + field_groups[0].1 = std::mem::take(&mut invalid_field_pairs); + } let mut field_groups = field_groups.into_iter(); // `vector()` uses EmptyMetric and keeps GreptimeDB's timestamp broadcast. let has_empty_metric_operand = left_is_empty_metric || right_is_empty_metric; @@ -1523,7 +1573,10 @@ impl PromPlanner { } let promql_annotations = self.promql_annotations.clone(); // These predicates always pass; they only evaluate otherwise-discarded pairs - // while collecting annotations. + // while collecting annotations. `+`/`-` over two alternative float/histogram + // operands and `%` over a mixed vector and a native-histogram vector already + // consumed their mixed pairs into their only output lane above, where the + // physical plan cannot short-circuit them away. let invalid_pair_predicates = invalid_field_pairs .into_iter() .filter(|_| promql_annotations.is_some()) diff --git a/src/query/src/promql/planner/test.rs b/src/query/src/promql/planner/test.rs index 19de1098288..3f9ebb3586c 100644 --- a/src/query/src/promql/planner/test.rs +++ b/src/query/src/promql/planner/test.rs @@ -7667,9 +7667,22 @@ async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { #[tokio::test] async fn test_mixed_binary_operator_reports_only_dropped_samples() { for (query, expected_rows, expected_infos) in [ + // `+`/`-` report the discarded mixed pairs from the float lane's coalesce fallbacks. ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), + ("(lf or on(tag) lh) - on(tag) (rf or on(tag) rh)", 0, 1), + // Compatible float/float and histogram/histogram pairs must not invent an annotation. ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), + ("(lf or on(tag) lh) - on(tag) (lf or on(tag) lh)", 2, 0), + // An unmatched operand has no dropped pair to report. + ( + "(lf or on(tag) lh) + on(tag) (fallback or on(tag) bad_native)", + 0, + 0, + ), ("(lf or on(tag) lh) % on(tag) lh", 0, 1), + // A non-bool comparison drops its unsupported pairs through the standalone annotation + // filter, not through the lane, so the dropped samples are still reported. + ("(lf or on(tag) lh) > on(tag) lh", 0, 1), ] { let state = build_query_engine_state(); let annotations = PromqlAnnotationCollector::default();