Jefffrey commented on code in PR #9959:
URL: https://github.com/apache/arrow-rs/pull/9959#discussion_r3223446844
##########
arrow-arith/src/temporal.rs:
##########
@@ -249,15 +248,12 @@ pub fn date_part(array: &dyn Array, part: DatePart) ->
Result<ArrayRef, ArrowErr
let new_array = array.with_values(values);
Ok(new_array)
}
- DataType::RunEndEncoded(k, _) => match k.data_type() {
- DataType::Int16 => ree_map!(array, Int16Type, |a| date_part(a,
part)),
- DataType::Int32 => ree_map!(array, Int32Type, |a| date_part(a,
part)),
- DataType::Int64 => ree_map!(array, Int64Type, |a| date_part(a,
part)),
- _ => Err(ArrowError::InvalidArgumentError(format!(
- "Invalid run-end type: {:?}",
- k.data_type()
- ))),
- },
+ DataType::RunEndEncoded(_, _) => {
+ let array = array.as_any_ree();
+ let values = date_part(array.values().as_ref(), part)?;
Review Comment:
```suggestion
let values = date_part(array.values(), part)?;
```
##########
arrow-string/src/length.rs:
##########
@@ -144,6 +138,10 @@ pub fn bit_length(array: &dyn Array) -> Result<ArrayRef,
ArrowError> {
let lengths = bit_length(d.values().as_ref())?;
return Ok(d.with_values(lengths));
}
+ if let Some(ree) = array.as_any_ree_opt() {
+ let lengths = bit_length(ree.values().as_ref())?;
Review Comment:
```suggestion
let lengths = bit_length(ree.values())?;
```
##########
arrow-array/src/array/run_array.rs:
##########
@@ -781,6 +781,52 @@ where
RunArrayIter::new(self)
}
}
+/// An array that can be downcast to a [`RunArray`] of any run end type and
any value type.
+///
+/// This can be used to efficiently implement kernels for all possible run end
+/// types without needing to create specialized implementations for each key
type.
+pub trait AnyRunEndArray: Array {
+ /// Returns the run ends of this array as a primitive array.
+ fn run_ends(&self) -> ArrayRef;
+
+ /// Returns the values of this array.
+ fn values(&self) -> &Arc<dyn Array>;
+
+ /// Returns a new run-end encoded array with the given values, preserving
the
+ /// existing run ends.
+ fn with_values(&self, values: ArrayRef) -> ArrayRef;
+}
+
+impl<R: RunEndIndexType> AnyRunEndArray for RunArray<R> {
+ fn run_ends(&self) -> ArrayRef {
+ let data = unsafe {
+ ArrayDataBuilder::new(R::DATA_TYPE)
+ .len(self.run_ends.values().len())
+ .buffers(vec![self.run_ends.inner().inner().clone()])
+ .build_unchecked()
+ };
+ Arc::new(PrimitiveArray::<R>::from(data))
+ }
+
+ fn values(&self) -> &Arc<dyn Array> {
+ &self.values
+ }
+
+ fn with_values(&self, values: ArrayRef) -> ArrayRef {
+ debug_assert_eq!(values.len(), self.values.len());
+ let (run_ends_field, values_field) = match &self.data_type {
+ DataType::RunEndEncoded(r, v) => (r, v),
+ _ => unreachable!("RunArray should have type RunEndEncoded"),
+ };
+ let dt = DataType::RunEndEncoded(Arc::clone(run_ends_field),
Arc::clone(values_field));
Review Comment:
The assert should be regular and not a debug one, and for the values field
we'd need to recreate it according to the datatype on the incoming `values`
##########
arrow-string/src/length.rs:
##########
@@ -59,6 +58,10 @@ pub fn length(array: &dyn Array) -> Result<ArrayRef,
ArrowError> {
let lengths = length(d.values().as_ref())?;
return Ok(d.with_values(lengths));
}
+ if let Some(ree) = array.as_any_ree_opt() {
+ let lengths = length(ree.values().as_ref())?;
Review Comment:
```suggestion
let lengths = length(ree.values())?;
```
--
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]
