Support for median(distinct) aggregation function (#10226)

* Support for median(distinct) aggregation function * Reduce duplication
apache · Apr 26, 2024 · c22c299 · c22c299
1 parent ff5eb8f
commit c22c299
Show file tree

Hide file tree

Showing 6 changed files with 437 additions and 43 deletions.
diff --git a/datafusion/core/benches/aggregate_query_sql.rs b/datafusion/core/benches/aggregate_query_sql.rs
@@ -163,6 +163,16 @@ fn criterion_benchmark(c: &mut Criterion) {
             )
         })
     });
+
+    c.bench_function("aggregate_query_distinct_median", |b| {
+        b.iter(|| {
+            query(
+                ctx.clone(),
+                "SELECT MEDIAN(DISTINCT u64_wide), MEDIAN(DISTINCT u64_narrow) \
+                 FROM t",
+            )
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

diff --git a/datafusion/physical-expr/src/aggregate/build_in.rs b/datafusion/physical-expr/src/aggregate/build_in.rs
@@ -352,14 +352,12 @@ pub fn create_aggregate_expr(
                 "APPROX_MEDIAN(DISTINCT) aggregations are not available"
             );
         }
-        (AggregateFunction::Median, false) => Arc::new(expressions::Median::new(
+        (AggregateFunction::Median, distinct) => Arc::new(expressions::Median::new(
             input_phy_exprs[0].clone(),
             name,
             data_type,
+            distinct,
         )),
-        (AggregateFunction::Median, true) => {
-            return not_impl_err!("MEDIAN(DISTINCT) aggregations are not available");
-        }
         (AggregateFunction::FirstValue, _) => Arc::new(
             expressions::FirstValue::new(
                 input_phy_exprs[0].clone(),