From 838f0ed6d5d3cc8c329162715e9136267c3e440e Mon Sep 17 00:00:00 2001
From: anant <anant.v09@proton.me>
Date: Thu, 15 May 2025 08:55:44 +0530
Subject: [PATCH] Add `count distinct` aggregate function

---
 src/alerts/alerts_utils.rs | 3 ++-
 src/alerts/mod.rs          | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/alerts/alerts_utils.rs b/src/alerts/alerts_utils.rs
index e1642cb23..68bc6910e 100644
--- a/src/alerts/alerts_utils.rs
+++ b/src/alerts/alerts_utils.rs
@@ -19,7 +19,7 @@
 use arrow_array::{Float64Array, Int64Array, RecordBatch};
 use datafusion::{
     functions_aggregate::{
-        count::count,
+        count::{count, count_distinct},
         expr_fn::avg,
         min_max::{max, min},
         sum::sum,
@@ -387,6 +387,7 @@ fn match_aggregate_operation(agg: &AggregateConfig) -> Expr {
     let column = format!(r#""{}""#, agg.column);
     match agg.aggregate_function {
         AggregateFunction::Avg => avg(col(column)),
+        AggregateFunction::CountDistinct => count_distinct(col(column)),
         AggregateFunction::Count => count(col(column)),
         AggregateFunction::Min => min(col(column)),
         AggregateFunction::Max => max(col(column)),
diff --git a/src/alerts/mod.rs b/src/alerts/mod.rs
index 8d48ae44c..e9324e96b 100644
--- a/src/alerts/mod.rs
+++ b/src/alerts/mod.rs
@@ -290,6 +290,7 @@ impl Display for WhereConfigOperator {
 pub enum AggregateFunction {
     Avg,
     Count,
+    CountDistinct,
     Min,
     Max,
     Sum,
@@ -300,6 +301,7 @@ impl Display for AggregateFunction {
         match self {
             AggregateFunction::Avg => write!(f, "Avg"),
             AggregateFunction::Count => write!(f, "Count"),
+            AggregateFunction::CountDistinct => write!(f, "CountDistinct"),
             AggregateFunction::Min => write!(f, "Min"),
             AggregateFunction::Max => write!(f, "Max"),
             AggregateFunction::Sum => write!(f, "Sum"),