This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new b4c02d0e24 support length() on Run-end encoding arrays (#9838)
b4c02d0e24 is described below
commit b4c02d0e248c4550af1fa3229c6d7d05662f5326
Author: RIchard Baah <[email protected]>
AuthorDate: Mon Apr 27 15:58:32 2026 -0400
support length() on Run-end encoding arrays (#9838)
# Which issue does this PR close?
works towards closing #3520, adds run-end encoding support for legnth
operations when the values array is valid.
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
# Rationale for this change
REE is unsupported to length
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
Introduces a macro to handle each REE run-end index type, eliminating
duplicated logic across match arms. The macro uses **new_unchecked**
instead of **try_new** to skip redundant validation, since the input is
already a valid REE array and length has its own error handling, the
invariants are guaranteed to hold before the unchecked call.
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
Yes.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
# Are there any user-facing changes?
Yes, this allows callers to invoke **length()** directly on REE-encoded
columns without materializing or casting to a primitive array first.
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
arrow-string/src/length.rs | 62 ++++++++++++++++++++++++++++++++++++++++++++--
1 file changed, 60 insertions(+), 2 deletions(-)
diff --git a/arrow-string/src/length.rs b/arrow-string/src/length.rs
index f7461d37c4..ff98c5632b 100644
--- a/arrow-string/src/length.rs
+++ b/arrow-string/src/length.rs
@@ -22,6 +22,20 @@ use arrow_array::{cast::AsArray, types::*};
use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer};
use arrow_schema::{ArrowError, DataType};
use std::sync::Arc;
+macro_rules! ree_length {
+ ($array:expr, $run_type:ty, $k:expr, $v:expr) => {{
+ let ree = $array.as_run_opt::<$run_type>().unwrap();
+ let inner_value_lengths = length(ree.values().as_ref())?;
+ let out_ree = unsafe {
+ RunArray::<$run_type>::new_unchecked(
+ DataType::RunEndEncoded(Arc::clone($k), Arc::clone($v)),
+ ree.run_ends().clone(),
+ inner_value_lengths,
+ )
+ };
+ Ok(Arc::new(out_ree) as ArrayRef)
+ }};
+}
fn length_impl<P: ArrowPrimitiveType>(
offsets: &OffsetBuffer<P::Native>,
@@ -50,14 +64,14 @@ fn bit_length_impl<P: ArrowPrimitiveType>(
/// For string array and binary array, length is the number of bytes of each
value.
///
/// * this only accepts ListArray/LargeListArray,
StringArray/LargeStringArray/StringViewArray, BinaryArray/LargeBinaryArray,
FixedSizeListArray,
-/// and ListViewArray/LargeListViewArray, or DictionaryArray with above
Arrays as values
+/// and ListViewArray/LargeListViewArray, or DictionaryArray with above
Arrays as values, or
+/// RunEndEncoded arrays with above arrays as values
/// * length of null is null.
pub fn length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
if let Some(d) = array.as_any_dictionary_opt() {
let lengths = length(d.values().as_ref())?;
return Ok(d.with_values(lengths));
}
-
match array.data_type() {
DataType::List(_) => {
let list = array.as_list::<i32>();
@@ -116,6 +130,15 @@ pub fn length(array: &dyn Array) -> Result<ArrayRef,
ArrowError> {
list.nulls().cloned(),
)?))
}
+ DataType::RunEndEncoded(k, v) => match k.data_type() {
+ DataType::Int16 => ree_length!(array, Int16Type, &k, &v),
+ DataType::Int32 => ree_length!(array, Int32Type, &k, &v),
+ DataType::Int64 => ree_length!(array, Int64Type, &k, &v),
+ _ => Err(ArrowError::InvalidArgumentError(format!(
+ "Invalid run-end type: {:?}",
+ k.data_type()
+ ))),
+ },
other => Err(ArrowError::ComputeError(format!(
"length not supported for {other:?}"
))),
@@ -845,4 +868,39 @@ mod tests {
let result = bit_length(&array).unwrap();
assert_eq!(result.as_ref(), &Int32Array::from(vec![32; 4]));
}
+ #[test]
+ fn length_test_ree_string_values() {
+ use arrow_array::RunArray;
+ use arrow_array::types::Int32Type;
+
+ let string_values = StringArray::from(vec!["hello", "owl", "test",
"arrow", "a"]);
+ let run_ends = PrimitiveArray::<Int32Type>::from(vec![2i32, 5, 9, 11,
14]);
+ let ree_array = RunArray::<Int32Type>::try_new(&run_ends,
&string_values).unwrap();
+
+ let result = length(&ree_array).unwrap();
+ let result = result
+ .as_any()
+ .downcast_ref::<RunArray<Int32Type>>()
+ .unwrap();
+
+ let result_values = result
+ .values()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+
+ let expected: Int32Array = vec![5, 3, 4, 5, 1].into();
+ assert_eq!(&expected, result_values);
+ }
+ #[test]
+ fn length_test_ree_invalid_type_early_fail() {
+ use arrow_array::RunArray;
+ use arrow_array::types::Int32Type;
+
+ let uint64_values = UInt64Array::from(vec![1u64, 2, 3]);
+ let run_ends = PrimitiveArray::<Int32Type>::from(vec![1i32, 2, 3]);
+ let ree_array = RunArray::<Int32Type>::try_new(&run_ends,
&uint64_values).unwrap();
+
+ assert!(length(&ree_array).is_err());
+ }
}
