Skip to content

Commit 8f0280b

Browse files
authored
Document and enforce compatibility between an aggregate's accumulators (#25710)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #25660. ## Rationale for this change See #25660. This change only documents/enforces this requirement for built-in aggregate UDF functions. Anyone who implements their own aggregate UDFs do not need to adhere to this requirement since vanilla datafusion doesn't require it. This does add a public test that user's can use to make sure this requirement is satisfied if they desire. ## What changes are included in this PR? This change adds: - Internal documentation explaining this requirement and how to check for it using the provided test function. - `datafusion_functions_aggregate::testing::check_state_compatibility()`: A generic test that identifies cases where the `Accumulator` and `GroupsAccumulator` of an `AggregateUDF` differ. This is exposed as a public function that can be used to test this requirement on `AggregateUDFImpl`s implemented by downstream consumers. The test is run on all built-in default aggregation functions. - Fixes for `approx_distinct()` and `avg()`, which both violated this requirement - This takes the changes from #25659 to fix `approx_distinct()` - The `Accumulator` for `avg()` would return `0` for the count if it hasn't seen any non-null values. The `GroupsAccumulator` would return `null` instead. This change updates the `Accumulator`s to also return `null` for the count in this case. ## What is the testing strategy for this PR? This adds a test to enforce this new requirement in a generic way. <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code Briefly describe how this PR is tested, and point to the specific tests you added. For example: 'This new feature is covered by the `sqllogictest` cases added in `foo.slt`'. If this PR does not add tests, explain why. For example, if the change is already covered by existing tests, please mention it. You should also check the `codecov` bot reply on this PR to confirm the changed code is exercised. --> ## Are there any user-facing changes? This changes the intermediate state of the `avg()` accumulator to match the `avg()` groups accumulator, which some users may be affected by. <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. If there are any breaking changes to public APIs, please add the `api change` label. -->
1 parent 9dec775 commit 8f0280b

7 files changed

Lines changed: 909 additions & 64 deletions

File tree

‎datafusion/expr-common/src/accumulator.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ impl Drop for AggregateMetricTimer<'_> {
9797
/// `Accumulator`s are stateful objects that implement a single group. They
9898
/// aggregate values from multiple rows together into a final output aggregate.
9999
///
100-
/// [`GroupsAccumulator]` is an additional more performant (but also complex) API
100+
/// [`GroupsAccumulator`] is an additional more performant (but also complex) API
101101
/// that manages state for multiple groups at once.
102102
///
103103
/// An accumulator knows how to:
@@ -118,6 +118,7 @@ impl Drop for AggregateMetricTimer<'_> {
118118
/// [`state`]: Self::state
119119
/// [`evaluate`]: Self::evaluate
120120
/// [`merge_batch`]: Self::merge_batch
121+
/// [`GroupsAccumulator`]: crate::groups_accumulator::GroupsAccumulator
121122
/// [window function]: https://en.wikipedia.org/wiki/Window_function_(SQL)
122123
pub trait Accumulator: Send + Sync + Debug + std::any::Any {
123124
/// Supplies optional metrics owned by this aggregate expression.

‎datafusion/functions-aggregate/src/approx_distinct.rs‎

Lines changed: 70 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,65 @@ impl<A: Accumulator> Accumulator for ApproxDistinctBitmapWrapper<A> {
126126
}
127127
}
128128

129+
/// A validated, zero-copy view of a serialized `approx_distinct` partial state,
130+
/// as produced by [`GroupHll::serialize`] or by the per-group [`Accumulator`]s.
131+
enum SerializedHll<'a> {
132+
/// The raw [`NUM_REGISTERS`] registers of a dense sketch.
133+
Dense(&'a [u8; NUM_REGISTERS]),
134+
/// Little-endian hashes of at most [`SPARSE_LIMIT`] distinct values. An
135+
/// empty state decodes as an empty sparse state.
136+
Sparse(&'a [[u8; size_of::<u64>()]]),
137+
}
138+
139+
impl<'a> SerializedHll<'a> {
140+
fn decode(bytes: &'a [u8]) -> Result<Self> {
141+
if let Ok(registers) = <&[u8; NUM_REGISTERS]>::try_from(bytes) {
142+
return Ok(Self::Dense(registers));
143+
}
144+
let (chunks, rest) = bytes.as_chunks::<{ size_of::<u64>() }>();
145+
if !rest.is_empty() {
146+
return internal_err!(
147+
"approx_distinct: malformed sparse state: length {} is not a multiple of {}",
148+
bytes.len(),
149+
size_of::<u64>()
150+
);
151+
}
152+
if chunks.len() > SPARSE_LIMIT {
153+
return internal_err!(
154+
"approx_distinct: malformed sparse state: length {} exceeds sparse limit {}",
155+
bytes.len(),
156+
SPARSE_LIMIT * size_of::<u64>()
157+
);
158+
}
159+
Ok(Self::Sparse(chunks))
160+
}
161+
}
162+
163+
/// Merge the serialized partial states in `states` into `hll`.
164+
fn merge_states<T: Hash + ?Sized>(
165+
hll: &mut HyperLogLog<T>,
166+
states: &[ArrayRef],
167+
) -> Result<()> {
168+
assert_eq!(1, states.len(), "expect only 1 element in the states");
169+
let binary_array = downcast_value!(states[0], BinaryArray);
170+
for v in binary_array.iter() {
171+
let v = v.ok_or_else(|| {
172+
internal_datafusion_err!("Impossibly got empty binary array from states")
173+
})?;
174+
match SerializedHll::decode(v)? {
175+
SerializedHll::Dense(registers) => {
176+
hll.merge(&HyperLogLog::new_with_registers(*registers));
177+
}
178+
SerializedHll::Sparse(chunks) => {
179+
for chunk in chunks {
180+
hll.add_hashed(u64::from_le_bytes(*chunk));
181+
}
182+
}
183+
}
184+
}
185+
Ok(())
186+
}
187+
129188
#[derive(Debug)]
130189
struct HLLAccumulator {
131190
hll: HyperLogLog<u8>,
@@ -166,16 +225,7 @@ impl Accumulator for HLLAccumulator {
166225
}
167226

168227
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
169-
assert_eq!(1, states.len(), "expect only 1 element in the states");
170-
let binary_array = downcast_value!(states[0], BinaryArray);
171-
for v in binary_array.iter() {
172-
let v = v.ok_or_else(|| {
173-
internal_datafusion_err!("Impossibly got empty binary array from states")
174-
})?;
175-
let other = v.try_into()?;
176-
self.hll.merge(&other);
177-
}
178-
Ok(())
228+
merge_states(&mut self.hll, states)
179229
}
180230

181231
fn state(&mut self) -> Result<Vec<ScalarValue>> {
@@ -226,16 +276,7 @@ where
226276
}
227277

228278
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
229-
assert_eq!(1, states.len(), "expect only 1 element in the states");
230-
let binary_array = downcast_value!(states[0], BinaryArray);
231-
for v in binary_array.iter() {
232-
let v = v.ok_or_else(|| {
233-
internal_datafusion_err!("Impossibly got empty binary array from states")
234-
})?;
235-
let other = v.try_into()?;
236-
self.hll.merge(&other);
237-
}
238-
Ok(())
279+
merge_states(&mut self.hll, states)
239280
}
240281

241282
fn state(&mut self) -> Result<Vec<ScalarValue>> {
@@ -335,36 +376,17 @@ impl GroupHll {
335376
}
336377
}
337378

338-
/// Merge a serialized state (produced by [`Self::serialize`] or by the
339-
/// per-group [`Accumulator`]) into this sketch.
379+
/// Merge a serialized state (see [`SerializedHll`]) into this sketch.
340380
fn merge_serialized(&mut self, bytes: &[u8]) -> Result<isize> {
341-
if bytes.is_empty() {
342-
return Ok(0);
343-
}
344-
if bytes.len() == NUM_REGISTERS {
345-
let other: HyperLogLog<u8> = bytes.try_into()?;
346-
Ok(self.merge_dense(&other))
347-
} else {
348-
if !bytes.len().is_multiple_of(size_of::<u64>()) {
349-
return internal_err!(
350-
"approx_distinct: malformed sparse state: length {} is not a multiple of {}",
351-
bytes.len(),
352-
size_of::<u64>()
353-
);
381+
Ok(match SerializedHll::decode(bytes)? {
382+
SerializedHll::Dense(registers) => {
383+
self.merge_dense(&HyperLogLog::new_with_registers(*registers))
354384
}
355-
if bytes.len() > SPARSE_LIMIT * size_of::<u64>() {
356-
return internal_err!(
357-
"approx_distinct: malformed sparse state: length {} exceeds sparse limit {}",
358-
bytes.len(),
359-
SPARSE_LIMIT * size_of::<u64>()
360-
);
361-
}
362-
let mut delta = 0;
363-
for chunk in bytes.as_chunks::<{ size_of::<u64>() }>().0 {
364-
delta += self.add_hash(u64::from_le_bytes(*chunk));
365-
}
366-
Ok(delta)
367-
}
385+
SerializedHll::Sparse(chunks) => chunks
386+
.iter()
387+
.map(|chunk| self.add_hash(u64::from_le_bytes(*chunk)))
388+
.sum(),
389+
})
368390
}
369391

370392
/// Merge a dense sketch into this one, promoting to dense if necessary.

‎datafusion/functions-aggregate/src/average.rs‎

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -653,10 +653,14 @@ impl Accumulator for AvgAccumulator {
653653
}
654654

655655
fn state(&mut self) -> Result<Vec<ScalarValue>> {
656-
Ok(vec![
657-
ScalarValue::from(self.count),
658-
ScalarValue::Float64(self.sum),
659-
])
656+
// With no non-NULL values both fields are NULL, matching
657+
// `AvgGroupsAccumulator::state`
658+
let (count, sum) = if self.count == 0 {
659+
(None, None)
660+
} else {
661+
(Some(self.count), self.sum)
662+
};
663+
Ok(vec![ScalarValue::UInt64(count), ScalarValue::Float64(sum)])
660664
}
661665

662666
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
@@ -833,9 +837,16 @@ where
833837
}
834838

835839
fn state(&mut self) -> Result<Vec<ScalarValue>> {
840+
// With no non-NULL values both fields are NULL, matching
841+
// `AvgGroupsAccumulator::state`
842+
let (count, sum) = if self.count == 0 {
843+
(None, None)
844+
} else {
845+
(Some(self.count), self.sum)
846+
};
836847
Ok(vec![
837-
ScalarValue::from(self.count),
838-
ScalarValue::new_primitive::<S>(self.sum, &self.sum_data_type)?,
848+
ScalarValue::UInt64(count),
849+
ScalarValue::new_primitive::<S>(sum, &self.sum_data_type)?,
839850
])
840851
}
841852

@@ -916,14 +927,21 @@ impl Accumulator for DurationAvgAccumulator {
916927
}
917928

918929
fn state(&mut self) -> Result<Vec<ScalarValue>> {
930+
// With no non-NULL values both fields are NULL, matching
931+
// `AvgGroupsAccumulator::state`
932+
let (count, sum) = if self.count == 0 {
933+
(None, None)
934+
} else {
935+
(Some(self.count), self.sum)
936+
};
919937
let duration_value = match self.time_unit {
920-
TimeUnit::Second => ScalarValue::DurationSecond(self.sum),
921-
TimeUnit::Millisecond => ScalarValue::DurationMillisecond(self.sum),
922-
TimeUnit::Microsecond => ScalarValue::DurationMicrosecond(self.sum),
923-
TimeUnit::Nanosecond => ScalarValue::DurationNanosecond(self.sum),
938+
TimeUnit::Second => ScalarValue::DurationSecond(sum),
939+
TimeUnit::Millisecond => ScalarValue::DurationMillisecond(sum),
940+
TimeUnit::Microsecond => ScalarValue::DurationMicrosecond(sum),
941+
TimeUnit::Nanosecond => ScalarValue::DurationNanosecond(sum),
924942
};
925943

926-
Ok(vec![ScalarValue::from(self.count), duration_value])
944+
Ok(vec![ScalarValue::UInt64(count), duration_value])
927945
}
928946

929947
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {

‎datafusion/functions-aggregate/src/hyperloglog.rs‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
//!
3535
//! This module also borrows some code structure from [pdatastructs.rs](https://github.com/crepererum/pdatastructs.rs/blob/3997ed50f6b6871c9e53c4c5e0f48f431405fc63/src/hyperloglog.rs).
3636
37+
#[cfg(not(feature = "force_hash_collisions"))]
3738
use std::hash::BuildHasher;
3839
use std::hash::Hash;
3940
use std::marker::PhantomData;
@@ -92,11 +93,20 @@ where
9293

9394
/// The HLL hash state is shared through `datafusion_common::hash_utils`
9495
/// so sketches remain compatible across accumulators.
96+
#[cfg(not(feature = "force_hash_collisions"))]
9597
#[inline]
9698
fn hash_value(&self, obj: &T) -> u64 {
9799
HLL_HASH_STATE.hash_one(obj)
98100
}
99101

102+
/// Test version of `hash_value` that forces all hashes to collide to zero,
103+
/// matching `create_hashes`.
104+
#[cfg(feature = "force_hash_collisions")]
105+
#[inline]
106+
fn hash_value(&self, _obj: &T) -> u64 {
107+
0
108+
}
109+
100110
/// Adds an element to the HyperLogLog.
101111
pub fn add(&mut self, obj: &T) {
102112
let hash = self.hash_value(obj);
@@ -285,7 +295,9 @@ where
285295
}
286296
}
287297

288-
#[cfg(test)]
298+
/// These tests check cardinality estimates, which collapse to 1 under
299+
/// `force_hash_collisions` because every value hashes to zero.
300+
#[cfg(all(test, not(feature = "force_hash_collisions")))]
289301
mod tests {
290302
use super::{HyperLogLog, NUM_REGISTERS};
291303

0 commit comments

Comments
 (0)