You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by co...@apache.org on 2026/03/30 17:28:06 UTC
(datafusion-comet) branch main updated: feat: support SQL aggregate FILTER (WHERE ...) clause in native execution (#3835)
This is an automated email from the ASF dual-hosted git repository.
comphead pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new e5e452a2e feat: support SQL aggregate FILTER (WHERE ...) clause in native execution (#3835)
e5e452a2e is described below
commit e5e452a2e44d59b8fede66ee912a02e575b2c3b3
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Tue Mar 31 01:28:01 2026 +0800
feat: support SQL aggregate FILTER (WHERE ...) clause in native execution (#3835)
---
dev/diffs/4.0.1.diff | 15 +++
native/core/src/execution/planner.rs | 18 ++-
native/proto/src/proto/expr.proto | 4 +
native/spark-expr/src/agg_funcs/avg.rs | 9 +-
native/spark-expr/src/agg_funcs/avg_decimal.rs | 9 +-
native/spark-expr/src/agg_funcs/sum_decimal.rs | 65 ++++++++-
native/spark-expr/src/agg_funcs/sum_int.rs | 146 ++++++++++++++++++++-
.../org/apache/comet/serde/QueryPlanSerde.scala | 16 ++-
.../org/apache/spark/sql/comet/operators.scala | 6 -
.../expressions/aggregate/aggregate_filter.sql | 86 ++++++++++++
10 files changed, 350 insertions(+), 24 deletions(-)
diff --git a/dev/diffs/4.0.1.diff b/dev/diffs/4.0.1.diff
index a0b1e81d0..364c62a99 100644
--- a/dev/diffs/4.0.1.diff
+++ b/dev/diffs/4.0.1.diff
@@ -245,6 +245,21 @@ index aa3d02dc2fb..c4f878d9908 100644
-- Test cases with unicode_rtrim.
WITH t(c1) AS (SELECT replace(listagg(DISTINCT col1 COLLATE unicode_rtrim) COLLATE utf8_binary, ' ', '') FROM (VALUES ('xbc '), ('xbc '), ('a'), ('xbc'))) SELECT len(c1), regexp_count(c1, 'a'), regexp_count(c1, 'xbc') FROM t;
WITH t(c1) AS (SELECT listagg(col1) WITHIN GROUP (ORDER BY col1 COLLATE unicode_rtrim) FROM (VALUES ('abc '), ('abc\n'), ('abc'), ('x'))) SELECT replace(replace(c1, ' ', ''), '\n', '$') FROM t;
+diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
+index 0000000..0000000 100644
+--- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
++++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
+@@ -6,6 +6,10 @@
+ -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/aggregates.sql#L352-L605
+
+ -- Test aggregate operator with codegen on and off.
++
++-- Floating-point precision difference between DataFusion and JVM for FILTER aggregates
++--SET spark.comet.enabled = false
++
+ --CONFIG_DIM1 spark.sql.codegen.wholeStage=true
+ --CONFIG_DIM1 spark.sql.codegen.wholeStage=false,spark.sql.codegen.factoryMode=CODEGEN_ONLY
+ --CONFIG_DIM1 spark.sql.codegen.wholeStage=false,spark.sql.codegen.factoryMode=NO_CODEGEN
diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql
index 3a409eea348..26e9aaf215c 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 5af31fcc2..0f96c829e 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -973,14 +973,28 @@ impl PhysicalPlanner {
.map(|expr| self.create_agg_expr(expr, Arc::clone(&schema)))
.collect();
- let num_agg = agg.agg_exprs.len();
let aggr_expr = agg_exprs?.into_iter().map(Arc::new).collect();
+
+ // Build per-aggregate filter expressions from the FILTER (WHERE ...) clause.
+ // Filters are only present in Partial mode; Final/PartialMerge always get None.
+ let filter_exprs: Result<Vec<Option<Arc<dyn PhysicalExpr>>>, ExecutionError> = agg
+ .agg_exprs
+ .iter()
+ .map(|expr| {
+ if let Some(f) = expr.filter.as_ref() {
+ self.create_expr(f, Arc::clone(&schema)).map(Some)
+ } else {
+ Ok(None)
+ }
+ })
+ .collect();
+
let aggregate: Arc<dyn ExecutionPlan> = Arc::new(
datafusion::physical_plan::aggregates::AggregateExec::try_new(
mode,
group_by,
aggr_expr,
- vec![None; num_agg], // no filter expressions
+ filter_exprs?,
Arc::clone(&child.native_plan),
Arc::clone(&schema),
)?,
diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto
index 32cbc0ce1..c12b29df1 100644
--- a/native/proto/src/proto/expr.proto
+++ b/native/proto/src/proto/expr.proto
@@ -141,6 +141,10 @@ message AggExpr {
BloomFilterAgg bloomFilterAgg = 16;
}
+ // Optional filter expression for SQL FILTER (WHERE ...) clause.
+ // Only set in Partial aggregation mode; absent in Final/PartialMerge.
+ optional Expr filter = 89;
+
// Optional QueryContext for error reporting (contains SQL text and position)
optional QueryContext query_context = 90;
diff --git a/native/spark-expr/src/agg_funcs/avg.rs b/native/spark-expr/src/agg_funcs/avg.rs
index d1d71cca2..3760b4250 100644
--- a/native/spark-expr/src/agg_funcs/avg.rs
+++ b/native/spark-expr/src/agg_funcs/avg.rs
@@ -245,7 +245,7 @@ where
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
- _opt_filter: Option<&arrow::array::BooleanArray>,
+ opt_filter: Option<&arrow::array::BooleanArray>,
total_num_groups: usize,
) -> Result<()> {
assert_eq!(values.len(), 1, "single argument to update_batch");
@@ -257,7 +257,7 @@ where
self.sums.resize(total_num_groups, T::default_value());
let iter = group_indices.iter().zip(data.iter());
- if values.null_count() == 0 {
+ if opt_filter.is_none() && values.null_count() == 0 {
for (&group_index, &value) in iter {
let sum = &mut self.sums[group_index];
// No overflow checking - Infinity is a valid result
@@ -266,6 +266,11 @@ where
}
} else {
for (idx, (&group_index, &value)) in iter.enumerate() {
+ if let Some(f) = opt_filter {
+ if !f.is_valid(idx) || !f.value(idx) {
+ continue;
+ }
+ }
if values.is_null(idx) {
continue;
}
diff --git a/native/spark-expr/src/agg_funcs/avg_decimal.rs b/native/spark-expr/src/agg_funcs/avg_decimal.rs
index 08e335f42..9e8a31afa 100644
--- a/native/spark-expr/src/agg_funcs/avg_decimal.rs
+++ b/native/spark-expr/src/agg_funcs/avg_decimal.rs
@@ -504,7 +504,7 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator {
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
- _opt_filter: Option<&arrow::array::BooleanArray>,
+ opt_filter: Option<&arrow::array::BooleanArray>,
total_num_groups: usize,
) -> Result<()> {
assert_eq!(values.len(), 1, "single argument to update_batch");
@@ -517,12 +517,17 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator {
ensure_bit_capacity(&mut self.is_not_null, total_num_groups);
let iter = group_indices.iter().zip(data.iter());
- if values.null_count() == 0 {
+ if opt_filter.is_none() && values.null_count() == 0 {
for (&group_index, &value) in iter {
self.update_single(group_index, value)?;
}
} else {
for (idx, (&group_index, &value)) in iter.enumerate() {
+ if let Some(f) = opt_filter {
+ if !f.is_valid(idx) || !f.value(idx) {
+ continue;
+ }
+ }
if values.is_null(idx) {
continue;
}
diff --git a/native/spark-expr/src/agg_funcs/sum_decimal.rs b/native/spark-expr/src/agg_funcs/sum_decimal.rs
index 56a735493..46db7f36b 100644
--- a/native/spark-expr/src/agg_funcs/sum_decimal.rs
+++ b/native/spark-expr/src/agg_funcs/sum_decimal.rs
@@ -446,7 +446,6 @@ impl GroupsAccumulator for SumDecimalGroupsAccumulator {
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> DFResult<()> {
- assert!(opt_filter.is_none(), "opt_filter is not supported yet");
assert_eq!(values.len(), 1);
let values = values[0].as_primitive::<Decimal128Type>();
let data = values.values();
@@ -454,12 +453,17 @@ impl GroupsAccumulator for SumDecimalGroupsAccumulator {
self.resize_helper(total_num_groups);
let iter = group_indices.iter().zip(data.iter());
- if values.null_count() == 0 {
+ if opt_filter.is_none() && values.null_count() == 0 {
for (&group_index, &value) in iter {
self.update_single(group_index, value)?;
}
} else {
for (idx, (&group_index, &value)) in iter.enumerate() {
+ if let Some(f) = opt_filter {
+ if !f.is_valid(idx) || !f.value(idx) {
+ continue;
+ }
+ }
if values.is_null(idx) {
continue;
}
@@ -540,7 +544,10 @@ impl GroupsAccumulator for SumDecimalGroupsAccumulator {
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> DFResult<()> {
- assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ debug_assert!(
+ opt_filter.is_none(),
+ "opt_filter is not supported in merge_batch"
+ );
self.resize_helper(total_num_groups);
@@ -712,4 +719,56 @@ mod tests {
let schema = Schema::new(fields);
RecordBatch::try_new(Arc::new(schema), columns).unwrap()
}
+
+ #[test]
+ fn test_update_batch_with_filter() {
+ use arrow::array::Decimal128Array;
+ use datafusion::logical_expr::{EmitTo, GroupsAccumulator};
+
+ let data_type = DataType::Decimal128(10, 2);
+ let mut acc = SumDecimalGroupsAccumulator::new(
+ data_type.clone(),
+ 10,
+ EvalMode::Legacy,
+ None,
+ crate::create_query_context_map(),
+ );
+
+ // values: [100, 200, 300, 400], filter: [T, F, T, F] => sum = 100+300 = 400
+ let values: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![100i128, 200, 300, 400]).with_data_type(data_type.clone()),
+ );
+ let filter = BooleanArray::from(vec![true, false, true, false]);
+ acc.update_batch(&[values], &[0, 0, 0, 0], Some(&filter), 1)
+ .unwrap();
+
+ let result = acc.evaluate(EmitTo::All).unwrap();
+ let result = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
+ assert_eq!(result.value(0), 400);
+ }
+
+ #[test]
+ fn test_update_batch_filter_null_treated_as_exclude() {
+ use arrow::array::Decimal128Array;
+ use datafusion::logical_expr::{EmitTo, GroupsAccumulator};
+
+ let data_type = DataType::Decimal128(10, 2);
+ let mut acc = SumDecimalGroupsAccumulator::new(
+ data_type.clone(),
+ 10,
+ EvalMode::Legacy,
+ None,
+ crate::create_query_context_map(),
+ );
+
+ let values: ArrayRef =
+ Arc::new(Decimal128Array::from(vec![10i128, 20, 30]).with_data_type(data_type.clone()));
+ let filter = BooleanArray::from(vec![Some(true), None, Some(true)]);
+ acc.update_batch(&[values], &[0, 0, 0], Some(&filter), 1)
+ .unwrap();
+
+ let result = acc.evaluate(EmitTo::All).unwrap();
+ let result = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
+ assert_eq!(result.value(0), 40); // 10 + 30 = 40
+ }
}
diff --git a/native/spark-expr/src/agg_funcs/sum_int.rs b/native/spark-expr/src/agg_funcs/sum_int.rs
index 2ea07c743..781528521 100644
--- a/native/spark-expr/src/agg_funcs/sum_int.rs
+++ b/native/spark-expr/src/agg_funcs/sum_int.rs
@@ -457,12 +457,18 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorLegacy {
int_array: &PrimitiveArray<T>,
group_indices: &[usize],
sums: &mut [Option<i64>],
+ opt_filter: Option<&BooleanArray>,
) -> DFResult<()>
where
T: ArrowPrimitiveType,
T::Native: ArrowNativeType,
{
for (i, &group_index) in group_indices.iter().enumerate() {
+ if let Some(f) = opt_filter {
+ if !f.is_valid(i) || !f.value(i) {
+ continue;
+ }
+ }
if !int_array.is_null(i) {
let v = int_array.value(i).to_i64().ok_or_else(|| {
DataFusionError::Internal("Failed to convert value to i64".to_string())
@@ -473,7 +479,6 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorLegacy {
Ok(())
}
- debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
let values = &values[0];
self.sums.resize(total_num_groups, None);
@@ -482,21 +487,25 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorLegacy {
as_primitive_array::<Int64Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
DataType::Int32 => update_groups_sum(
as_primitive_array::<Int32Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
DataType::Int16 => update_groups_sum(
as_primitive_array::<Int16Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
DataType::Int8 => update_groups_sum(
as_primitive_array::<Int8Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
_ => {
return Err(DataFusionError::Internal(format!(
@@ -534,7 +543,10 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorLegacy {
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> DFResult<()> {
- debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ debug_assert!(
+ opt_filter.is_none(),
+ "opt_filter is not supported in merge_batch"
+ );
if values.len() != 1 {
return Err(DataFusionError::Internal(format!(
@@ -589,12 +601,18 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorAnsi {
int_array: &PrimitiveArray<T>,
group_indices: &[usize],
sums: &mut [Option<i64>],
+ opt_filter: Option<&BooleanArray>,
) -> DFResult<()>
where
T: ArrowPrimitiveType,
T::Native: ArrowNativeType,
{
for (i, &group_index) in group_indices.iter().enumerate() {
+ if let Some(f) = opt_filter {
+ if !f.is_valid(i) || !f.value(i) {
+ continue;
+ }
+ }
if !int_array.is_null(i) {
let v = int_array.value(i).to_i64().ok_or_else(|| {
DataFusionError::Internal("Failed to convert value to i64".to_string())
@@ -608,7 +626,6 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorAnsi {
Ok(())
}
- debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
let values = &values[0];
self.sums.resize(total_num_groups, None);
@@ -617,21 +634,25 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorAnsi {
as_primitive_array::<Int64Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
DataType::Int32 => update_groups_sum(
as_primitive_array::<Int32Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
DataType::Int16 => update_groups_sum(
as_primitive_array::<Int16Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
DataType::Int8 => update_groups_sum(
as_primitive_array::<Int8Type>(values),
group_indices,
&mut self.sums,
+ opt_filter,
)?,
_ => {
return Err(DataFusionError::Internal(format!(
@@ -669,7 +690,10 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorAnsi {
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> DFResult<()> {
- debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ debug_assert!(
+ opt_filter.is_none(),
+ "opt_filter is not supported in merge_batch"
+ );
if values.len() != 1 {
return Err(DataFusionError::Internal(format!(
@@ -737,12 +761,18 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorTry {
group_indices: &[usize],
sums: &mut [Option<i64>],
has_all_nulls: &mut [bool],
+ opt_filter: Option<&BooleanArray>,
) -> DFResult<()>
where
T: ArrowPrimitiveType,
T::Native: ArrowNativeType,
{
for (i, &group_index) in group_indices.iter().enumerate() {
+ if let Some(f) = opt_filter {
+ if !f.is_valid(i) || !f.value(i) {
+ continue;
+ }
+ }
if !int_array.is_null(i) {
// Skip if this group already overflowed
if !has_all_nulls[group_index] && sums[group_index].is_none() {
@@ -760,8 +790,6 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorTry {
}
Ok(())
}
-
- debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
let values = &values[0];
self.sums.resize(total_num_groups, Some(0));
self.has_all_nulls.resize(total_num_groups, true);
@@ -772,24 +800,28 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorTry {
group_indices,
&mut self.sums,
&mut self.has_all_nulls,
+ opt_filter,
)?,
DataType::Int32 => update_groups_sum(
as_primitive_array::<Int32Type>(values),
group_indices,
&mut self.sums,
&mut self.has_all_nulls,
+ opt_filter,
)?,
DataType::Int16 => update_groups_sum(
as_primitive_array::<Int16Type>(values),
group_indices,
&mut self.sums,
&mut self.has_all_nulls,
+ opt_filter,
)?,
DataType::Int8 => update_groups_sum(
as_primitive_array::<Int8Type>(values),
group_indices,
&mut self.sums,
&mut self.has_all_nulls,
+ opt_filter,
)?,
_ => {
return Err(DataFusionError::Internal(format!(
@@ -842,7 +874,10 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorTry {
opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> DFResult<()> {
- debug_assert!(opt_filter.is_none(), "opt_filter is not supported yet");
+ debug_assert!(
+ opt_filter.is_none(),
+ "opt_filter is not supported in merge_batch"
+ );
if values.len() != 2 {
return Err(DataFusionError::Internal(format!(
@@ -900,3 +935,100 @@ impl GroupsAccumulator for SumIntGroupsAccumulatorTry {
std::mem::size_of_val(self)
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::Int64Array;
+ use datafusion::logical_expr::{EmitTo, GroupsAccumulator};
+
+ fn run_update_batch_with_filter(
+ acc: &mut dyn GroupsAccumulator,
+ values: Vec<i64>,
+ groups: Vec<usize>,
+ filter: Vec<bool>,
+ num_groups: usize,
+ ) -> Vec<Option<i64>> {
+ let values: ArrayRef = Arc::new(Int64Array::from(values));
+ let filter = BooleanArray::from(filter);
+ acc.update_batch(&[values], &groups, Some(&filter), num_groups)
+ .unwrap();
+ acc.evaluate(EmitTo::All)
+ .unwrap()
+ .as_primitive::<Int64Type>()
+ .iter()
+ .collect()
+ }
+
+ #[test]
+ fn test_legacy_update_batch_with_filter() {
+ let mut acc = SumIntGroupsAccumulatorLegacy::new();
+ // values: [1, 2, 3, 4, 5], filter: [T, F, T, F, T] => sum = 1+3+5 = 9
+ let result = run_update_batch_with_filter(
+ &mut acc,
+ vec![1, 2, 3, 4, 5],
+ vec![0, 0, 0, 0, 0],
+ vec![true, false, true, false, true],
+ 1,
+ );
+ assert_eq!(result, vec![Some(9)]);
+ }
+
+ #[test]
+ fn test_legacy_update_batch_filter_null_treated_as_exclude() {
+ let mut acc = SumIntGroupsAccumulatorLegacy::new();
+ let values: ArrayRef = Arc::new(Int64Array::from(vec![10i64, 20, 30]));
+ // null filter entry should be treated as exclude
+ let filter = BooleanArray::from(vec![Some(true), None, Some(true)]);
+ acc.update_batch(&[values], &[0, 0, 0], Some(&filter), 1)
+ .unwrap();
+ let result: Vec<Option<i64>> = acc
+ .evaluate(EmitTo::All)
+ .unwrap()
+ .as_primitive::<Int64Type>()
+ .iter()
+ .collect();
+ assert_eq!(result, vec![Some(40)]); // 10 + 30 = 40
+ }
+
+ #[test]
+ fn test_ansi_update_batch_with_filter() {
+ let mut acc = SumIntGroupsAccumulatorAnsi::new();
+ let result = run_update_batch_with_filter(
+ &mut acc,
+ vec![10, 20, 30, 40],
+ vec![0, 1, 0, 1],
+ vec![true, true, false, true],
+ 2,
+ );
+ // group 0: 10 (30 filtered out); group 1: 20+40 = 60
+ assert_eq!(result, vec![Some(10), Some(60)]);
+ }
+
+ #[test]
+ fn test_try_update_batch_with_filter() {
+ let mut acc = SumIntGroupsAccumulatorTry::new();
+ let result = run_update_batch_with_filter(
+ &mut acc,
+ vec![1, 2, 3, 4, 5],
+ vec![0, 0, 0, 0, 0],
+ vec![true, false, true, false, true],
+ 1,
+ );
+ assert_eq!(result, vec![Some(9)]); // 1+3+5 = 9
+ }
+
+ #[test]
+ fn test_no_filter_still_works() {
+ let mut acc = SumIntGroupsAccumulatorLegacy::new();
+ let values: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 3]));
+ acc.update_batch(&[values], &[0, 0, 0], None, 1).unwrap();
+ let result: Vec<Option<i64>> = acc
+ .evaluate(EmitTo::All)
+ .unwrap()
+ .as_primitive::<Int64Type>()
+ .iter()
+ .collect();
+ assert_eq!(result, vec![Some(6)]);
+ }
+}
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 02a76f69f..2ce398c8f 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -517,13 +517,25 @@ object QueryPlanSerde extends Logging with CometExprShim {
}
// Attach QueryContext and expr_id to the aggregate expression
- protoAggExprOpt.map { protoAggExpr =>
+ protoAggExprOpt.flatMap { protoAggExpr =>
val builder = protoAggExpr.toBuilder
builder.setExprId(nextExprId())
+
+ // Serialize FILTER (WHERE ...) clause if present.
+ // The filter is only meaningful in Partial mode; Final/PartialMerge never set it.
+ if (aggExpr.filter.isDefined && aggExpr.mode == Partial) {
+ val filterProto = exprToProto(aggExpr.filter.get, inputs, binding)
+ if (filterProto.isEmpty) {
+ withInfo(aggExpr, aggExpr.filter.get)
+ return None
+ }
+ builder.setFilter(filterProto.get)
+ }
+
extractQueryContext(fn).foreach { ctx =>
builder.setQueryContext(ctx)
}
- builder.build()
+ Some(builder.build())
}
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index 2965e4698..1f5e7b667 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -1362,12 +1362,6 @@ trait CometBaseAggregate {
return None
}
- // Aggregate expressions with filter are not supported yet.
- if (aggregateExpressions.exists(_.filter.isDefined)) {
- withInfo(aggregate, "Aggregate expression with filter is not supported")
- return None
- }
-
if (groupingExpressions.exists(expr =>
expr.dataType match {
case _: MapType => true
diff --git a/spark/src/test/resources/sql-tests/expressions/aggregate/aggregate_filter.sql b/spark/src/test/resources/sql-tests/expressions/aggregate/aggregate_filter.sql
new file mode 100644
index 000000000..c8787868a
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/aggregate/aggregate_filter.sql
@@ -0,0 +1,86 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you 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.
+
+-- Tests for SQL aggregate FILTER (WHERE ...) clause support.
+-- See https://github.com/apache/datafusion-comet/issues/XXXX
+
+statement
+CREATE TABLE test_agg_filter(
+ grp string,
+ i int,
+ l long,
+ d decimal(10, 2),
+ flag boolean
+) USING parquet
+
+statement
+INSERT INTO test_agg_filter VALUES
+ ('a', 1, 10, 1.00, true),
+ ('a', 2, 20, 2.00, false),
+ ('a', 3, 30, 3.00, true),
+ ('b', 4, 40, 4.00, false),
+ ('b', 5, 50, 5.00, true),
+ ('b', NULL, NULL, NULL, true)
+
+-- Basic FILTER on SUM(int)
+query
+SELECT SUM(i) FILTER (WHERE flag = true) FROM test_agg_filter
+
+-- FILTER on SUM with GROUP BY
+query
+SELECT grp, SUM(i) FILTER (WHERE flag = true) FROM test_agg_filter GROUP BY grp ORDER BY grp
+
+-- FILTER on SUM(long)
+query
+SELECT SUM(l) FILTER (WHERE flag = true) FROM test_agg_filter
+
+-- FILTER on SUM(decimal)
+query
+SELECT SUM(d) FILTER (WHERE flag = true) FROM test_agg_filter
+
+-- Multiple aggregates: one with filter, one without
+query
+SELECT SUM(i), SUM(i) FILTER (WHERE flag = true) FROM test_agg_filter
+
+-- FILTER with NULL rows: NULLs should not be included even when filter passes
+query
+SELECT grp, SUM(i) FILTER (WHERE flag = true) FROM test_agg_filter GROUP BY grp ORDER BY grp
+
+-- FILTER with COUNT
+query
+SELECT COUNT(*) FILTER (WHERE flag = true) FROM test_agg_filter
+
+-- FILTER with COUNT GROUP BY
+query
+SELECT grp, COUNT(*) FILTER (WHERE flag = true) FROM test_agg_filter GROUP BY grp ORDER BY grp
+
+-- FILTER on AVG(int)
+query
+SELECT AVG(i) FILTER (WHERE flag = true) FROM test_agg_filter
+
+-- FILTER on AVG with GROUP BY
+query
+SELECT grp, AVG(i) FILTER (WHERE flag = true) FROM test_agg_filter GROUP BY grp ORDER BY grp
+
+-- FILTER on AVG(decimal): final cast back to decimal may differ from Spark due to rounding,
+-- so use spark_answer_only mode to validate correctness without checking operator coverage
+query spark_answer_only
+SELECT AVG(d) FILTER (WHERE flag = true) FROM test_agg_filter
+
+query spark_answer_only
+SELECT grp, AVG(d) FILTER (WHERE flag = true) FROM test_agg_filter GROUP BY grp ORDER BY grp
+
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org