Dandandan commented on code in PR #7376:
URL: https://github.com/apache/arrow-datafusion/pull/7376#discussion_r1302163472
##########
datafusion/physical-expr/src/aggregate/median.rs:
##########
@@ -106,159 +126,75 @@ impl PartialEq<dyn Any> for Median {
}
}
-#[derive(Debug)]
/// The median accumulator accumulates the raw input values
/// as `ScalarValue`s
///
/// The intermediate state is represented as a List of scalar values updated by
/// `merge_batch` and a `Vec` of `ArrayRef` that are converted to scalar values
/// in the final evaluation step so that we avoid expensive conversions and
/// allocations during `update_batch`.
-struct MedianAccumulator {
+struct MedianAccumulator<T: ArrowNumericType> {
data_type: DataType,
- arrays: Vec<ArrayRef>,
- all_values: Vec<ScalarValue>,
+ all_values: Vec<T::Native>,
}
-impl Accumulator for MedianAccumulator {
+impl<T: ArrowNumericType> std::fmt::Debug for MedianAccumulator<T> {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "MedianAccumulator({})", self.data_type)
+ }
+}
+
+impl<T: ArrowNumericType> Accumulator for MedianAccumulator<T> {
fn state(&self) -> Result<Vec<ScalarValue>> {
- let all_values = to_scalar_values(&self.arrays)?;
+ let all_values = self
+ .all_values
+ .iter()
+ .map(|x| ScalarValue::new_primitive::<T>(Some(*x),
&self.data_type))
+ .collect();
let state = ScalarValue::new_list(Some(all_values),
self.data_type.clone());
Ok(vec![state])
}
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- assert_eq!(values.len(), 1);
- let array = &values[0];
-
- // Defer conversions to scalar values to final evaluation.
- assert_eq!(array.data_type(), &self.data_type);
- self.arrays.push(array.clone());
-
+ let values = values[0].as_primitive::<T>();
+ self.all_values.reserve(values.len() - values.null_count());
+ self.all_values.extend(values.iter().flatten());
Ok(())
}
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- assert_eq!(states.len(), 1);
-
- let array = &states[0];
- assert!(matches!(array.data_type(), DataType::List(_)));
- for index in 0..array.len() {
- match ScalarValue::try_from_array(array, index)? {
- ScalarValue::List(Some(mut values), _) => {
- self.all_values.append(&mut values);
- }
- ScalarValue::List(None, _) => {} // skip empty state
- v => {
- return internal_err!(
- "unexpected state in median. Expected DataType::List,
got {v:?}"
- )
- }
- }
+ let array = states[0].as_list::<i32>();
+ for v in array.iter().flatten() {
+ self.update_batch(&[v])?
}
Ok(())
}
fn evaluate(&self) -> Result<ScalarValue> {
- let batch_values = to_scalar_values(&self.arrays)?;
-
- if !self
- .all_values
- .iter()
- .chain(batch_values.iter())
- .any(|v| !v.is_null())
- {
- return ScalarValue::try_from(&self.data_type);
- }
-
- // Create an array of all the non null values and find the
- // sorted indexes
- let array = ScalarValue::iter_to_array(
- self.all_values
- .iter()
- .chain(batch_values.iter())
- // ignore null values
- .filter(|v| !v.is_null())
- .cloned(),
- )?;
-
- // find the mid point
- let len = array.len();
- let mid = len / 2;
-
- // only sort up to the top size/2 elements
- let limit = Some(mid + 1);
- let options = None;
- let indices = sort_to_indices(&array, options, limit)?;
-
- // pick the relevant indices in the original arrays
- let result = if len >= 2 && len % 2 == 0 {
- // even number of values, average the two mid points
- let s1 = scalar_at_index(&array, &indices, mid - 1)?;
- let s2 = scalar_at_index(&array, &indices, mid)?;
- match s1.add(s2)? {
- ScalarValue::Int8(Some(v)) => ScalarValue::Int8(Some(v / 2)),
- ScalarValue::Int16(Some(v)) => ScalarValue::Int16(Some(v / 2)),
- ScalarValue::Int32(Some(v)) => ScalarValue::Int32(Some(v / 2)),
- ScalarValue::Int64(Some(v)) => ScalarValue::Int64(Some(v / 2)),
- ScalarValue::UInt8(Some(v)) => ScalarValue::UInt8(Some(v / 2)),
- ScalarValue::UInt16(Some(v)) => ScalarValue::UInt16(Some(v /
2)),
- ScalarValue::UInt32(Some(v)) => ScalarValue::UInt32(Some(v /
2)),
- ScalarValue::UInt64(Some(v)) => ScalarValue::UInt64(Some(v /
2)),
- ScalarValue::Float32(Some(v)) => ScalarValue::Float32(Some(v /
2.0)),
- ScalarValue::Float64(Some(v)) => ScalarValue::Float64(Some(v /
2.0)),
- ScalarValue::Decimal128(Some(v), p, s) => {
- ScalarValue::Decimal128(Some(v / 2), p, s)
- }
- v => {
- return internal_err!("Unsupported type in
MedianAccumulator: {v:?}")
- }
- }
+ // TODO: evaluate could pass &mut self
+ let mut d = self.all_values.clone();
+ let cmp = |x: &T::Native, y: &T::Native| x.compare(*y);
+
+ let len = d.len();
+ let median = if len == 0 {
+ None
+ } else if len % 2 == 0 {
+ let (low, high, _) = d.select_nth_unstable_by(len / 2, cmp);
+ let (_, low, _) = low.select_nth_unstable_by(low.len() - 1, cmp);
Review Comment:
🎉
##########
datafusion/physical-expr/src/aggregate/median.rs:
##########
@@ -106,159 +126,75 @@ impl PartialEq<dyn Any> for Median {
}
}
-#[derive(Debug)]
/// The median accumulator accumulates the raw input values
/// as `ScalarValue`s
///
/// The intermediate state is represented as a List of scalar values updated by
/// `merge_batch` and a `Vec` of `ArrayRef` that are converted to scalar values
/// in the final evaluation step so that we avoid expensive conversions and
/// allocations during `update_batch`.
-struct MedianAccumulator {
+struct MedianAccumulator<T: ArrowNumericType> {
data_type: DataType,
- arrays: Vec<ArrayRef>,
- all_values: Vec<ScalarValue>,
+ all_values: Vec<T::Native>,
}
-impl Accumulator for MedianAccumulator {
+impl<T: ArrowNumericType> std::fmt::Debug for MedianAccumulator<T> {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "MedianAccumulator({})", self.data_type)
+ }
+}
+
+impl<T: ArrowNumericType> Accumulator for MedianAccumulator<T> {
fn state(&self) -> Result<Vec<ScalarValue>> {
- let all_values = to_scalar_values(&self.arrays)?;
+ let all_values = self
+ .all_values
+ .iter()
+ .map(|x| ScalarValue::new_primitive::<T>(Some(*x),
&self.data_type))
+ .collect();
let state = ScalarValue::new_list(Some(all_values),
self.data_type.clone());
Ok(vec![state])
}
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- assert_eq!(values.len(), 1);
- let array = &values[0];
-
- // Defer conversions to scalar values to final evaluation.
- assert_eq!(array.data_type(), &self.data_type);
- self.arrays.push(array.clone());
-
+ let values = values[0].as_primitive::<T>();
+ self.all_values.reserve(values.len() - values.null_count());
+ self.all_values.extend(values.iter().flatten());
Ok(())
}
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- assert_eq!(states.len(), 1);
-
- let array = &states[0];
- assert!(matches!(array.data_type(), DataType::List(_)));
- for index in 0..array.len() {
- match ScalarValue::try_from_array(array, index)? {
- ScalarValue::List(Some(mut values), _) => {
- self.all_values.append(&mut values);
- }
- ScalarValue::List(None, _) => {} // skip empty state
- v => {
- return internal_err!(
- "unexpected state in median. Expected DataType::List,
got {v:?}"
- )
- }
- }
+ let array = states[0].as_list::<i32>();
+ for v in array.iter().flatten() {
+ self.update_batch(&[v])?
}
Ok(())
}
fn evaluate(&self) -> Result<ScalarValue> {
- let batch_values = to_scalar_values(&self.arrays)?;
-
- if !self
- .all_values
- .iter()
- .chain(batch_values.iter())
- .any(|v| !v.is_null())
- {
- return ScalarValue::try_from(&self.data_type);
- }
-
- // Create an array of all the non null values and find the
- // sorted indexes
- let array = ScalarValue::iter_to_array(
- self.all_values
- .iter()
- .chain(batch_values.iter())
- // ignore null values
- .filter(|v| !v.is_null())
- .cloned(),
- )?;
-
- // find the mid point
- let len = array.len();
- let mid = len / 2;
-
- // only sort up to the top size/2 elements
- let limit = Some(mid + 1);
- let options = None;
- let indices = sort_to_indices(&array, options, limit)?;
-
- // pick the relevant indices in the original arrays
- let result = if len >= 2 && len % 2 == 0 {
- // even number of values, average the two mid points
- let s1 = scalar_at_index(&array, &indices, mid - 1)?;
- let s2 = scalar_at_index(&array, &indices, mid)?;
- match s1.add(s2)? {
- ScalarValue::Int8(Some(v)) => ScalarValue::Int8(Some(v / 2)),
- ScalarValue::Int16(Some(v)) => ScalarValue::Int16(Some(v / 2)),
- ScalarValue::Int32(Some(v)) => ScalarValue::Int32(Some(v / 2)),
- ScalarValue::Int64(Some(v)) => ScalarValue::Int64(Some(v / 2)),
- ScalarValue::UInt8(Some(v)) => ScalarValue::UInt8(Some(v / 2)),
- ScalarValue::UInt16(Some(v)) => ScalarValue::UInt16(Some(v /
2)),
- ScalarValue::UInt32(Some(v)) => ScalarValue::UInt32(Some(v /
2)),
- ScalarValue::UInt64(Some(v)) => ScalarValue::UInt64(Some(v /
2)),
- ScalarValue::Float32(Some(v)) => ScalarValue::Float32(Some(v /
2.0)),
- ScalarValue::Float64(Some(v)) => ScalarValue::Float64(Some(v /
2.0)),
- ScalarValue::Decimal128(Some(v), p, s) => {
- ScalarValue::Decimal128(Some(v / 2), p, s)
- }
- v => {
- return internal_err!("Unsupported type in
MedianAccumulator: {v:?}")
- }
- }
+ // TODO: evaluate could pass &mut self
+ let mut d = self.all_values.clone();
+ let cmp = |x: &T::Native, y: &T::Native| x.compare(*y);
+
+ let len = d.len();
+ let median = if len == 0 {
+ None
+ } else if len % 2 == 0 {
+ let (low, high, _) = d.select_nth_unstable_by(len / 2, cmp);
+ let (_, low, _) = low.select_nth_unstable_by(low.len() - 1, cmp);
Review Comment:
🎉
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
