Skip to content

Commit b4c02d0

Browse files
authored
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. -->
1 parent 710e68e commit b4c02d0

1 file changed

Lines changed: 60 additions & 2 deletions

File tree

arrow-string/src/length.rs

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,20 @@ use arrow_array::{cast::AsArray, types::*};
2222
use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer};
2323
use arrow_schema::{ArrowError, DataType};
2424
use std::sync::Arc;
25+
macro_rules! ree_length {
26+
($array:expr, $run_type:ty, $k:expr, $v:expr) => {{
27+
let ree = $array.as_run_opt::<$run_type>().unwrap();
28+
let inner_value_lengths = length(ree.values().as_ref())?;
29+
let out_ree = unsafe {
30+
RunArray::<$run_type>::new_unchecked(
31+
DataType::RunEndEncoded(Arc::clone($k), Arc::clone($v)),
32+
ree.run_ends().clone(),
33+
inner_value_lengths,
34+
)
35+
};
36+
Ok(Arc::new(out_ree) as ArrayRef)
37+
}};
38+
}
2539

2640
fn length_impl<P: ArrowPrimitiveType>(
2741
offsets: &OffsetBuffer<P::Native>,
@@ -50,14 +64,14 @@ fn bit_length_impl<P: ArrowPrimitiveType>(
5064
/// For string array and binary array, length is the number of bytes of each value.
5165
///
5266
/// * this only accepts ListArray/LargeListArray, StringArray/LargeStringArray/StringViewArray, BinaryArray/LargeBinaryArray, FixedSizeListArray,
53-
/// and ListViewArray/LargeListViewArray, or DictionaryArray with above Arrays as values
67+
/// and ListViewArray/LargeListViewArray, or DictionaryArray with above Arrays as values, or
68+
/// RunEndEncoded arrays with above arrays as values
5469
/// * length of null is null.
5570
pub fn length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
5671
if let Some(d) = array.as_any_dictionary_opt() {
5772
let lengths = length(d.values().as_ref())?;
5873
return Ok(d.with_values(lengths));
5974
}
60-
6175
match array.data_type() {
6276
DataType::List(_) => {
6377
let list = array.as_list::<i32>();
@@ -116,6 +130,15 @@ pub fn length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
116130
list.nulls().cloned(),
117131
)?))
118132
}
133+
DataType::RunEndEncoded(k, v) => match k.data_type() {
134+
DataType::Int16 => ree_length!(array, Int16Type, &k, &v),
135+
DataType::Int32 => ree_length!(array, Int32Type, &k, &v),
136+
DataType::Int64 => ree_length!(array, Int64Type, &k, &v),
137+
_ => Err(ArrowError::InvalidArgumentError(format!(
138+
"Invalid run-end type: {:?}",
139+
k.data_type()
140+
))),
141+
},
119142
other => Err(ArrowError::ComputeError(format!(
120143
"length not supported for {other:?}"
121144
))),
@@ -845,4 +868,39 @@ mod tests {
845868
let result = bit_length(&array).unwrap();
846869
assert_eq!(result.as_ref(), &Int32Array::from(vec![32; 4]));
847870
}
871+
#[test]
872+
fn length_test_ree_string_values() {
873+
use arrow_array::RunArray;
874+
use arrow_array::types::Int32Type;
875+
876+
let string_values = StringArray::from(vec!["hello", "owl", "test", "arrow", "a"]);
877+
let run_ends = PrimitiveArray::<Int32Type>::from(vec![2i32, 5, 9, 11, 14]);
878+
let ree_array = RunArray::<Int32Type>::try_new(&run_ends, &string_values).unwrap();
879+
880+
let result = length(&ree_array).unwrap();
881+
let result = result
882+
.as_any()
883+
.downcast_ref::<RunArray<Int32Type>>()
884+
.unwrap();
885+
886+
let result_values = result
887+
.values()
888+
.as_any()
889+
.downcast_ref::<Int32Array>()
890+
.unwrap();
891+
892+
let expected: Int32Array = vec![5, 3, 4, 5, 1].into();
893+
assert_eq!(&expected, result_values);
894+
}
895+
#[test]
896+
fn length_test_ree_invalid_type_early_fail() {
897+
use arrow_array::RunArray;
898+
use arrow_array::types::Int32Type;
899+
900+
let uint64_values = UInt64Array::from(vec![1u64, 2, 3]);
901+
let run_ends = PrimitiveArray::<Int32Type>::from(vec![1i32, 2, 3]);
902+
let ree_array = RunArray::<Int32Type>::try_new(&run_ends, &uint64_values).unwrap();
903+
904+
assert!(length(&ree_array).is_err());
905+
}
848906
}

0 commit comments

Comments
 (0)