Support `avg(distinct)` for `float64` type #17255

Jefffrey · 2025-08-20T05:11:32Z

Which issue does this PR close?

Relates to #2408

Decimal support not included so can't close issue yet

Rationale for this change

Building on old PR #15413, to get it over the line. Opened as new PR as other PR was too old and wasn't sure if should push to the original author branch; for now pushed to my own branch but preserved the original commits.

What changes are included in this PR?

From original PR:

mv DistinctSumAccumulator to common so that it can be used in Float64DistinctAvgAccumulator
implement Float64DistinctAvgAccumulator using DistinctSumAccumulator
tested in aggregate.slt

Additional changes made by me:

Fixed error query error DataFusion error: Arrow error: Invalid argument error: number of columns\(1\) must match number of fields\(2\) in schema (Support Avg distinct for float64 type #15413 (comment)) which was caused by wrong state fields (and also but not disabling group accumulator support if distinct)
Update the SLT test cases (removed the decimal test cases as those aren't relevant for this float PR)

Are these changes tested?

Added SLT tests, and also regenerated the extended tests: apache/datafusion-testing#11

Are there any user-facing changes?

Not sure if doc changes are required, as this is mainly for SQL and it seems we don't explicitly say avg(distinct) is disallowed so don't need to update anything saying it works now?

… and remove unused sum_distinct module

This reverts commit 1fa2db8.

Jefffrey · 2025-08-20T05:16:02Z

datafusion/functions-aggregate/src/average.rs

-        Ok(vec![
-            Field::new(
-                format_state_name(args.name, "count"),
-                DataType::UInt64,
-                true,
-            ),
-            Field::new(
-                format_state_name(args.name, "sum"),
-                args.input_fields[0].data_type().clone(),
-                true,
-            ),
-        ]
-        .into_iter()
-        .map(Arc::new)
-        .collect())
+        if args.is_distinct {
+            // Copied from datafusion_functions_aggregate::sum::Sum::state_fields
+            // since the accumulator uses DistinctSumAccumulator internally.
+            Ok(vec![Field::new_list(
+                format_state_name(args.name, "sum distinct"),
+                Field::new_list_field(args.return_type().clone(), true),
+                false,
+            )
+            .into()])
+        } else {
+            Ok(vec![
+                Field::new(
+                    format_state_name(args.name, "count"),
+                    DataType::UInt64,
+                    true,
+                ),
+                Field::new(
+                    format_state_name(args.name, "sum"),
+                    args.input_fields[0].data_type().clone(),
+                    true,
+                ),
+            ]
+            .into_iter()
+            .map(Arc::new)
+            .collect())
+        }


This is my main contribution in addition to the changes from the original PR, to fix that error with the differing field counts. I wonder if there's a better way to architect this, since wasn't clearly obvious that this is related to the state of the accumulator

I don't really understand this question -- the PR's code looks good to me

Sorry, what I meant was that in accumulators we can return the state() like so:

datafusion/datafusion/functions-aggregate/src/sum.rs

Lines 432 to 450 in be3842b

fn state(&mut self) -> Result<Vec<ScalarValue>> {

// 1. Stores aggregate state in `ScalarValue::List`

// 2. Constructs `ScalarValue::List` state from distinct numeric stored in hash set

let state_out = {

let distinct_values = self

.values

.iter()

.map(|value| {

ScalarValue::new_primitive::<T>(Some(value.0), &self.data_type)

})

.collect::<Result<Vec<_>>>()?;

vec![ScalarValue::List(ScalarValue::new_list_nullable(

&distinct_values,

&self.data_type,

))]

};

Ok(state_out)

}

However this must align with state_fields() of the parent aggregate UDF:

datafusion/datafusion/functions-aggregate/src/sum.rs

Lines 204 to 221 in be3842b

fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {

if args.is_distinct {

Ok(vec![Field::new_list(

format_state_name(args.name, "sum distinct"),

// See COMMENTS.md to understand why nullable is set to true

Field::new_list_field(args.return_type().clone(), true),

false,

)

.into()])

} else {

Ok(vec![Field::new(

format_state_name(args.name, "sum"),

args.return_type().clone(),

true,

)

.into()])

}

}

But this isn't clearly obvious at compile time, and during runtime we only hit this issue for certain test cases (for this distinct avg PR). So I was wondering if there was a better way to enforce this at compile time. Hope that clears it up.

Jefffrey · 2025-08-20T07:35:04Z

Looks like still getting same error for some of the extended tests:

# Datafusion - Datafusion expected results:
query error DataFusion error: Arrow error: Invalid argument error: number of columns\(5\) must match number of fields\(4\) in schema
SELECT ALL - AVG ( ALL + col0 ) AS col1 FROM tab0 GROUP BY col0 HAVING + AVG ( DISTINCT - col2 ) IS NULL

Will look into this

Edit: seems to be related to group accumulator support

…test case

Jefffrey · 2025-08-20T08:57:34Z

datafusion/functions-aggregate/src/average.rs

            args.return_field.data_type(),
            DataType::Float64 | DataType::Decimal128(_, _) | DataType::Duration(_)
-        )
+        ) && !args.is_distinct


Similar to how sum handles it:

datafusion/datafusion/functions-aggregate/src/sum.rs

Lines 223 to 225 in 8c58f53

fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {

!args.is_distinct

}

Jefffrey · 2025-08-20T09:44:01Z

Looks like still getting same error for some of the extended tests:
# Datafusion - Datafusion expected results:
query error DataFusion error: Arrow error: Invalid argument error: number of columns\(5\) must match number of fields\(4\) in schema
SELECT ALL - AVG ( ALL + col0 ) AS col1 FROM tab0 GROUP BY col0 HAVING + AVG ( DISTINCT - col2 ) IS NULL
Will look into this

Edit: seems to be related to group accumulator support

Fixed by bc121fb

Omega359 · 2025-08-22T19:40:05Z

This will I assume require regenerating the extended slt files in datafusion-testing?

alamb

Thank you @Jefffrey -- I think the Pr looks quite good. I think it should just have a few more tests, but otherwise 👌

datafusion/functions-aggregate/src/average.rs

alamb · 2025-08-22T19:57:00Z

datafusion/sqllogictest/test_files/aggregate.slt

+  (1, 1),
+  (2, 2),
+  (3, 3),
+  (4, 4),


Could you update this test so:

The input isn't in order

Add a test for floating point values

Test for an input that includes at least one null value

the values in b are different than the values in b

Will work on adding these cases 👍

Addressed those points in this commit: 3abb4b7

alamb · 2025-08-22T19:57:35Z

datafusion/functions-aggregate/src/average.rs

-        Ok(vec![
-            Field::new(
-                format_state_name(args.name, "count"),
-                DataType::UInt64,
-                true,
-            ),
-            Field::new(
-                format_state_name(args.name, "sum"),
-                args.input_fields[0].data_type().clone(),
-                true,
-            ),
-        ]
-        .into_iter()
-        .map(Arc::new)
-        .collect())
+        if args.is_distinct {
+            // Copied from datafusion_functions_aggregate::sum::Sum::state_fields
+            // since the accumulator uses DistinctSumAccumulator internally.
+            Ok(vec![Field::new_list(
+                format_state_name(args.name, "sum distinct"),
+                Field::new_list_field(args.return_type().clone(), true),
+                false,
+            )
+            .into()])
+        } else {
+            Ok(vec![
+                Field::new(
+                    format_state_name(args.name, "count"),
+                    DataType::UInt64,
+                    true,
+                ),
+                Field::new(
+                    format_state_name(args.name, "sum"),
+                    args.input_fields[0].data_type().clone(),
+                    true,
+                ),
+            ]
+            .into_iter()
+            .map(Arc::new)
+            .collect())
+        }


I don't really understand this question -- the PR's code looks good to me

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Jefffrey · 2025-08-23T04:19:23Z

This will I assume require regenerating the extended slt files in datafusion-testing?

Yep, refer to apache/datafusion-testing#11

alamb · 2025-08-23T11:28:50Z

I think all we need now is to merge apache/datafusion-testing#11 and then update the datafusion-testing pin on this PR.

I am testing locally with the changes from apache/datafusion-testing#11 using:

INCLUDE_SQLITE=true cargo test --profile release-nonlto --test sqllogictests

alamb · 2025-08-23T11:29:13Z

datafusion/sqllogictest/test_files/aggregate.slt

 SELECT array_agg(a_varchar order by a_varchar) WITHIN GROUP (ORDER BY a_varchar)
 FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar);
+
+# distinct average


alamb · 2025-08-23T11:35:54Z

   Finished `release-nonlto` profile [optimized] target(s) in 2m 09s
     Running bin/sqllogictests.rs (target/release-nonlto/deps/sqllogictests-78c77c9b80ab3916)
Completed 942 test files in 4 minutes
andrewlamb@Andrews-MacBook-Pro-3:~/Software/datafusion$

👍

alamb · 2025-08-23T11:37:24Z

I updated the pin and merged up from main

Jefffrey · 2025-08-23T11:39:41Z

I updated the pin and merged up from main

Cheers 👍

alamb · 2025-08-24T10:35:28Z

🚀

qazxcdswe123 and others added 15 commits March 25, 2025 17:01

chore: mv DistinctSumAccumulator to common

01fa485

feat: add avg distinct support for float64 type

3eb166b

chore: fmt

6ae50ab

refactor: update import for DataType in Float64DistinctAvgAccumulator…

4a8868d

… and remove unused sum_distinct module

Merge remote-tracking branch 'apache/main' into avg-distinct-float

ec2a394

update

1fa2db8

Revert "update"

b0b4e43

This reverts commit 1fa2db8.

Merge branch 'main' into pr_15413

ae06e74

Fix submodules

e1f16e4

Fix broken import

6c952bf

Fix state_fields for distinct avg

af66e5f

Merge branch 'main' into pr_15413

bfcb019

Update SLT test files

228679b

Added null case for SLT test

7b7ffbf

Merge branch 'main' into pr_15413

194fa34

github-actions bot added sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels Aug 20, 2025

Jefffrey mentioned this pull request Aug 20, 2025

Regenerate tests due to new avg(distinct) support for floats apache/datafusion-testing#11

Merged

Jefffrey commented Aug 20, 2025

View reviewed changes

Jefffrey marked this pull request as ready for review August 20, 2025 06:41

Jefffrey marked this pull request as draft August 20, 2025 08:44

Disable group accumulator support for avg(distinct) and add group by …

bc121fb

…test case

Jefffrey commented Aug 20, 2025

View reviewed changes

Jefffrey marked this pull request as ready for review August 20, 2025 09:44

alamb approved these changes Aug 22, 2025

View reviewed changes

Fix state field name for avg distinct

b9c458f

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Mix up tests

3abb4b7

alamb reviewed Aug 23, 2025

View reviewed changes

alamb added 2 commits August 23, 2025 07:36

Update datafusion-tesitng pin

2fd3a33

Merge remote-tracking branch 'apache/main' into pr_15413

1d450c7

Jefffrey merged commit 241e47d into apache:main Aug 24, 2025
27 checks passed

Jefffrey deleted the pr_15413 branch August 24, 2025 01:03

Jefffrey mentioned this pull request Sep 14, 2025

avg(distinct) support for decimal types #17560

Merged

Jefffrey mentioned this pull request Oct 19, 2025

No way to get the schema for sliding accumulator state #14701

Open

	fn state(&mut self) -> Result<Vec<ScalarValue>> {
	// 1. Stores aggregate state in `ScalarValue::List`
	// 2. Constructs `ScalarValue::List` state from distinct numeric stored in hash set
	let state_out = {
	let distinct_values = self
	.values
	.iter()
	.map(\|value\| {
	ScalarValue::new_primitive::<T>(Some(value.0), &self.data_type)
	})
	.collect::<Result<Vec<_>>>()?;

	vec![ScalarValue::List(ScalarValue::new_list_nullable(
	&distinct_values,
	&self.data_type,
	))]
	};
	Ok(state_out)
	}

	fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
	if args.is_distinct {
	Ok(vec![Field::new_list(
	format_state_name(args.name, "sum distinct"),
	// See COMMENTS.md to understand why nullable is set to true
	Field::new_list_field(args.return_type().clone(), true),
	false,
	)
	.into()])
	} else {
	Ok(vec![Field::new(
	format_state_name(args.name, "sum"),
	args.return_type().clone(),
	true,
	)
	.into()])
	}
	}

	fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
	!args.is_distinct
	}

Support avg(distinct) for float64 type #17255

Support avg(distinct) for float64 type #17255

Uh oh!

Conversation

Jefffrey commented Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey commented Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey commented Aug 20, 2025

Uh oh!

Omega359 commented Aug 22, 2025

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey commented Aug 23, 2025

Uh oh!

alamb commented Aug 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Aug 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alamb commented Aug 23, 2025

Uh oh!

Jefffrey commented Aug 23, 2025

Uh oh!

Uh oh!

alamb commented Aug 24, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Support `avg(distinct)` for `float64` type #17255

Support `avg(distinct)` for `float64` type #17255

Jefffrey commented Aug 20, 2025 •

edited

Loading

Jefffrey commented Aug 20, 2025 •

edited

Loading

alamb commented Aug 23, 2025 •

edited

Loading