From 715f5066019d80637dfefcf8e31d6e8c397d58f1 Mon Sep 17 00:00:00 2001 From: Mason Hall Date: Tue, 22 Sep 2026 15:54:31 -0400 Subject: [PATCH 1/5] handle grouped accumulator state in ungrouped approx distinct --- .../src/approx_distinct.rs | 36 ++++++++++++++++--- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 74f116cf48725..2b447f2b9b317 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -126,6 +126,36 @@ impl Accumulator for ApproxDistinctBitmapWrapper { } } +fn merge_serialized(hll: &mut HyperLogLog, bytes: &[u8]) -> Result<()> { + if bytes.is_empty() { + return Ok(()); + } + if bytes.len() == NUM_REGISTERS { + let other: HyperLogLog = bytes.try_into()?; + hll.merge(&other); + } else { + if !bytes.len().is_multiple_of(size_of::()) { + return internal_err!( + "approx_distinct: malformed sparse state: length {} is not a multiple of {}", + bytes.len(), + size_of::() + ); + } + if bytes.len() > SPARSE_LIMIT * size_of::() { + return internal_err!( + "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}", + bytes.len(), + SPARSE_LIMIT * size_of::() + ); + } + for chunk in bytes.chunks_exact(size_of::()) { + let h = u64::from_le_bytes(chunk.try_into().unwrap()); + hll.add_hashed(h); + } + } + Ok(()) +} + #[derive(Debug)] struct HLLAccumulator { hll: HyperLogLog, @@ -172,8 +202,7 @@ impl Accumulator for HLLAccumulator { let v = v.ok_or_else(|| { internal_datafusion_err!("Impossibly got empty binary array from states") })?; - let other = v.try_into()?; - self.hll.merge(&other); + merge_serialized(&mut self.hll, v)?; } Ok(()) } @@ -232,8 +261,7 @@ where let v = v.ok_or_else(|| { internal_datafusion_err!("Impossibly got empty binary array from states") })?; - let other = v.try_into()?; - self.hll.merge(&other); + merge_serialized(&mut self.hll, v)?; } Ok(()) } From c9e145dc415f15160f6fa19877d3e2e32b0f6dd7 Mon Sep 17 00:00:00 2001 From: Mason Hall Date: Tue, 22 Sep 2026 16:07:57 -0400 Subject: [PATCH 2/5] add tests, fmt --- .../src/approx_distinct.rs | 62 ++++++++++++++++++- 1 file changed, 60 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 2b447f2b9b317..de0af2dd452a1 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -126,7 +126,10 @@ impl Accumulator for ApproxDistinctBitmapWrapper { } } -fn merge_serialized(hll: &mut HyperLogLog, bytes: &[u8]) -> Result<()> { +fn merge_serialized( + hll: &mut HyperLogLog, + bytes: &[u8], +) -> Result<()> { if bytes.is_empty() { return Ok(()); } @@ -1039,7 +1042,7 @@ mod tests { use arrow::array::{ AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, - IntervalYearMonthArray, StringViewArray, + IntervalYearMonthArray, StringArray, StringViewArray, }; use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256}; use std::sync::Arc; @@ -1383,6 +1386,61 @@ mod tests { ); assert_eq!(distinct_count(&mut acc_single), 3); } + + /// Grouped state (empty, sparse and dense rows) must merge into the + /// ungrouped accumulator. + #[test] + fn hll_acc_merges_grouped_state() { + let sparse = 10; + let dense = SPARSE_LIMIT * 4; + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..sparse + dense).map(|i| format!("value-{i}")), + )); + let group_indices: Vec = (0..sparse + dense) + .map(|i| if i < sparse { 1 } else { 2 }) + .collect(); + + let mut grouped = HllGroupsAccumulator::new(); + grouped + .update_batch(std::slice::from_ref(&values), &group_indices, None, 3) + .unwrap(); + let state = grouped.state(EmitTo::All).unwrap(); + + let mut direct = HLLAccumulator::new(); + direct.update_batch(std::slice::from_ref(&values)).unwrap(); + + let mut merged = HLLAccumulator::new(); + merged.merge_batch(&state).unwrap(); + + assert_eq!(distinct_count(&mut merged), distinct_count(&mut direct)); + } + + /// Grouped state (empty, sparse and dense rows) must merge into the + /// ungrouped numeric accumulator. + #[test] + fn numeric_acc_merges_grouped_state() { + let sparse = 10; + let dense = SPARSE_LIMIT * 4; + let values: ArrayRef = + Arc::new(Int64Array::from_iter_values(0..(sparse + dense) as i64)); + let group_indices: Vec = (0..sparse + dense) + .map(|i| if i < sparse { 1 } else { 2 }) + .collect(); + + let mut grouped = HllGroupsAccumulator::new(); + grouped + .update_batch(std::slice::from_ref(&values), &group_indices, None, 3) + .unwrap(); + let state = grouped.state(EmitTo::All).unwrap(); + + let mut direct = NumericHLLAccumulator::::new(); + direct.update_batch(std::slice::from_ref(&values)).unwrap(); + + let mut merged = NumericHLLAccumulator::::new(); + merged.merge_batch(&state).unwrap(); + + assert_eq!(merged.evaluate().unwrap(), direct.evaluate().unwrap()); + } } fn h(v: u64) -> u64 { From 21b7c0cea2ffcdee294d66834d75299e48014c07 Mon Sep 17 00:00:00 2001 From: Mason Hall Date: Wed, 23 Sep 2026 15:42:56 -0400 Subject: [PATCH 3/5] fix clippy --- datafusion/functions-aggregate/src/approx_distinct.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index de0af2dd452a1..a47b46682cc71 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -151,8 +151,8 @@ fn merge_serialized( SPARSE_LIMIT * size_of::() ); } - for chunk in bytes.chunks_exact(size_of::()) { - let h = u64::from_le_bytes(chunk.try_into().unwrap()); + for chunk in bytes.as_chunks::<{ size_of::() }>().0 { + let h = u64::from_le_bytes(*chunk); hll.add_hashed(h); } } From abf16b73ed38df26704df2df665749d9434eb47b Mon Sep 17 00:00:00 2001 From: Mason Hall Date: Tue, 29 Sep 2026 13:50:29 -0400 Subject: [PATCH 4/5] factor out duplicate code --- .../src/approx_distinct.rs | 115 ++++++++---------- 1 file changed, 53 insertions(+), 62 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index a47b46682cc71..3ce2e5005705b 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -126,34 +126,60 @@ impl Accumulator for ApproxDistinctBitmapWrapper { } } -fn merge_serialized( - hll: &mut HyperLogLog, - bytes: &[u8], -) -> Result<()> { - if bytes.is_empty() { - return Ok(()); - } - if bytes.len() == NUM_REGISTERS { - let other: HyperLogLog = bytes.try_into()?; - hll.merge(&other); - } else { - if !bytes.len().is_multiple_of(size_of::()) { +/// A validated, zero-copy view of a serialized `approx_distinct` partial state, +/// as produced by [`GroupHll::serialize`] or by the per-group [`Accumulator`]s. +enum SerializedHll<'a> { + /// The raw [`NUM_REGISTERS`] registers of a dense sketch. + Dense(&'a [u8; NUM_REGISTERS]), + /// Little-endian hashes of at most [`SPARSE_LIMIT`] distinct values. An + /// empty state decodes as an empty sparse state. + Sparse(&'a [[u8; size_of::()]]), +} + +impl<'a> SerializedHll<'a> { + fn decode(bytes: &'a [u8]) -> Result { + if let Ok(registers) = <&[u8; NUM_REGISTERS]>::try_from(bytes) { + return Ok(Self::Dense(registers)); + } + let (chunks, rest) = bytes.as_chunks::<{ size_of::() }>(); + if !rest.is_empty() { return internal_err!( "approx_distinct: malformed sparse state: length {} is not a multiple of {}", bytes.len(), size_of::() ); } - if bytes.len() > SPARSE_LIMIT * size_of::() { + if chunks.len() > SPARSE_LIMIT { return internal_err!( "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}", bytes.len(), SPARSE_LIMIT * size_of::() ); } - for chunk in bytes.as_chunks::<{ size_of::() }>().0 { - let h = u64::from_le_bytes(*chunk); - hll.add_hashed(h); + Ok(Self::Sparse(chunks)) + } +} + +/// Merge the serialized partial states in `states` into `hll`. +fn merge_states( + hll: &mut HyperLogLog, + states: &[ArrayRef], +) -> Result<()> { + assert_eq!(1, states.len(), "expect only 1 element in the states"); + let binary_array = downcast_value!(states[0], BinaryArray); + for v in binary_array.iter() { + let v = v.ok_or_else(|| { + internal_datafusion_err!("Impossibly got empty binary array from states") + })?; + match SerializedHll::decode(v)? { + SerializedHll::Dense(registers) => { + hll.merge(&HyperLogLog::new_with_registers(*registers)); + } + SerializedHll::Sparse(chunks) => { + for chunk in chunks { + hll.add_hashed(u64::from_le_bytes(*chunk)); + } + } } } Ok(()) @@ -199,15 +225,7 @@ impl Accumulator for HLLAccumulator { } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - assert_eq!(1, states.len(), "expect only 1 element in the states"); - let binary_array = downcast_value!(states[0], BinaryArray); - for v in binary_array.iter() { - let v = v.ok_or_else(|| { - internal_datafusion_err!("Impossibly got empty binary array from states") - })?; - merge_serialized(&mut self.hll, v)?; - } - Ok(()) + merge_states(&mut self.hll, states) } fn state(&mut self) -> Result> { @@ -258,15 +276,7 @@ where } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - assert_eq!(1, states.len(), "expect only 1 element in the states"); - let binary_array = downcast_value!(states[0], BinaryArray); - for v in binary_array.iter() { - let v = v.ok_or_else(|| { - internal_datafusion_err!("Impossibly got empty binary array from states") - })?; - merge_serialized(&mut self.hll, v)?; - } - Ok(()) + merge_states(&mut self.hll, states) } fn state(&mut self) -> Result> { @@ -366,36 +376,17 @@ impl GroupHll { } } - /// Merge a serialized state (produced by [`Self::serialize`] or by the - /// per-group [`Accumulator`]) into this sketch. + /// Merge a serialized state (see [`SerializedHll`]) into this sketch. fn merge_serialized(&mut self, bytes: &[u8]) -> Result { - if bytes.is_empty() { - return Ok(0); - } - if bytes.len() == NUM_REGISTERS { - let other: HyperLogLog = bytes.try_into()?; - Ok(self.merge_dense(&other)) - } else { - if !bytes.len().is_multiple_of(size_of::()) { - return internal_err!( - "approx_distinct: malformed sparse state: length {} is not a multiple of {}", - bytes.len(), - size_of::() - ); - } - if bytes.len() > SPARSE_LIMIT * size_of::() { - return internal_err!( - "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}", - bytes.len(), - SPARSE_LIMIT * size_of::() - ); - } - let mut delta = 0; - for chunk in bytes.as_chunks::<{ size_of::() }>().0 { - delta += self.add_hash(u64::from_le_bytes(*chunk)); + Ok(match SerializedHll::decode(bytes)? { + SerializedHll::Dense(registers) => { + self.merge_dense(&HyperLogLog::new_with_registers(*registers)) } - Ok(delta) - } + SerializedHll::Sparse(chunks) => chunks + .iter() + .map(|chunk| self.add_hash(u64::from_le_bytes(*chunk))) + .sum(), + }) } /// Merge a dense sketch into this one, promoting to dense if necessary. From be2076ee1f01ee06a51b46c5b3fa137f2c76929f Mon Sep 17 00:00:00 2001 From: Mason Hall Date: Tue, 6 Oct 2026 11:43:44 -0400 Subject: [PATCH 5/5] remove tests --- .../src/approx_distinct.rs | 57 +------------------ 1 file changed, 1 insertion(+), 56 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 3ce2e5005705b..dc9270ebe41df 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -1033,7 +1033,7 @@ mod tests { use arrow::array::{ AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, - IntervalYearMonthArray, StringArray, StringViewArray, + IntervalYearMonthArray, StringViewArray, }; use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256}; use std::sync::Arc; @@ -1377,61 +1377,6 @@ mod tests { ); assert_eq!(distinct_count(&mut acc_single), 3); } - - /// Grouped state (empty, sparse and dense rows) must merge into the - /// ungrouped accumulator. - #[test] - fn hll_acc_merges_grouped_state() { - let sparse = 10; - let dense = SPARSE_LIMIT * 4; - let values: ArrayRef = Arc::new(StringArray::from_iter_values( - (0..sparse + dense).map(|i| format!("value-{i}")), - )); - let group_indices: Vec = (0..sparse + dense) - .map(|i| if i < sparse { 1 } else { 2 }) - .collect(); - - let mut grouped = HllGroupsAccumulator::new(); - grouped - .update_batch(std::slice::from_ref(&values), &group_indices, None, 3) - .unwrap(); - let state = grouped.state(EmitTo::All).unwrap(); - - let mut direct = HLLAccumulator::new(); - direct.update_batch(std::slice::from_ref(&values)).unwrap(); - - let mut merged = HLLAccumulator::new(); - merged.merge_batch(&state).unwrap(); - - assert_eq!(distinct_count(&mut merged), distinct_count(&mut direct)); - } - - /// Grouped state (empty, sparse and dense rows) must merge into the - /// ungrouped numeric accumulator. - #[test] - fn numeric_acc_merges_grouped_state() { - let sparse = 10; - let dense = SPARSE_LIMIT * 4; - let values: ArrayRef = - Arc::new(Int64Array::from_iter_values(0..(sparse + dense) as i64)); - let group_indices: Vec = (0..sparse + dense) - .map(|i| if i < sparse { 1 } else { 2 }) - .collect(); - - let mut grouped = HllGroupsAccumulator::new(); - grouped - .update_batch(std::slice::from_ref(&values), &group_indices, None, 3) - .unwrap(); - let state = grouped.state(EmitTo::All).unwrap(); - - let mut direct = NumericHLLAccumulator::::new(); - direct.update_batch(std::slice::from_ref(&values)).unwrap(); - - let mut merged = NumericHLLAccumulator::::new(); - merged.merge_batch(&state).unwrap(); - - assert_eq!(merged.evaluate().unwrap(), direct.evaluate().unwrap()); - } } fn h(v: u64) -> u64 {