Jefffrey commented on code in PR #9409:
URL: https://github.com/apache/arrow-rs/pull/9409#discussion_r2889540484
##########
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:
Personally I'd put an `unreachable!()` here to make it clear its an invalid
path; `as_primitive_array()` below already panics on invalid input
##########
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),
Review Comment:
ditto
##########
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);
+ Ok(acc.add_wrapping(val.mul_wrapping(V::Native::usize_as(len))))
+ });
+ sum
+ }
+
+ /// Computes the sum (erroring on overflow) of the array values.
+ pub(super) fn sum_checked<I: RunEndIndexType, V: ArrowNumericType>(
+ array: &dyn Array,
+ ) -> Result<Option<V::Native>, ArrowError> {
+ let Some(ree) = downcast::<I, V>(array) else {
+ return Err(ArrowError::InvalidArgumentError(
+ "Input array is not a TypedRunArray<'_, _,
PrimitiveArray<T>".to_string(),
Review Comment:
I think we can make the error easier to read by removing mention of
`TypedRunArray`; just mention it's not a RunArray of a primitive for example
##########
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:
I guess this was a debugging println?
--
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]
