-
Notifications
You must be signed in to change notification settings - Fork 2.3k
fix: correct percentile_cont(DISTINCT) accumulation and sliding-window retract #23913
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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::{ | ||
|
|
@@ -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; | ||
|
|
@@ -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>, | ||
| 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, | ||
| } | ||
| } | ||
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Adopted the same null-free fast path in |
||
| 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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Switched to |
||
| } | ||
|
|
||
| fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { | ||
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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?
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I left
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
| } | ||
| Ok(()) | ||
| } | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This needs an
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added |
||
| ---- | ||
| 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; | ||
| ---- | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
GenericDistinctBufferpreviously used foldhash, but this will switch to using siphash, which is a lot slower.There was a problem hiding this comment.
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
RandomStatethatGenericDistinctBufferuses, instead of the std default SipHash. Fixed in #23946.