This is an automated email from the ASF dual-hosted git repository.
kamille pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1e0531f93d perf: Improve `median` with no grouping by 2X (#14399)
1e0531f93d is described below
commit 1e0531f93d4c0ecfa5ebdaa76d61a44ded8dfb42
Author: Yongting You <[email protected]>
AuthorDate: Sun Feb 2 21:50:58 2025 +0800
perf: Improve `median` with no grouping by 2X (#14399)
* improve median() no grouping case
* review
---
datafusion/functions-aggregate/src/median.rs | 28 +++++++++++++++++++++-------
1 file changed, 21 insertions(+), 7 deletions(-)
diff --git a/datafusion/functions-aggregate/src/median.rs
b/datafusion/functions-aggregate/src/median.rs
index defbbe737a..ba6b63260e 100644
--- a/datafusion/functions-aggregate/src/median.rs
+++ b/datafusion/functions-aggregate/src/median.rs
@@ -242,14 +242,28 @@ impl<T: ArrowNumericType> Debug for MedianAccumulator<T> {
impl<T: ArrowNumericType> Accumulator for MedianAccumulator<T> {
fn state(&mut self) -> Result<Vec<ScalarValue>> {
- let all_values = self
- .all_values
- .iter()
- .map(|x| ScalarValue::new_primitive::<T>(Some(*x),
&self.data_type))
- .collect::<Result<Vec<_>>>()?;
+ // Convert `all_values` to `ListArray` and return a single List
ScalarValue
- let arr = ScalarValue::new_list_nullable(&all_values, &self.data_type);
- Ok(vec![ScalarValue::List(arr)])
+ // Build offsets
+ let offsets =
+ OffsetBuffer::new(ScalarBuffer::from(vec![0, self.all_values.len()
as i32]));
+
+ // Build inner array
+ let values_array = PrimitiveArray::<T>::new(
+ ScalarBuffer::from(std::mem::take(&mut self.all_values)),
+ None,
+ )
+ .with_data_type(self.data_type.clone());
+
+ // Build the result list array
+ let list_array = ListArray::new(
+ Arc::new(Field::new_list_field(self.data_type.clone(), true)),
+ offsets,
+ Arc::new(values_array),
+ None,
+ );
+
+ Ok(vec![ScalarValue::List(Arc::new(list_array))])
}
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
