Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 48 additions & 11 deletions datafusion/functions-aggregate/src/percentile_cont.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator;
use crate::min_max::{max_udaf, min_udaf};
use datafusion_common::{
Result, ScalarValue, exec_datafusion_err, internal_datafusion_err,
utils::take_function_args,
utils::{SingleRowListArrayBuilder, take_function_args},
};
use datafusion_expr::utils::format_state_name;
use datafusion_expr::{
Expand All @@ -54,7 +54,7 @@ use datafusion_expr::{
};
use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate;
use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filtered_null_mask;
use datafusion_functions_aggregate_common::utils::{GenericDistinctBuffer, Hashable};
use datafusion_functions_aggregate_common::utils::Hashable;
use datafusion_macros::user_doc;

use crate::utils::validate_percentile_expr;
Expand Down Expand Up @@ -662,16 +662,24 @@ where
}
}

/// Sliding-window–capable accumulator for `percentile_cont(DISTINCT ...)`.
///
/// Distinct values are tracked with a per-value multiplicity count (how many
/// rows currently in the window carry that value) rather than a plain set, so
/// that `retract_batch` only drops a value once *all* of its occurrences have
/// left the window frame. The percentile is then computed over the set of keys
/// with a positive count.
#[derive(Debug)]
struct DistinctPercentileContAccumulator<T: ArrowNumericType> {
distinct_values: GenericDistinctBuffer<T>,
/// Distinct value -> number of in-window rows carrying it.
counts: HashMap<Hashable<T::Native>, usize>,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GenericDistinctBuffer previously used foldhash, but this will switch to using siphash, which is a lot slower.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — switched the count map to the foldhash RandomState that GenericDistinctBuffer uses, instead of the std default SipHash. Fixed in #23946.

percentile: f64,
}

impl<T: ArrowNumericType + Debug> DistinctPercentileContAccumulator<T> {
fn new(percentile: f64) -> Self {
Self {
distinct_values: GenericDistinctBuffer::new(T::DATA_TYPE),
counts: HashMap::default(),
percentile,
}
}
Expand All @@ -684,26 +692,50 @@ where
f64: AsPrimitive<T::Native>,
{
fn state(&mut self) -> Result<Vec<ScalarValue>> {
self.distinct_values.state()
// Emit the distinct keys as a single List scalar, matching the state
// shape declared in `state_fields` (a List of the input type). Counts
// are window-local bookkeeping and are intentionally not serialized:
// cross-partition merges only need the distinct key set.
let arr = Arc::new(
PrimitiveArray::<T>::from_iter_values(self.counts.keys().map(|v| v.0))
.with_data_type(T::DATA_TYPE),
);
Ok(vec![
SingleRowListArrayBuilder::new(arr).build_list_scalar(),
])
}

fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
self.distinct_values.update_batch(values)
// `values` may carry extra argument columns (e.g. the percentile
// literal); only the first column holds the aggregated values.
let arr = values[0].as_primitive::<T>();
for value in arr.iter().flatten() {
*self.counts.entry(Hashable(value)).or_default() += 1;
}
Comment on lines +711 to +714

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GenericDistinctBuffer has a fast-path for null-free arrays; might be worth adopting here as well.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adopted the same null-free fast path in update_batch and retract_batch. Fixed in #23946.

Ok(())
}

fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
self.distinct_values.merge_batch(states)
let list = states[0].as_list::<i32>();
for values in list.iter().flatten() {
let arr = values.as_primitive::<T>();
for value in arr.iter().flatten() {
*self.counts.entry(Hashable(value)).or_default() += 1;
}
}
Ok(())
}

fn evaluate(&mut self) -> Result<ScalarValue> {
let mut values: Vec<T::Native> =
self.distinct_values.values.iter().map(|v| v.0).collect();
let mut values: Vec<T::Native> = self.counts.keys().map(|v| v.0).collect();
let value = calculate_percentile::<T>(&mut values, self.percentile);
ScalarValue::new_primitive::<T>(value, &T::DATA_TYPE)
}

fn size(&self) -> usize {
size_of_val(self) + self.distinct_values.size()
size_of_val(self)
+ self.counts.capacity()
* (size_of::<Hashable<T::Native>>() + size_of::<usize>())
Comment on lines -706 to +738

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

estimate_memory_size::<(Hashable<T::Native>, usize)>(self.counts.capacity(), size_of_val(self)) is better I think.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switched to estimate_memory_size::<(Hashable<T::Native>, usize)>(self.counts.capacity(), size_of_val(self)). Fixed in #23946.

}

fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
Expand All @@ -713,7 +745,12 @@ where

let arr = values[0].as_primitive::<T>();
for value in arr.iter().flatten() {
self.distinct_values.values.remove(&Hashable(value));
if let Some(count) = self.counts.get_mut(&Hashable(value)) {
*count -= 1;
if *count == 0 {
self.counts.remove(&Hashable(value));
}
}
Comment on lines +748 to +753

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wdyt of raising an error if the caller attempts to retract a value that isn't in the map?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I left retract_batch silently skipping a value that isn't in the map rather than erroring. The sliding-window executor only retracts values it previously passed to update_batch, so a missing key shouldn't happen in practice — but if it ever did (a bug elsewhere), erroring would turn a latent inconsistency into a hard query failure, whereas skipping keeps the window's remaining counts intact. I'm not strongly attached either way; if you'd prefer a debug_assert! (catches it in tests/debug builds without failing production queries), I'm happy to add that instead. What do you think?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we try to retract a missing value, seems that we are likely to eventually produce incorrect query results, so I'd argue it is better to raise an error here rather than to silently proceed and return the wrong answer. I'd personally argue for internal_err!; but if you disagree, then I'd say debug_assert! is better than nothing 😊

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed — returning a wrong answer is worse than failing loudly here. Went with internal_err! in #23946: retracting a value not in the count map now surfaces as an error rather than silently proceeding. Thanks!

}
Ok(())
}
Expand Down
27 changes: 27 additions & 0 deletions datafusion/sqllogictest/test_files/aggregate.slt
Original file line number Diff line number Diff line change
Expand Up @@ -1104,6 +1104,33 @@ ORDER BY tags, timestamp;
statement ok
DROP TABLE median_window_test;

# Regression: percentile_cont(DISTINCT ...) used to forward the extra
# percentile-argument column into the distinct-values buffer (which asserts a
# single input array), panicking on every distinct query. Plain aggregate:
statement ok
CREATE TABLE distinct_pct(id INT, x DOUBLE) AS VALUES
(1, 5), (2, 5), (3, 9);

query R
SELECT percentile_cont(DISTINCT x, 0.5) FROM distinct_pct;
----
7

# Regression: distinct sliding-window percentile must count value multiplicity
# on retract. Row 3's frame is {5, 9}; the row-1 `5` leaves the frame but the
# row-2 `5` remains, so the distinct set is still {5, 9} (median 7), not {9}.
query IR
SELECT id, percentile_cont(DISTINCT x, 0.5)
OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)
FROM distinct_pct;
Comment on lines +1123 to +1125

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs an ORDER BY to be deterministic, or else use rowsort SLT directive.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added ORDER BY id to make the row order deterministic. Fixed in #23946.

----
1 5
2 5
3 7

statement ok
DROP TABLE distinct_pct;

query RT
select approx_median(arrow_cast(col_f32, 'Float16')), arrow_typeof(approx_median(arrow_cast(col_f32, 'Float16'))) from median_table;
----
Expand Down