Repository: GreptimeTeam/greptimedb
Pull request #5652: feat(promql): supports quantile and count_values
Author: login: killme2008
Labels: name: docs-not-required
description: This change does not impact docs.
Created: 2025-03-05T16:02:58Z | Merged: 2025-03-10T07:02:53Z

I hereby agree to the terms of the [GreptimeDB CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md).

## Refer to a related PR or issue link (optional)

#5525 

## What's changed and what's your intention?

<!--    
 __!!! DO NOT LEAVE THIS BLOCK EMPTY !!!__

Please explain IN DETAIL what the changes are in this PR and why they are needed:

- Summarize your change (**mandatory**)
- How does this PR work? Need a brief introduction for the changed logic (optional)
- Describe clearly one logical change and avoid lazy messages (optional)
- Describe any limitations of the current code (optional)
- Describe if this PR will break **API or data compatibility**  (optional)
-->

Impl `quantile` and `count_values` operator, for example:

```
quantile(0.8, http_total_requests)
count_values('status_code', http_total_requests)
```

## PR Checklist
Please convert it to a draft if some of the following conditions are not met.

- [x] I have written the necessary rustdoc comments.
- [x] I have added the necessary unit tests and integration tests.
- [x] This PR requires documentation updates.
- [x] API changes are backward compatible.
- [x] Schema or data changes are backward compatible.



Diff:
diff --git a/Cargo.lock b/Cargo.lock
index cea12295428e..639639e72d2c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8760,6 +8760,7 @@ dependencies = [
  "common-recordbatch",
  "common-telemetry",
  "datafusion",
+ "datafusion-common",
  "datafusion-expr",
  "datatypes",
  "futures",
diff --git a/src/promql/Cargo.toml b/src/promql/Cargo.toml
index 0970b0d38e93..f93cb8beb99b 100644
--- a/src/promql/Cargo.toml
+++ b/src/promql/Cargo.toml
@@ -16,6 +16,7 @@ common-macro.workspace = true
 common-recordbatch.workspace = true
 common-telemetry.workspace = true
 datafusion.workspace = true
+datafusion-common.workspace = true
 datafusion-expr.workspace = true
 datatypes.workspace = true
 futures.workspace = true
diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs
index 4209a517c617..dade00ea7bbf 100644
--- a/src/promql/src/functions.rs
+++ b/src/promql/src/functions.rs
@@ -20,6 +20,7 @@ mod holt_winters;
 mod idelta;
 mod predict_linear;
 mod quantile;
+mod quantile_aggr;
 mod resets;
 mod round;
 #[cfg(test)]
@@ -39,6 +40,7 @@ pub use holt_winters::HoltWinters;
 pub use idelta::IDelta;
 pub use predict_linear::PredictLinear;
 pub use quantile::QuantileOverTime;
+pub use quantile_aggr::quantile_udaf;
 pub use resets::Resets;
 pub use round::Round;
 
diff --git a/src/promql/src/functions/quantile.rs b/src/promql/src/functions/quantile.rs
index 19f346ee182a..4900ab445370 100644
--- a/src/promql/src/functions/quantile.rs
+++ b/src/promql/src/functions/quantile.rs
@@ -125,7 +125,7 @@ impl QuantileOverTime {
 }
 
 /// Refer to <https://github.com/prometheus/prometheus/blob/6e2905a4d4ff9b47b1f6d201333f5bd53633f921/promql/quantile.go#L357-L386>
-fn quantile_impl(values: &[f64], quantile: f64) -> Option<f64> {
+pub(crate) fn quantile_impl(values: &[f64], quantile: f64) -> Option<f64> {
     if quantile.is_nan() || values.is_empty() {
         return Some(f64::NAN);
     }
diff --git a/src/promql/src/functions/quantile_aggr.rs b/src/promql/src/functions/quantile_aggr.rs
new file mode 100644
index 000000000000..e19cfeb21afd
--- /dev/null
+++ b/src/promql/src/functions/quantile_aggr.rs
@@ -0,0 +1,297 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, AsArray};
+use datafusion::common::cast::{as_list_array, as_primitive_array, as_struct_array};
+use datafusion::error::Result as DfResult;
+use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility};
+use datafusion::prelude::create_udaf;
+use datafusion_common::ScalarValue;
+use datatypes::arrow::array::{ListArray, StructArray};
+use datatypes::arrow::datatypes::{DataType, Field, Float64Type};
+
+use crate::functions::quantile::quantile_impl;
+
+const QUANTILE_NAME: &str = "quantile";
+
+const VALUES_FIELD_NAME: &str = "values";
+const DEFAULT_LIST_FIELD_NAME: &str = "item";
+
+#[derive(Debug, Default)]
+pub struct QuantileAccumulator {
+    q: f64,
+    values: Vec<Option<f64>>,
+}
+
+/// Create a quantile `AggregateUDF` for PromQL quantile operator,
+/// which calculates φ-quantile (0 ≤ φ ≤ 1) over dimensions
+pub fn quantile_udaf(q: f64) -> Arc<AggregateUDF> {
+    Arc::new(create_udaf(
+        QUANTILE_NAME,
+        // Input type: (values)
+        vec![DataType::Float64],
+        // Output type: the φ-quantile
+        Arc::new(DataType::Float64),
+        Volatility::Immutable,
+        // Create the accumulator
+        Arc::new(move |_| Ok(Box::new(QuantileAccumulator::new(q)))),
+        // Intermediate state types
+        Arc::new(vec![DataType::Struct(
+            vec![Field::new(
+                VALUES_FIELD_NAME,
+                DataType::List(Arc::new(Field::new(
+                    DEFAULT_LIST_FIELD_NAME,
+                    DataType::Float64,
+                    true,
+                ))),
+                false,
+            )]
+            .into(),
+        )]),
+    ))
+}
+
+impl QuantileAccumulator {
+    pub fn new(q: f64) -> Self {
+        Self {
+            q,
+            ..Default::default()
+        }
+    }
+}
+
+impl DfAccumulator for QuantileAccumulator {
+    fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> {
+        let f64_array = values[0].as_primitive::<Float64Type>();
+
+        self.values.extend(f64_array);
+
+        Ok(())
+    }
+
+    fn evaluate(&mut self) -> DfResult<ScalarValue> {
+        let values: Vec<_> = self.values.iter().map(|v| v.unwrap_or(0.0)).collect();
+
+        let result = quantile_impl(&values, self.q);
+
+        ScalarValue::new_primitive::<Float64Type>(result, &DataType::Float64)
+    }
+
+    fn size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.values.capacity() * std::mem::size_of::<Option<f64>>()
+    }
+
+    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
+        let values_array = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
+            Some(self.values.clone()),
+        ]));
+
+        let state_struct = StructArray::new(
+            vec![Field::new(
+                VALUES_FIELD_NAME,
+                DataType::List(Arc::new(Field::new(
+                    DEFAULT_LIST_FIELD_NAME,
+                    DataType::Float64,
+                    true,
+                ))),
+                false,
+            )]
+            .into(),
+            vec![values_array],
+            None,
+        );
+
+        Ok(vec![ScalarValue::Struct(Arc::new(state_struct))])
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        for state in states {
+            let state = as_struct_array(state)?;
+
+            for list in as_list_array(state.column(0))?.iter().flatten() {
+                let f64_array = as_primitive_array::<Float64Type>(&list)?.clone();
+                self.values.extend(&f64_array);
+            }
+        }
+
+        Ok(())
+    }
+}
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use datafusion::arrow::array::{ArrayRef, Float64Array};
+    use datafusion_common::ScalarValue;
+
+    use super::*;
+
+    fn create_f64_array(values: Vec<Option<f64>>) -> ArrayRef {
+        Arc::new(Float64Array::from(values)) as ArrayRef
+    }
+
+    #[test]
+    fn test_quantile_accumulator_empty() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+
+        let result = accumulator.evaluate().unwrap();
+
+        match result {
+            ScalarValue::Float64(_) => (),
+            _ => panic!("Expected Float64 scalar value"),
+        }
+    }
+
+    #[test]
+    fn test_quantile_accumulator_single_value() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(10.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let result = accumulator.evaluate().unwrap();
+
+        assert_eq!(result, ScalarValue::Float64(Some(10.0)));
+    }
+
+    #[test]
+    fn test_quantile_accumulator_multiple_values() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0), Some(5.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let result = accumulator.evaluate().unwrap();
+
+        assert_eq!(result, ScalarValue::Float64(Some(3.0)));
+    }
+
+    #[test]
+    fn test_quantile_accumulator_with_nulls() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(1.0), None, Some(3.0), Some(4.0), Some(5.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+
+        let result = accumulator.evaluate().unwrap();
+        assert_eq!(result, ScalarValue::Float64(Some(3.0)));
+    }
+
+    #[test]
+    fn test_quantile_accumulator_multiple_batches() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input1 = create_f64_array(vec![Some(1.0), Some(2.0)]);
+        let input2 = create_f64_array(vec![Some(3.0), Some(4.0), Some(5.0)]);
+
+        accumulator.update_batch(&[input1]).unwrap();
+        accumulator.update_batch(&[input2]).unwrap();
+
+        let result = accumulator.evaluate().unwrap();
+        assert_eq!(result, ScalarValue::Float64(Some(3.0)));
+    }
+
+    #[test]
+    fn test_quantile_accumulator_different_quantiles() {
+        let mut min_accumulator = QuantileAccumulator::new(0.0);
+        let input = create_f64_array(vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0), Some(5.0)]);
+        min_accumulator.update_batch(&[input.clone()]).unwrap();
+        assert_eq!(
+            min_accumulator.evaluate().unwrap(),
+            ScalarValue::Float64(Some(1.0))
+        );
+
+        let mut q1_accumulator = QuantileAccumulator::new(0.25);
+        q1_accumulator.update_batch(&[input.clone()]).unwrap();
+        assert_eq!(
+            q1_accumulator.evaluate().unwrap(),
+            ScalarValue::Float64(Some(2.0))
+        );
+
+        let mut q3_accumulator = QuantileAccumulator::new(0.75);
+        q3_accumulator.update_batch(&[input.clone()]).unwrap();
+        assert_eq!(
+            q3_accumulator.evaluate().unwrap(),
+            ScalarValue::Float64(Some(4.0))
+        );
+
+        let mut max_accumulator = QuantileAccumulator::new(1.0);
+        max_accumulator.update_batch(&[input]).unwrap();
+        assert_eq!(
+            max_accumulator.evaluate().unwrap(),
+            ScalarValue::Float64(Some(5.0))
+        );
+    }
+
+    #[test]
+    fn test_quantile_accumulator_size() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(1.0), Some(2.0), Some(3.0)]);
+
+        let initial_size = accumulator.size();
+        accumulator.update_batch(&[input]).unwrap();
+        let after_update_size = accumulator.size();
+
+        assert!(after_update_size >= initial_size);
+    }
+
+    #[test]
+    fn test_quantile_accumulator_state_and_merge() -> DfResult<()> {
+        let mut acc1 = QuantileAccumulator::new(0.5);
+        let input1 = create_f64_array(vec![Some(1.0), Some(2.0)]);
+        acc1.update_batch(&[input1])?;
+
+        let state1 = acc1.state()?;
+
+        let mut acc2 = QuantileAccumulator::new(0.5);
+        let input2 = create_f64_array(vec![Some(3.0), Some(4.0), Some(5.0)]);
+        acc2.update_batch(&[input2])?;
+
+        let mut struct_builders = vec![];
+        for scalar in &state1 {
+            if let ScalarValue::Struct(struct_array) = scalar {
+                struct_builders.push(struct_array.clone() as ArrayRef);
+            }
+        }
+
+        acc2.merge_batch(&struct_builders)?;
+
+        let result = acc2.evaluate()?;
+
+        assert_eq!(result, ScalarValue::Float64(Some(3.0)));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_quantile_accumulator_with_extreme_values() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(f64::MAX), Some(f64::MIN), Some(0.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let _result = accumulator.evaluate().unwrap();
+    }
+
+    #[test]
+    fn test_quantile_udaf_creation() {
+        let q = 0.5;
+        let udaf = quantile_udaf(q);
+
+        assert_eq!(udaf.name(), QUANTILE_NAME);
+        assert_eq!(udaf.return_type(&[]).unwrap(), DataType::Float64);
+    }
+}
diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs
index 22477d5049de..d86fd1c769a9 100644
--- a/src/query/src/promql/planner.rs
+++ b/src/query/src/promql/planner.rs
@@ -51,8 +51,8 @@ use promql::extension_plan::{
     RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn,
 };
 use promql::functions::{
-    AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, HoltWinters, IDelta,
-    Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime,
+    quantile_udaf, AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, HoltWinters,
+    IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime,
     QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime,
 };
 use promql_parser::label::{MatchOp, Matcher, Matchers, METRIC_NAME};
@@ -266,7 +266,10 @@ impl PromPlanner {
         aggr_expr: &AggregateExpr,
     ) -> Result<LogicalPlan> {
         let AggregateExpr {
-            op, expr, modifier, ..
+            op,
+            expr,
+            modifier,
+            param,
         } = aggr_expr;
 
         let input = self.prom_expr_to_plan(expr, session_state).await?;
@@ -277,19 +280,40 @@ impl PromPlanner {
             _ => {
                 // calculate columns to group by
                 // Need to append time index column into group by columns
-                let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
+                let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?;
                 // convert op and value columns to aggregate exprs
-                let aggr_exprs = self.create_aggregate_exprs(*op, &input)?;
+                let (aggr_exprs, prev_field_exprs) =
+                    self.create_aggregate_exprs(*op, param, &input)?;
 
                 // create plan
-                let group_sort_expr = group_exprs
-                    .clone()
-                    .into_iter()
-                    .map(|expr| expr.sort(true, false));
-                LogicalPlanBuilder::from(input)
-                    .aggregate(group_exprs.clone(), aggr_exprs)
-                    .context(DataFusionPlanningSnafu)?
-                    .sort(group_sort_expr)
+                let builder = LogicalPlanBuilder::from(input);
+                let builder = if op.id() == token::T_COUNT_VALUES {
+                    let label = Self::get_param_value_as_str(*op, param)?;
+                    // `count_values` must be grouped by fields,
+                    // and project the fields to the new label.
+                    group_exprs.extend(prev_field_exprs.clone());
+                    let project_fields = self
+                        .create_field_column_exprs()?
+                        .into_iter()
+                        .chain(self.create_tag_column_exprs()?)
+                        .chain(Some(self.create_time_index_column_expr()?))
+                        .chain(prev_field_exprs.into_iter().map(|expr| expr.alias(label)));
+
+                    builder
+                        .aggregate(group_exprs.clone(), aggr_exprs)
+                        .context(DataFusionPlanningSnafu)?
+                        .project(project_fields)
+                        .context(DataFusionPlanningSnafu)?
+                } else {
+                    builder
+                        .aggregate(group_exprs.clone(), aggr_exprs)
+                        .context(DataFusionPlanningSnafu)?
+                };
+
+                let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false));
+
+                builder
+                    .sort(sort_expr)
                     .context(DataFusionPlanningSnafu)?
                     .build()
                     .context(DataFusionPlanningSnafu)
@@ -312,18 +336,7 @@ impl PromPlanner {
 
         let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?;
 
-        let param = param
-            .as_deref()
-            .with_context(|| FunctionInvalidArgumentSnafu {
-                fn_name: (*op).to_string(),
-            })?;
-
-        let PromExpr::NumberLiteral(NumberLiteral { val }) = param else {
-            return FunctionInvalidArgumentSnafu {
-                fn_name: (*op).to_string(),
-            }
-            .fail();
-        };
+        let val = Self::get_param_value_as_f64(*op, param)?;
 
         // convert op and value columns to window exprs.
         let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?;
@@ -341,7 +354,7 @@ impl PromPlanner {
                 let predicate = DfExpr::BinaryExpr(BinaryExpr {
                     left: Box::new(col(rank)),
                     op: Operator::LtEq,
-                    right: Box::new(lit(*val)),
+                    right: Box::new(lit(val)),
                 });
 
                 match expr {
@@ -1931,32 +1944,44 @@ impl PromPlanner {
         })
     }
 
-    /// Create [DfExpr::AggregateFunction] expr for each value column with given aggregate function.
+    /// Creates a set of DataFusion `DfExpr::AggregateFunction` expressions for each value column using the specified aggregate function.
     ///
-    /// # Side effect
+    /// # Side Effects
+    ///
+    /// This method modifies the value columns in the context by replacing them with the new columns
+    /// created by the aggregate function application.
+    ///
+    /// # Returns
+    ///
+    /// Returns a tuple of `(aggregate_expressions, previous_field_expressions)` where:
+    /// - `aggregate_expressions`: Expressions that apply the aggregate function to the original fields
+    /// - `previous_field_expressions`: Original field expressions before aggregation. This is non-empty
+    ///   only when the operation is `count_values`, as this operation requires preserving the original
+    ///   values for grouping.
     ///
-    /// This method will update value columns in context to the new value columns created by
-    /// aggregate function.
     fn create_aggregate_exprs(
         &mut self,
         op: TokenType,
+        param: &Option<Box<PromExpr>>,
         input_plan: &LogicalPlan,
-    ) -> Result<Vec<DfExpr>> {
+    ) -> Result<(Vec<DfExpr>, Vec<DfExpr>)> {
         let aggr = match op.id() {
             token::T_SUM => sum_udaf(),
+            token::T_QUANTILE => {
+                let q = Self::get_param_value_as_f64(op, param)?;
+                quantile_udaf(q)
+            }
             token::T_AVG => avg_udaf(),
-            token::T_COUNT => count_udaf(),
+            token::T_COUNT_VALUES | token::T_COUNT => count_udaf(),
             token::T_MIN => min_udaf(),
             token::T_MAX => max_udaf(),
             token::T_GROUP => grouping_udaf(),
             token::T_STDDEV => stddev_pop_udaf(),
             token::T_STDVAR => var_pop_udaf(),
-            token::T_TOPK | token::T_BOTTOMK | token::T_COUNT_VALUES | token::T_QUANTILE => {
-                UnsupportedExprSnafu {
-                    name: format!("{op:?}"),
-                }
-                .fail()?
+            token::T_TOPK | token::T_BOTTOMK => UnsupportedExprSnafu {
+                name: format!("{op:?}"),
             }
+            .fail()?,
             _ => UnexpectedTokenSnafu { token: op }.fail()?,
         };
 
@@ -1966,19 +1991,41 @@ impl PromPlanner {
             .field_columns
             .iter()
             .map(|col| {
-                DfExpr::AggregateFunction(AggregateFunction {
+                Ok(DfExpr::AggregateFunction(AggregateFunction {
                     func: aggr.clone(),
                     args: vec![DfExpr::Column(Column::from_name(col))],
                     distinct: false,
                     filter: None,
                     order_by: None,
                     null_treatment: None,
-                })
+                }))
             })
-            .collect();
+            .collect::<Result<Vec<_>>>()?;
+
+        // if the aggregator is `count_values`, it must be grouped by current fields.
+        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
+            let prev_field_exprs: Vec<_> = self
+                .ctx
+                .field_columns
+                .iter()
+                .map(|col| DfExpr::Column(Column::from_name(col)))
+                .collect();
+
+            ensure!(
+                self.ctx.field_columns.len() == 1,
+                UnsupportedExprSnafu {
+                    name: "count_values on multi-value input"
+                }
+            );
+
+            prev_field_exprs
+        } else {
+            vec![]
+        };
 
-        // update value column name according to the aggregators
+        // update value column name according to the aggregators,
         let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
+
         let normalized_exprs =
             normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
         for expr in normalized_exprs {
@@ -1986,7 +2033,39 @@ impl PromPlanner {
         }
         self.ctx.field_columns = new_field_columns;
 
-        Ok(exprs)
+        Ok((exprs, prev_field_exprs))
+    }
+
+    fn get_param_value_as_str(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<&str> {
+        let param = param
+            .as_deref()
+            .with_context(|| FunctionInvalidArgumentSnafu {
+                fn_name: op.to_string(),
+            })?;
+        let PromExpr::StringLiteral(StringLiteral { val }) = param else {
+            return FunctionInvalidArgumentSnafu {
+                fn_name: op.to_string(),
+            }
+            .fail();
+        };
+
+        Ok(val)
+    }
+
+    fn get_param_value_as_f64(op: TokenType, param: &Option<Box<PromExpr>>) -> Result<f64> {
+        let param = param
+            .as_deref()
+            .with_context(|| FunctionInvalidArgumentSnafu {
+                fn_name: op.to_string(),
+            })?;
+        let PromExpr::NumberLiteral(NumberLiteral { val }) = param else {
+            return FunctionInvalidArgumentSnafu {
+                fn_name: op.to_string(),
+            }
+            .fail();
+        };
+
+        Ok(*val)
     }
 
     /// Create [DfExpr::WindowFunction] expr for each value column with given window function.
@@ -3342,30 +3421,6 @@ mod test {
         do_aggregate_expr_plan("stdvar", "var_pop").await;
     }
 
-    #[tokio::test]
-    #[should_panic]
-    async fn aggregate_top_k() {
-        do_aggregate_expr_plan("topk", "").await;
-    }
-
-    #[tokio::test]
-    #[should_panic]
-    async fn aggregate_bottom_k() {
-        do_aggregate_expr_plan("bottomk", "").await;
-    }
-
-    #[tokio::test]
-    #[should_panic]
-    async fn aggregate_count_values() {
-        do_aggregate_expr_plan("count_values", "").await;
-    }
-
-    #[tokio::test]
-    #[should_panic]
-    async fn aggregate_quantile() {
-        do_aggregate_expr_plan("quantile", "").await;
-    }
-
     // TODO(ruihang): add range fn tests once exprs are ready.
 
     // {
@@ -4248,4 +4303,98 @@ mod test {
 
         assert_eq!(plan.display_indent_schema().to_string(), expected);
     }
+
+    #[tokio::test]
+    async fn test_count_values_expr() {
+        let mut eval_stmt = EvalStmt {
+            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
+            start: UNIX_EPOCH,
+            end: UNIX_EPOCH
+                .checked_add(Duration::from_secs(100_000))
+                .unwrap(),
+            interval: Duration::from_secs(5),
+            lookback_delta: Duration::from_secs(1),
+        };
+        let case = r#"count_values('series', prometheus_tsdb_head_series{ip=~"(10\\.0\\.160\\.237:8080|10\\.0\\.160\\.237:9090)"}) by (ip)"#;
+
+        let prom_expr = parser::parse(case).unwrap();
+        eval_stmt.expr = prom_expr;
+        let table_provider = build_test_table_provider_with_fields(
+            &[
+                (
+                    DEFAULT_SCHEMA_NAME.to_string(),
+                    "prometheus_tsdb_head_series".to_string(),
+                ),
+                (
+                    DEFAULT_SCHEMA_NAME.to_string(),
+                    "http_server_requests_seconds_count".to_string(),
+                ),
+            ],
+            &["ip"],
+        )
+        .await;
+
+        let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_session_state())
+            .await
+            .unwrap();
+        let expected = r#"Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), series:Float64;N]
+  Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), series:Float64;N, greptime_value:Float64;N]
+    Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), series:Float64;N, greptime_value:Float64;N]
+      Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]
+        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+          PromSeriesNormalize: offset=[0], time index=[greptime_timestamp], filter NaN: [false] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+            PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+              Sort: prometheus_tsdb_head_series.ip DESC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp DESC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+                Filter: prometheus_tsdb_head_series.ip ~ Utf8("(10\.0\.160\.237:8080|10\.0\.160\.237:9090)") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-1000, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100001000, None) [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+                  TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]"#;
+
+        assert_eq!(plan.display_indent_schema().to_string(), expected);
+    }
+
+    #[tokio::test]
+    async fn test_quantile_expr() {
+        let mut eval_stmt = EvalStmt {
+            expr: PromExpr::NumberLiteral(NumberLiteral { val: 1.0 }),
+            start: UNIX_EPOCH,
+            end: UNIX_EPOCH
+                .checked_add(Duration::from_secs(100_000))
+                .unwrap(),
+            interval: Duration::from_secs(5),
+            lookback_delta: Duration::from_secs(1),
+        };
+        let case = r#"quantile(0.3, sum(prometheus_tsdb_head_series{ip=~"(10\\.0\\.160\\.237:8080|10\\.0\\.160\\.237:9090)"}) by (ip))"#;
+
+        let prom_expr = parser::parse(case).unwrap();
+        eval_stmt.expr = prom_expr;
+        let table_provider = build_test_table_provider_with_fields(
+            &[
+                (
+                    DEFAULT_SCHEMA_NAME.to_string(),
+                    "prometheus_tsdb_head_series".to_string(),
+                ),
+                (
+                    DEFAULT_SCHEMA_NAME.to_string(),
+                    "http_server_requests_seconds_count".to_string(),
+                ),
+            ],
+            &["ip"],
+        )
+        .await;
+
+        let plan = PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_session_state())
+            .await
+            .unwrap();
+        let expected = r#"Sort: prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [greptime_timestamp:Timestamp(Millisecond, None), quantile(sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]
+  Aggregate: groupBy=[[prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[quantile(sum(prometheus_tsdb_head_series.greptime_value))]] [greptime_timestamp:Timestamp(Millisecond, None), quantile(sum(prometheus_tsdb_head_series.greptime_value)):Float64;N]
+    Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]
+      Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp]], aggr=[[sum(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), sum(prometheus_tsdb_head_series.greptime_value):Float64;N]
+        PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+          PromSeriesNormalize: offset=[0], time index=[greptime_timestamp], filter NaN: [false] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+            PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+              Sort: prometheus_tsdb_head_series.ip DESC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp DESC NULLS LAST [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+                Filter: prometheus_tsdb_head_series.ip ~ Utf8("(10\.0\.160\.237:8080|10\.0\.160\.237:9090)") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-1000, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100001000, None) [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]
+                  TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(Millisecond, None), greptime_value:Float64;N]"#;
+
+        assert_eq!(plan.display_indent_schema().to_string(), expected);
+    }
 }
diff --git a/tests/cases/standalone/common/promql/count_values.result b/tests/cases/standalone/common/promql/count_values.result
new file mode 100644
index 000000000000..e2f65eae342b
--- /dev/null
+++ b/tests/cases/standalone/common/promql/count_values.result
@@ -0,0 +1,67 @@
+CREATE TABLE http_requests (
+  ts timestamp(3) time index,
+  host STRING,
+  idc STRING,
+  val BIGINT,
+  PRIMARY KEY(host, idc),
+);
+
+Affected Rows: 0
+
+INSERT INTO TABLE http_requests VALUES
+    (0,     'host1', "idc1", 200),
+    (0,     'host2', "idc1", 200),
+    (0,     'host3', "idc2", 200),
+    (0,     'host4', "idc2", 401),
+    (5000,  'host1', "idc1", 404),
+    (5000,  'host2', "idc1", 401),
+    (5000,  'host3', "idc2", 404),
+    (5000,  'host4', "idc2", 500),
+    (10000, 'host1', "idc1", 200),
+    (10000, 'host2', "idc1", 200),
+    (10000, 'host3', "idc2", 201),
+    (10000, 'host4', "idc2", 201),
+    (15000, 'host1', "idc1", 500),
+    (15000, 'host2', "idc1", 500),
+    (15000, 'host3', "idc2", 500),
+    (15000, 'host4', "idc2", 500);
+
+Affected Rows: 16
+
+TQL EVAL (0, 15, '5s') count_values("status_code", http_requests);
+
++--------------------------+---------------------+-------------+
+| count(http_requests.val) | ts                  | status_code |
++--------------------------+---------------------+-------------+
+| 3                        | 1970-01-01T00:00:00 | 200         |
+| 1                        | 1970-01-01T00:00:00 | 401         |
+| 1                        | 1970-01-01T00:00:05 | 401         |
+| 2                        | 1970-01-01T00:00:05 | 404         |
+| 1                        | 1970-01-01T00:00:05 | 500         |
+| 2                        | 1970-01-01T00:00:10 | 200         |
+| 2                        | 1970-01-01T00:00:10 | 201         |
+| 4                        | 1970-01-01T00:00:15 | 500         |
++--------------------------+---------------------+-------------+
+
+TQL EVAL (0, 15, '5s') count_values("status_code", http_requests) by (idc);
+
++--------------------------+------+---------------------+-------------+
+| count(http_requests.val) | idc  | ts                  | status_code |
++--------------------------+------+---------------------+-------------+
+| 2                        | idc1 | 1970-01-01T00:00:00 | 200         |
+| 1                        | idc1 | 1970-01-01T00:00:05 | 401         |
+| 1                        | idc1 | 1970-01-01T00:00:05 | 404         |
+| 2                        | idc1 | 1970-01-01T00:00:10 | 200         |
+| 2                        | idc1 | 1970-01-01T00:00:15 | 500         |
+| 1                        | idc2 | 1970-01-01T00:00:00 | 200         |
+| 1                        | idc2 | 1970-01-01T00:00:00 | 401         |
+| 1                        | idc2 | 1970-01-01T00:00:05 | 404         |
+| 1                        | idc2 | 1970-01-01T00:00:05 | 500         |
+| 2                        | idc2 | 1970-01-01T00:00:10 | 201         |
+| 2                        | idc2 | 1970-01-01T00:00:15 | 500         |
++--------------------------+------+---------------------+-------------+
+
+DROP TABLE http_requests;
+
+Affected Rows: 0
+
diff --git a/tests/cases/standalone/common/promql/count_values.sql b/tests/cases/standalone/common/promql/count_values.sql
new file mode 100644
index 000000000000..09c2b136b83f
--- /dev/null
+++ b/tests/cases/standalone/common/promql/count_values.sql
@@ -0,0 +1,31 @@
+CREATE TABLE http_requests (
+  ts timestamp(3) time index,
+  host STRING,
+  idc STRING,
+  val BIGINT,
+  PRIMARY KEY(host, idc),
+);
+
+INSERT INTO TABLE http_requests VALUES
+    (0,     'host1', "idc1", 200),
+    (0,     'host2', "idc1", 200),
+    (0,     'host3', "idc2", 200),
+    (0,     'host4', "idc2", 401),
+    (5000,  'host1', "idc1", 404),
+    (5000,  'host2', "idc1", 401),
+    (5000,  'host3', "idc2", 404),
+    (5000,  'host4', "idc2", 500),
+    (10000, 'host1', "idc1", 200),
+    (10000, 'host2', "idc1", 200),
+    (10000, 'host3', "idc2", 201),
+    (10000, 'host4', "idc2", 201),
+    (15000, 'host1', "idc1", 500),
+    (15000, 'host2', "idc1", 500),
+    (15000, 'host3', "idc2", 500),
+    (15000, 'host4', "idc2", 500);
+
+TQL EVAL (0, 15, '5s') count_values("status_code", http_requests);
+
+TQL EVAL (0, 15, '5s') count_values("status_code", http_requests) by (idc);
+
+DROP TABLE http_requests;
diff --git a/tests/cases/standalone/common/promql/quantile.result b/tests/cases/standalone/common/promql/quantile.result
new file mode 100644
index 000000000000..8676bbbb777b
--- /dev/null
+++ b/tests/cases/standalone/common/promql/quantile.result
@@ -0,0 +1,71 @@
+CREATE TABLE test (
+  ts timestamp(3) time index,
+  host STRING,
+  idc STRING,
+  val BIGINT,
+  PRIMARY KEY(host, idc),
+);
+
+Affected Rows: 0
+
+INSERT INTO TABLE test VALUES
+    (0,     'host1', "idc1", 1),
+    (0,     'host2', "idc1", 2),
+    (0,     'host3', "idc2", 3),
+    (0,     'host4', "idc2", 4),
+    (5000,  'host1', "idc1", 5),
+    (5000,  'host2', "idc1", 6),
+    (5000,  'host3', "idc2", 7),
+    (5000,  'host4', "idc2", 8),
+    (10000, 'host1', "idc1", 9),
+    (10000, 'host2', "idc1", 10),
+    (10000, 'host3', "idc2", 11),
+    (10000, 'host4', "idc2", 12),
+    (15000, 'host1', "idc1", 13),
+    (15000, 'host2', "idc1", 14),
+    (15000, 'host3', "idc2", 15),
+    (15000, 'host4', "idc2", 16);
+
+Affected Rows: 16
+
+TQL EVAL (0, 15, '5s') quantile(0.5, test);
+
++---------------------+--------------------+
+| ts                  | quantile(test.val) |
++---------------------+--------------------+
+| 1970-01-01T00:00:00 | 2.5                |
+| 1970-01-01T00:00:05 | 6.5                |
+| 1970-01-01T00:00:10 | 10.5               |
+| 1970-01-01T00:00:15 | 14.5               |
++---------------------+--------------------+
+
+TQL EVAL (0, 15, '5s') quantile(0.5, test) by (idc);
+
++------+---------------------+--------------------+
+| idc  | ts                  | quantile(test.val) |
++------+---------------------+--------------------+
+| idc1 | 1970-01-01T00:00:00 | 1.5                |
+| idc1 | 1970-01-01T00:00:05 | 5.5                |
+| idc1 | 1970-01-01T00:00:10 | 9.5                |
+| idc1 | 1970-01-01T00:00:15 | 13.5               |
+| idc2 | 1970-01-01T00:00:00 | 3.5                |
+| idc2 | 1970-01-01T00:00:05 | 7.5                |
+| idc2 | 1970-01-01T00:00:10 | 11.5               |
+| idc2 | 1970-01-01T00:00:15 | 15.5               |
++------+---------------------+--------------------+
+
+TQL EVAL (0, 15, '5s') quantile(0.5, sum(test) by (idc));
+
++---------------------+-------------------------+
+| ts                  | quantile(sum(test.val)) |
++---------------------+-------------------------+
+| 1970-01-01T00:00:00 | 5.0                     |
+| 1970-01-01T00:00:05 | 13.0                    |
+| 1970-01-01T00:00:10 | 21.0                    |
+| 1970-01-01T00:00:15 | 29.0                    |
++---------------------+-------------------------+
+
+DROP TABLE test;
+
+Affected Rows: 0
+
diff --git a/tests/cases/standalone/common/promql/quantile.sql b/tests/cases/standalone/common/promql/quantile.sql
new file mode 100644
index 000000000000..4bef39c7570f
--- /dev/null
+++ b/tests/cases/standalone/common/promql/quantile.sql
@@ -0,0 +1,33 @@
+CREATE TABLE test (
+  ts timestamp(3) time index,
+  host STRING,
+  idc STRING,
+  val BIGINT,
+  PRIMARY KEY(host, idc),
+);
+
+INSERT INTO TABLE test VALUES
+    (0,     'host1', "idc1", 1),
+    (0,     'host2', "idc1", 2),
+    (0,     'host3', "idc2", 3),
+    (0,     'host4', "idc2", 4),
+    (5000,  'host1', "idc1", 5),
+    (5000,  'host2', "idc1", 6),
+    (5000,  'host3', "idc2", 7),
+    (5000,  'host4', "idc2", 8),
+    (10000, 'host1', "idc1", 9),
+    (10000, 'host2', "idc1", 10),
+    (10000, 'host3', "idc2", 11),
+    (10000, 'host4', "idc2", 12),
+    (15000, 'host1', "idc1", 13),
+    (15000, 'host2', "idc1", 14),
+    (15000, 'host3', "idc2", 15),
+    (15000, 'host4', "idc2", 16);
+
+TQL EVAL (0, 15, '5s') quantile(0.5, test);
+
+TQL EVAL (0, 15, '5s') quantile(0.5, test) by (idc);
+
+TQL EVAL (0, 15, '5s') quantile(0.5, sum(test) by (idc));
+
+DROP TABLE test;


Reviews:
user: login: waynexia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T01:08:20Z
user: login: waynexia
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-07T01:08:54Z
user: login: waynexia
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T01:20:30Z
user: login: killme2008
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T05:34:04Z
user: login: killme2008
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-07T05:34:12Z
user: login: fengjiachun
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-10T06:41:29Z

Review comments:
diff_hunk: @@ -0,0 +1,298 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, AsArray};
+use datafusion::common::cast::{as_list_array, as_primitive_array, as_struct_array};
+use datafusion::error::Result as DfResult;
+use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility};
+use datafusion::prelude::create_udaf;
+use datafusion_common::ScalarValue;
+use datatypes::arrow::array::{ListArray, StructArray};
+use datatypes::arrow::datatypes::{DataType, Field, Float64Type};
+
+use crate::functions::quantile::quantile_impl;
+
+const QUANTILE_NAME: &str = "quantile";
+
+const VALUES_FIELD_NAME: &str = "values";
+const DEFAULT_LIST_FIELD_NAME: &str = "item";
+
+#[derive(Debug, Default)]
+pub struct QuantileAccumulator {
+    q: f64,
+    values: Vec<Option<f64>>,
+}
+
+/// Create a quantile `AggregateUDF` for PromQL quantile operator,
+/// which calculates φ-quantile (0 ≤ φ ≤ 1) over dimensions
+pub fn quantile_udaf(q: f64) -> Arc<AggregateUDF> {
+    Arc::new(create_udaf(
+        QUANTILE_NAME,
+        // Input type: (values)
+        vec![DataType::Float64],
+        // Output type: the φ-quantile
+        Arc::new(DataType::Float64),
+        Volatility::Immutable,
+        // Create the accumulator
+        Arc::new(move |_| Ok(Box::new(QuantileAccumulator::new(q)))),
+        // Intermediate state types
+        Arc::new(vec![DataType::Struct(
+            vec![Field::new(
+                VALUES_FIELD_NAME,
+                DataType::List(Arc::new(Field::new(
+                    DEFAULT_LIST_FIELD_NAME,
+                    DataType::Float64,
+                    true,
+                ))),
+                false,
+            )]
+            .into(),
+        )]),
+    ))
+}
+
+impl QuantileAccumulator {
+    pub fn new(q: f64) -> Self {
+        Self {
+            q,
+            ..Default::default()
+        }
+    }
+}
+
+impl DfAccumulator for QuantileAccumulator {
+    fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> {
+        let f64_array = values[0].as_primitive::<Float64Type>();
+
+        self.values.extend(f64_array);
+
+        Ok(())
+    }
+
+    fn evaluate(&mut self) -> DfResult<ScalarValue> {
+        let values: Vec<_> = self.values.iter().map(|v| v.unwrap_or(0.0)).collect();
+
+        let result = quantile_impl(&values, self.q);
+
+        ScalarValue::new_primitive::<Float64Type>(result, &DataType::Float64)
+    }
+
+    fn size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.values.capacity() * std::mem::size_of::<Option<f64>>()
+    }
+
+    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
+        let values_array = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
+            Some(self.values.clone()),
+        ]));
+
+        let state_struct = StructArray::new(
+            vec![Field::new(
+                VALUES_FIELD_NAME,
+                DataType::List(Arc::new(Field::new(
+                    DEFAULT_LIST_FIELD_NAME,
+                    DataType::Float64,
+                    true,
+                ))),
+                false,
+            )]
+            .into(),
+            vec![values_array],
+            None,
+        );
+
+        Ok(vec![ScalarValue::Struct(Arc::new(state_struct))])
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        for state in states {
+            let state = as_struct_array(state)?;
+
+            for list in as_list_array(state.column(0))?.iter().flatten() {
+                let f64_array = as_primitive_array::<Float64Type>(&list)?.clone();
+                self.values.extend(&f64_array);
+            }
+        }
+
+        Ok(())
+    }
+}
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use datafusion::arrow::array::{ArrayRef, Float64Array};
+    use datafusion_common::ScalarValue;
+
+    use super::*;
+
+    fn create_f64_array(values: Vec<Option<f64>>) -> ArrayRef {
+        Arc::new(Float64Array::from(values)) as ArrayRef
+    }
+
+    #[test]
+    fn test_quantile_accumulator_empty() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+
+        let result = accumulator.evaluate().unwrap();
+
+        match result {
+            ScalarValue::Float64(_) => (),
+            _ => panic!("Expected Float64 scalar value"),
+        }
+    }
+
+    #[test]
+    fn test_quantile_accumulator_single_value() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(10.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let result = accumulator.evaluate().unwrap();
+
+        assert_eq!(result, ScalarValue::Float64(Some(10.0)));
+    }
+
+    #[test]
+    fn test_quantile_accumulator_multiple_values() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0), Some(5.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let result = accumulator.evaluate().unwrap();
+
+        // 中位数应该是 3.0
path: src/promql/src/functions/quantile_aggr.rs
user: login: waynexia
body: 👀 maybe you can try prompting in English
created_at: 2025-03-07T01:08:20Z
author_association: MEMBER
diff_hunk: @@ -1966,27 +1982,81 @@ impl PromPlanner {
             .field_columns
             .iter()
             .map(|col| {
-                DfExpr::AggregateFunction(AggregateFunction {
+                Ok(DfExpr::AggregateFunction(AggregateFunction {
                     func: aggr.clone(),
                     args: vec![DfExpr::Column(Column::from_name(col))],
                     distinct: false,
                     filter: None,
                     order_by: None,
                     null_treatment: None,
-                })
+                }))
             })
-            .collect();
+            .collect::<Result<Vec<_>>>()?;
+
+        // if the aggregator is `count_values`, it must be grouped by current fields.
+        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
+            let prev_field_exprs: Vec<_> = self
+                .ctx
+                .field_columns
+                .iter()
+                .map(|col| DfExpr::Column(Column::from_name(col)))
+                .collect();
 
-        // update value column name according to the aggregators
+            ensure!(
+                self.ctx.field_columns.len() == 1,
+                UnsupportedExprSnafu {
+                    name: "count_values on multi-value input"
+                }
+            );
+
+            prev_field_exprs
+        } else {
+            vec![]
+        };
+
+        // update value column name according to the aggregators,
         let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
+
         let normalized_exprs =
             normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
         for expr in normalized_exprs {
             new_field_columns.push(expr.schema_name().to_string());
         }
         self.ctx.field_columns = new_field_columns;
 
-        Ok(exprs)
+        Ok((exprs, prev_field_exprs))
path: src/query/src/promql/planner.rs
user: login: waynexia
body: add some documents for this new return value
created_at: 2025-03-07T01:18:00Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,298 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, AsArray};
+use datafusion::common::cast::{as_list_array, as_primitive_array, as_struct_array};
+use datafusion::error::Result as DfResult;
+use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility};
+use datafusion::prelude::create_udaf;
+use datafusion_common::ScalarValue;
+use datatypes::arrow::array::{ListArray, StructArray};
+use datatypes::arrow::datatypes::{DataType, Field, Float64Type};
+
+use crate::functions::quantile::quantile_impl;
+
+const QUANTILE_NAME: &str = "quantile";
+
+const VALUES_FIELD_NAME: &str = "values";
+const DEFAULT_LIST_FIELD_NAME: &str = "item";
+
+#[derive(Debug, Default)]
+pub struct QuantileAccumulator {
+    q: f64,
+    values: Vec<Option<f64>>,
+}
+
+/// Create a quantile `AggregateUDF` for PromQL quantile operator,
+/// which calculates φ-quantile (0 ≤ φ ≤ 1) over dimensions
+pub fn quantile_udaf(q: f64) -> Arc<AggregateUDF> {
+    Arc::new(create_udaf(
+        QUANTILE_NAME,
+        // Input type: (values)
+        vec![DataType::Float64],
+        // Output type: the φ-quantile
+        Arc::new(DataType::Float64),
+        Volatility::Immutable,
+        // Create the accumulator
+        Arc::new(move |_| Ok(Box::new(QuantileAccumulator::new(q)))),
+        // Intermediate state types
+        Arc::new(vec![DataType::Struct(
+            vec![Field::new(
+                VALUES_FIELD_NAME,
+                DataType::List(Arc::new(Field::new(
+                    DEFAULT_LIST_FIELD_NAME,
+                    DataType::Float64,
+                    true,
+                ))),
+                false,
+            )]
+            .into(),
+        )]),
+    ))
+}
+
+impl QuantileAccumulator {
+    pub fn new(q: f64) -> Self {
+        Self {
+            q,
+            ..Default::default()
+        }
+    }
+}
+
+impl DfAccumulator for QuantileAccumulator {
+    fn update_batch(&mut self, values: &[ArrayRef]) -> DfResult<()> {
+        let f64_array = values[0].as_primitive::<Float64Type>();
+
+        self.values.extend(f64_array);
+
+        Ok(())
+    }
+
+    fn evaluate(&mut self) -> DfResult<ScalarValue> {
+        let values: Vec<_> = self.values.iter().map(|v| v.unwrap_or(0.0)).collect();
+
+        let result = quantile_impl(&values, self.q);
+
+        ScalarValue::new_primitive::<Float64Type>(result, &DataType::Float64)
+    }
+
+    fn size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.values.capacity() * std::mem::size_of::<Option<f64>>()
+    }
+
+    fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
+        let values_array = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
+            Some(self.values.clone()),
+        ]));
+
+        let state_struct = StructArray::new(
+            vec![Field::new(
+                VALUES_FIELD_NAME,
+                DataType::List(Arc::new(Field::new(
+                    DEFAULT_LIST_FIELD_NAME,
+                    DataType::Float64,
+                    true,
+                ))),
+                false,
+            )]
+            .into(),
+            vec![values_array],
+            None,
+        );
+
+        Ok(vec![ScalarValue::Struct(Arc::new(state_struct))])
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> DfResult<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        for state in states {
+            let state = as_struct_array(state)?;
+
+            for list in as_list_array(state.column(0))?.iter().flatten() {
+                let f64_array = as_primitive_array::<Float64Type>(&list)?.clone();
+                self.values.extend(&f64_array);
+            }
+        }
+
+        Ok(())
+    }
+}
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use datafusion::arrow::array::{ArrayRef, Float64Array};
+    use datafusion_common::ScalarValue;
+
+    use super::*;
+
+    fn create_f64_array(values: Vec<Option<f64>>) -> ArrayRef {
+        Arc::new(Float64Array::from(values)) as ArrayRef
+    }
+
+    #[test]
+    fn test_quantile_accumulator_empty() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+
+        let result = accumulator.evaluate().unwrap();
+
+        match result {
+            ScalarValue::Float64(_) => (),
+            _ => panic!("Expected Float64 scalar value"),
+        }
+    }
+
+    #[test]
+    fn test_quantile_accumulator_single_value() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(10.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let result = accumulator.evaluate().unwrap();
+
+        assert_eq!(result, ScalarValue::Float64(Some(10.0)));
+    }
+
+    #[test]
+    fn test_quantile_accumulator_multiple_values() {
+        let mut accumulator = QuantileAccumulator::new(0.5);
+        let input = create_f64_array(vec![Some(1.0), Some(2.0), Some(3.0), Some(4.0), Some(5.0)]);
+
+        accumulator.update_batch(&[input]).unwrap();
+        let result = accumulator.evaluate().unwrap();
+
+        // 中位数应该是 3.0
path: src/promql/src/functions/quantile_aggr.rs
user: login: killme2008
body: hmmm, forgot to delete it.
created_at: 2025-03-07T05:34:04Z
author_association: MEMBER
diff_hunk: @@ -1966,27 +1982,81 @@ impl PromPlanner {
             .field_columns
             .iter()
             .map(|col| {
-                DfExpr::AggregateFunction(AggregateFunction {
+                Ok(DfExpr::AggregateFunction(AggregateFunction {
                     func: aggr.clone(),
                     args: vec![DfExpr::Column(Column::from_name(col))],
                     distinct: false,
                     filter: None,
                     order_by: None,
                     null_treatment: None,
-                })
+                }))
             })
-            .collect();
+            .collect::<Result<Vec<_>>>()?;
+
+        // if the aggregator is `count_values`, it must be grouped by current fields.
+        let prev_field_exprs = if op.id() == token::T_COUNT_VALUES {
+            let prev_field_exprs: Vec<_> = self
+                .ctx
+                .field_columns
+                .iter()
+                .map(|col| DfExpr::Column(Column::from_name(col)))
+                .collect();
 
-        // update value column name according to the aggregators
+            ensure!(
+                self.ctx.field_columns.len() == 1,
+                UnsupportedExprSnafu {
+                    name: "count_values on multi-value input"
+                }
+            );
+
+            prev_field_exprs
+        } else {
+            vec![]
+        };
+
+        // update value column name according to the aggregators,
         let mut new_field_columns = Vec::with_capacity(self.ctx.field_columns.len());
+
         let normalized_exprs =
             normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?;
         for expr in normalized_exprs {
             new_field_columns.push(expr.schema_name().to_string());
         }
         self.ctx.field_columns = new_field_columns;
 
-        Ok(exprs)
+        Ok((exprs, prev_field_exprs))
path: src/query/src/promql/planner.rs
user: login: killme2008
body: Ok
created_at: 2025-03-07T05:34:12Z
author_association: MEMBER

PR comments:
user: login: coderabbitai[bot]
created_at: 2025-03-05T16:03:06Z
body: <!-- This is an auto-generated comment: summarize by coderabbit.ai -->
<!-- This is an auto-generated comment: skip review by coderabbit.ai -->

> [!IMPORTANT]
> ## Review skipped
> 
> Auto reviews are disabled on this repository.
> 
> Please check the settings in the CodeRabbit UI or the `.coderabbit.yaml` file in this repository. To trigger a single review, invoke the `@coderabbitai review` command.
> 
> You can disable this status message by setting the `reviews.review_status` to `false` in the CodeRabbit configuration file.

<!-- end of auto-generated comment: skip review by coderabbit.ai -->
<!-- tips_start -->

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

<details>
<summary>❤️ Share</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai)
- [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)

</details>

<details>
<summary>🪧 Tips</summary>

### Chat

There are 3 ways to chat with [CodeRabbit](https://coderabbit.ai?utm_source=oss&utm_medium=github&utm_campaign=GreptimeTeam/greptimedb&utm_content=5652):

- Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  - `I pushed a fix in commit <commit_id>, please review it.`
  - `Generate unit testing code for this file.`
  - `Open a follow-up GitHub issue for this discussion.`
- Files and specific lines of code (under the "Files changed" tab): Tag `@coderabbitai` in a new review comment at the desired location with your query. Examples:
  - `@coderabbitai generate unit testing code for this file.`
  -	`@coderabbitai modularize this function.`
- PR comments: Tag `@coderabbitai` in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  - `@coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.`
  - `@coderabbitai read src/utils.ts and generate unit testing code.`
  - `@coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.`
  - `@coderabbitai help me debug CodeRabbit configuration file.`

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

### CodeRabbit Commands (Invoked using PR comments)

- `@coderabbitai pause` to pause the reviews on a PR.
- `@coderabbitai resume` to resume the paused reviews.
- `@coderabbitai review` to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
- `@coderabbitai full review` to do a full review from scratch and review all the files again.
- `@coderabbitai summary` to regenerate the summary of the PR.
- `@coderabbitai generate docstrings` to [generate docstrings](https://docs.coderabbit.ai/finishing-touches/docstrings) for this PR.
- `@coderabbitai resolve` resolve all the CodeRabbit review comments.
- `@coderabbitai configuration` to show the current CodeRabbit configuration for the repository.
- `@coderabbitai help` to get help.

### Other keywords and placeholders

- Add `@coderabbitai ignore` anywhere in the PR description to prevent this PR from being reviewed.
- Add `@coderabbitai summary` to generate the high-level summary at a specific location in the PR description.
- Add `@coderabbitai` anywhere in the PR title to generate the title automatically.

### Documentation and Community

- Visit our [Documentation](https://docs.coderabbit.ai) for detailed information on how to use CodeRabbit.
- Join our [Discord Community](http://discord.gg/coderabbit) to get help, request features, and share feedback.
- Follow us on [X/Twitter](https://twitter.com/coderabbitai) for updates and announcements.

</details>

<!-- tips_end -->
author_association: CONTRIBUTOR
user: login: killme2008
created_at: 2025-03-06T07:51:29Z
body: It's ready for review, please take a look @WenyXu @waynexia @evenyag 
author_association: MEMBER