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 af1c24a48d minor: drive-by refactors for dicts in substring & filter
(#10264)
af1c24a48d is described below
commit af1c24a48d137a62fccde3de63d9b0b225a985f3
Author: Jeffrey Vo <[email protected]>
AuthorDate: Thu Jul 2 20:07:42 2026 +0900
minor: drive-by refactors for dicts in substring & filter (#10264)
# Which issue does this PR close?
<!--
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.
-->
- Part of https://github.com/apache/arrow-rs/issues/9298
# Rationale for this change
<!--
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.
-->
Some drive-by simplifications whilst I was checking some parts of the
codebase.
# What changes are included in this PR?
<!--
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.
-->
In substring utilize `as_any_dictionary()` to cut down on lines of code
(and also codegen)
In filter, follow similar approach to take when reconstructing a
dictionary
# Are these changes tested?
<!--
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)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
Existing tests
# Are there any user-facing changes?
<!--
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.
-->
No
---
arrow-select/src/filter.rs | 18 ++++------
arrow-string/src/substring.rs | 80 ++++++++-----------------------------------
2 files changed, 20 insertions(+), 78 deletions(-)
diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs
index 1de22e79c0..8b69811d4a 100644
--- a/arrow-select/src/filter.rs
+++ b/arrow-select/src/filter.rs
@@ -1007,20 +1007,14 @@ fn filter_fixed_size_binary(
}
/// `filter` implementation for dictionaries
-fn filter_dict<T>(array: &DictionaryArray<T>, predicate: &FilterPredicate) ->
DictionaryArray<T>
-where
- T: ArrowDictionaryKeyType,
- T::Native: num_traits::Num,
-{
- let builder = filter_primitive::<T>(array.keys(), predicate)
- .into_data()
- .into_builder()
- .data_type(array.data_type().clone())
- .child_data(vec![array.values().to_data()]);
-
+fn filter_dict<K: ArrowDictionaryKeyType>(
+ array: &DictionaryArray<K>,
+ predicate: &FilterPredicate,
+) -> DictionaryArray<K> {
// SAFETY:
// Keys were valid before, filtered subset is therefore still valid
- DictionaryArray::from(unsafe { builder.build_unchecked() })
+ let new_keys = filter_primitive(array.keys(), predicate);
+ unsafe { DictionaryArray::new_unchecked(new_keys, array.values().clone()) }
}
/// `filter` implementation for structs
diff --git a/arrow-string/src/substring.rs b/arrow-string/src/substring.rs
index 1146956345..152f1f152d 100644
--- a/arrow-string/src/substring.rs
+++ b/arrow-string/src/substring.rs
@@ -20,6 +20,7 @@
//! [GenericStringArray], [GenericBinaryArray], [FixedSizeBinaryArray],
[DictionaryArray]
use arrow_array::builder::BufferBuilder;
+use arrow_array::cast::AsArray;
use arrow_array::types::*;
use arrow_array::*;
use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBuffer};
@@ -73,54 +74,17 @@ pub fn substring(
start: i64,
length: Option<u64>,
) -> Result<ArrayRef, ArrowError> {
- macro_rules! substring_dict {
- ($kt: ident, $($t: ident: $gt: ident), *) => {
- match $kt.as_ref() {
- $(
- &DataType::$t => {
- let dict = array
- .as_any()
- .downcast_ref::<DictionaryArray<$gt>>()
- .unwrap_or_else(|| {
- panic!("Expect 'DictionaryArray<{}>' but got
array of data type {:?}",
- stringify!($gt), array.data_type())
- });
- let values = substring(dict.values(), start, length)?;
- Ok(Arc::new(dict.with_values(values)))
- },
- )*
- t => panic!("Unsupported dictionary key type: {}", t)
- }
- }
- }
-
match array.data_type() {
- DataType::Dictionary(kt, _) => {
- substring_dict!(
- kt,
- Int8: Int8Type,
- Int16: Int16Type,
- Int32: Int32Type,
- Int64: Int64Type,
- UInt8: UInt8Type,
- UInt16: UInt16Type,
- UInt32: UInt32Type,
- UInt64: UInt64Type
- )
+ DataType::Dictionary(_, _) => {
+ let dictionary = array.as_any_dictionary();
+ let values = substring(dictionary.values(), start, length)?;
+ Ok(Arc::new(dictionary.with_values(values)))
+ }
+ DataType::LargeBinary => {
+ byte_substring(array.as_binary::<i64>(), start, length.map(|e| e
as i64))
}
- DataType::LargeBinary => byte_substring(
- array
- .as_any()
- .downcast_ref::<LargeBinaryArray>()
- .expect("A large binary is expected"),
- start,
- length.map(|e| e as i64),
- ),
DataType::Binary => byte_substring(
- array
- .as_any()
- .downcast_ref::<BinaryArray>()
- .expect("A binary is expected"),
+ array.as_binary::<i32>(),
start as i32,
length.map(|e| e as i32),
),
@@ -128,29 +92,13 @@ pub fn substring(
let old_len: usize = (*old_len)
.try_into()
.expect("negative FixedSizeBinary value length");
- fixed_size_binary_substring(
- array
- .as_any()
- .downcast_ref::<FixedSizeBinaryArray>()
- .expect("a fixed size binary is expected"),
- old_len,
- start,
- length,
- )
+ fixed_size_binary_substring(array.as_fixed_size_binary(), old_len,
start, length)
+ }
+ DataType::LargeUtf8 => {
+ byte_substring(array.as_string::<i64>(), start, length.map(|e| e
as i64))
}
- DataType::LargeUtf8 => byte_substring(
- array
- .as_any()
- .downcast_ref::<LargeStringArray>()
- .expect("A large string is expected"),
- start,
- length.map(|e| e as i64),
- ),
DataType::Utf8 => byte_substring(
- array
- .as_any()
- .downcast_ref::<StringArray>()
- .expect("A string is expected"),
+ array.as_string::<i32>(),
start as i32,
length.map(|e| e as i32),
),
