brunal commented on code in PR #9409:
URL: https://github.com/apache/arrow-rs/pull/9409#discussion_r2899648753
##########
arrow-arith/src/aggregate.rs:
##########
@@ -567,14 +567,22 @@ where
Some(sum)
}
+ DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() {
+ DataType::Int16 => ree::sum_wrapping::<types::Int16Type,
T>(&array),
+ DataType::Int32 => ree::sum_wrapping::<types::Int32Type,
T>(&array),
+ DataType::Int64 => ree::sum_wrapping::<types::Int64Type,
T>(&array),
+ _ => None,
Review Comment:
Done
##########
arrow-arith/src/aggregate.rs:
##########
@@ -603,10 +611,111 @@ where
Ok(Some(sum))
}
+ DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() {
+ DataType::Int16 => ree::sum_checked::<types::Int16Type, T>(&array),
+ DataType::Int32 => ree::sum_checked::<types::Int32Type, T>(&array),
+ DataType::Int64 => ree::sum_checked::<types::Int64Type, T>(&array),
+ _ => Ok(None),
+ },
_ => sum_checked::<T>(as_primitive_array(&array)),
}
}
+// Logic for summing run-end-encoded arrays.
+mod ree {
+ use std::convert::Infallible;
+
+ use arrow_array::cast::AsArray;
+ use arrow_array::types::RunEndIndexType;
+ use arrow_array::{Array, ArrowNativeTypeOp, ArrowNumericType,
PrimitiveArray, TypedRunArray};
+ use arrow_buffer::ArrowNativeType;
+ use arrow_schema::ArrowError;
+
+ /// Downcasts an array to a TypedRunArray.
+ fn downcast<'a, I: RunEndIndexType, V: ArrowNumericType>(
+ array: &'a dyn Array,
+ ) -> Option<TypedRunArray<'a, I, PrimitiveArray<V>>> {
+ let array = array.as_run_opt::<I>()?;
+ // We only support RunArray wrapping primitive types.
+ array.downcast::<PrimitiveArray<V>>()
+ }
+
+ /// Computes the sum (wrapping) of the array values.
+ pub(super) fn sum_wrapping<I: RunEndIndexType, V: ArrowNumericType>(
+ array: &dyn Array,
+ ) -> Option<V::Native> {
+ let ree = downcast::<I, V>(array)?;
+ let Ok(sum) = fold(ree, |acc, val, len| -> Result<V::Native,
Infallible> {
+ println!("Adding {:?}x{} to {:?}", val, len, acc);
Review Comment:
Done
--
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]
