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

@@ -205,6 +205,13 @@ impl CountNestAggrRule {
     fn collect_required_instant_columns_into(plan: &LogicalPlan, required: &mut HashSet<String>) {
         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 <- <input>` 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<Option<LogicalPlan>> {
+        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<String> {
+        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<String> {
+        let mut columns = Vec::new();
+        plan.apply(|node| {
+            if let LogicalPlan::Extension(extension) = node
+                && extension.node.as_any().is::<SeriesDivide>()
+            {
+                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"]);
+    }
+}

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())

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();