From c4dafb58269ff15b1a1ad50affec4458c83f87cf Mon Sep 17 00:00:00 2001
From: shuiyisong <113876041+shuiyisong@users.noreply.github.com>
Date: Tue, 1 Sep 2026 13:03:00 +0000
Subject: [PATCH] fix(promql): resolve derived labels in aggregation arithmetic
 (#8994)

Signed-off-by: shuiyisong <xixing.sys@gmail.com>
---
 src/query/src/promql/planner.rs               | 41 ++++++++++++++++++-
 .../standalone/common/promql/label.result     | 17 ++++++++
 .../cases/standalone/common/promql/label.sql  |  4 ++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs
index de0181af7f..770f5e5ee2 100644
--- a/src/query/src/promql/planner.rs
+++ b/src/query/src/promql/planner.rs
@@ -6611,12 +6611,20 @@ impl PromPlanner {
         let preserve_field_names =
             Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns);
         let table_ref = self.ctx.table_name.clone().map(TableReference::bare);
+        // Derived labels can be unqualified even when the context still names the source table.
+        let input_schema = input.schema().clone();
         let non_field_columns_iter = self
             .ctx
             .tag_columns
             .iter()
             .chain(self.ctx.time_index_column.iter())
-            .map(|col| Ok(DfExpr::Column(Column::new(table_ref.clone(), col))));
+            .map(|col| {
+                input_schema
+                    .qualified_field_with_name(table_ref.as_ref(), col)
+                    .or_else(|_| input_schema.qualified_field_with_unqualified_name(col))
+                    .map(|field| DfExpr::Column(field.into()))
+                    .context(DataFusionPlanningSnafu)
+            });
         let tsid_iter =
             Self::optional_tsid_projection(input.schema(), table_ref.as_ref(), self.ctx.use_tsid)
                 .into_iter()
@@ -11973,6 +11981,37 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo:
         assert_eq!(format!("\n{ret}"), expected, "\n{}", ret);
     }
 
+    #[tokio::test]
+    async fn label_replace_aggregation_queries_plan_successfully() {
+        let aggregate =
+            r#"sum by (foo) (label_replace(some_metric, "foo", "$1", "tag_0", "(.*)"))"#;
+        let queries = [
+            aggregate.to_string(),
+            format!("{aggregate} <= 10"),
+            format!("{aggregate} * 0.8"),
+            format!("0.8 * {aggregate}"),
+            format!("{aggregate} <= {aggregate} * 0.8"),
+        ];
+        let state = build_query_engine_state();
+        let mut failures = Vec::new();
+
+        for query in queries {
+            let table_provider = build_test_table_provider(
+                &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())],
+                1,
+                1,
+            )
+            .await;
+            if let Err(error) =
+                PromPlanner::stmt_to_plan(table_provider, &build_eval_stmt(&query), &state).await
+            {
+                failures.push(format!("{query}: {error:?}"));
+            }
+        }
+
+        assert!(failures.is_empty(), "{}", failures.join("\n"));
+    }
+
     #[tokio::test]
     async fn test_matchers_to_expr() {
         let mut eval_stmt = EvalStmt {
diff --git a/tests/cases/standalone/common/promql/label.result b/tests/cases/standalone/common/promql/label.result
index f83bbd38e6..bd9228a324 100644
--- a/tests/cases/standalone/common/promql/label.result
+++ b/tests/cases/standalone/common/promql/label.result
@@ -314,6 +314,23 @@ TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-", "idc", "h
 | 1970-01-01T00:00:15 | 3   | idc2:zone1-host1 | host1 | idc2:zone1 |
 +---------------------+-----+------------------+-------+------------+
 
+-- Issue 8969 --
+-- SQLNESS SORT_RESULT 3 1
+TQL EVAL (0, 15, '5s') sum by (foo) (label_replace(test, "foo", "$1", "host", "(.*)")) * 0.8;
+
++-------+---------------------+------------------------------+
+| foo   | ts                  | sum(test.val) * Float64(0.8) |
++-------+---------------------+------------------------------+
+| host1 | 1970-01-01T00:00:00 | 0.8                          |
+| host1 | 1970-01-01T00:00:05 | 3.2                          |
+| host1 | 1970-01-01T00:00:10 | 7.2                          |
+| host1 | 1970-01-01T00:00:15 | 12.8                         |
+| host2 | 1970-01-01T00:00:00 | 1.6                          |
+| host2 | 1970-01-01T00:00:05 | 4.800000000000001            |
+| host2 | 1970-01-01T00:00:10 | 9.600000000000001            |
+| host2 | 1970-01-01T00:00:15 | 16.0                         |
++-------+---------------------+------------------------------+
+
 DROP TABLE test;
 
 Affected Rows: 0
diff --git a/tests/cases/standalone/common/promql/label.sql b/tests/cases/standalone/common/promql/label.sql
index 3fb20792e0..6a77e9b521 100644
--- a/tests/cases/standalone/common/promql/label.sql
+++ b/tests/cases/standalone/common/promql/label.sql
@@ -96,6 +96,10 @@ TQL EVAL (0, 15, '5s') label_replace(test{host="host1"}, "new_idc", "idc99", "id
 -- SQLNESS SORT_RESULT 3 1
 TQL EVAL (0, 15, '5s') label_join(test{host="host1"}, "new_host", "-", "idc", "host") == 3;
 
+-- Issue 8969 --
+-- SQLNESS SORT_RESULT 3 1
+TQL EVAL (0, 15, '5s') sum by (foo) (label_replace(test, "foo", "$1", "host", "(.*)")) * 0.8;
+
 DROP TABLE test;
 
 CREATE TABLE test (