This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch 57_maintenance
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/57_maintenance by this push:
new fadd9eea5e [57_maintenance] fix(arrow-array): preserve type params in
`PrimitiveArray::into_builder` (#11807)
fadd9eea5e is described below
commit fadd9eea5e533623238d93c73b8ca48feca7511a
Author: Jeffrey Vo <[email protected]>
AuthorDate: Fri Oct 23 00:20:07 2026 +0900
[57_maintenance] fix(arrow-array): preserve type params in
`PrimitiveArray::into_builder` (#11807)
- Backport #11310 and #11321
- For inclusion in https://github.com/apache/arrow-rs/issues/11666
Cherry-pick of
https://github.com/apache/arrow-rs/commit/2f3cd82bea9f5baabdcc5b76efb4624620eeeb15
and
https://github.com/apache/arrow-rs/commit/38054b7713eade3bd8092bbf0fc3eaad69d4ce65
These are combined into a single PR since the fixes work on the same
function
## NOTE
We also backport #9473 183f8c1c5361ac5f026d6fbfa8e99a2920dcb652 since we
needed it for a test, and I think fix makes more sense with this in
mind. It was introduced in 58 so only 57 needs this extra backport
---------
Co-authored-by: Neil Conway <[email protected]>
Co-authored-by: Bruno <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>
---
arrow-array/src/array/dictionary_array.rs | 22 +++++++-
arrow-array/src/array/primitive_array.rs | 24 ++++++++-
arrow-array/src/builder/primitive_run_builder.rs | 65 +++++++++++++++++++++++-
3 files changed, 106 insertions(+), 5 deletions(-)
diff --git a/arrow-array/src/array/dictionary_array.rs
b/arrow-array/src/array/dictionary_array.rs
index a68ddb22a5..fc712f7c6e 100644
--- a/arrow-array/src/array/dictionary_array.rs
+++ b/arrow-array/src/array/dictionary_array.rs
@@ -1052,7 +1052,7 @@ impl<K: ArrowDictionaryKeyType> AnyDictionaryArray for
DictionaryArray<K> {
mod tests {
use super::*;
use crate::cast::as_dictionary_array;
- use crate::{Int8Array, Int16Array, Int32Array, RunArray};
+ use crate::{Decimal128Array, Int8Array, Int16Array, Int32Array, RunArray};
use arrow_buffer::{Buffer, ToByteSlice};
#[test]
@@ -1439,6 +1439,26 @@ mod tests {
assert_eq!(expected, returned);
}
+ #[test]
+ fn test_unary_mut_keeps_value_type() {
+ let values = || {
+ Decimal128Array::from(vec![123, 456])
+ .with_precision_and_scale(5, 2)
+ .unwrap()
+ };
+ let keys = Int8Array::from_iter_values([0, 1, 0]);
+ let dict = DictionaryArray::new(keys, Arc::new(values()));
+ let data_type = dict.data_type().clone();
+ let updated = dict.unary_mut::<_, Decimal128Type>(|x| x + 1).unwrap();
+ assert_eq!(updated.data_type(), &data_type);
+
+ // When the keys are shared, the original array is returned
+ let keys = Int8Array::from_iter_values([0, 1, 0]);
+ let dict = DictionaryArray::new(keys.clone(), Arc::new(values()));
+ let returned = dict.unary_mut::<_, Decimal128Type>(|x| x +
1).unwrap_err();
+ assert_eq!(returned, DictionaryArray::new(keys, Arc::new(values())));
+ }
+
#[test]
fn test_occupancy() {
let keys = Int32Array::new((100..200).collect(), None);
diff --git a/arrow-array/src/array/primitive_array.rs
b/arrow-array/src/array/primitive_array.rs
index 0ef3520282..3b8974e4a5 100644
--- a/arrow-array/src/array/primitive_array.rs
+++ b/arrow-array/src/array/primitive_array.rs
@@ -1124,6 +1124,7 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
/// `Err(self)`
pub fn into_builder(self) -> Result<PrimitiveBuilder<T>, Self> {
let len = self.len();
+ let data_type = self.data_type.clone();
let data = self.into_data();
let null_bit_buffer = data.nulls().map(|b| b.inner().sliced());
@@ -1163,9 +1164,9 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
};
match try_mutable_buffers {
- Ok(builder) => Ok(builder),
+ Ok(builder) => Ok(builder.with_data_type(data_type)),
Err((buffer, null_bit_buffer)) => {
- let builder = ArrayData::builder(T::DATA_TYPE)
+ let builder = ArrayData::builder(data_type)
.len(len)
.add_buffer(buffer)
.null_bit_buffer(null_bit_buffer);
@@ -2807,6 +2808,25 @@ mod tests {
}
}
+ #[test]
+ fn test_into_builder_keeps_data_type() {
+ let array = Decimal128Array::from(vec![1, 2])
+ .with_precision_and_scale(10, 2)
+ .unwrap();
+ let data_type = array.data_type().clone();
+
+ // Shared buffers: the array is returned
+ let shared = array.clone();
+ let array = array.into_builder().unwrap_err();
+ assert_eq!(array.data_type(), &data_type);
+ drop(shared);
+
+ // Unshared buffers: the buffers are reused by the builder
+ let mut builder = array.into_builder().unwrap();
+ builder.append_value(3);
+ assert_eq!(builder.finish().data_type(), &data_type);
+ }
+
#[test]
fn test_into_builder_on_sliced_array() {
let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
diff --git a/arrow-array/src/builder/primitive_run_builder.rs
b/arrow-array/src/builder/primitive_run_builder.rs
index cc97246f19..0914944030 100644
--- a/arrow-array/src/builder/primitive_run_builder.rs
+++ b/arrow-array/src/builder/primitive_run_builder.rs
@@ -108,6 +108,20 @@ where
prev_run_end_index: 0,
}
}
+
+ /// Overrides the data type of the values child array.
+ ///
+ /// By default, `V::DATA_TYPE` is used (via [`PrimitiveBuilder`]). This
+ /// allows setting the timezone of a Timestamp, the precision & scale of a
+ /// Decimal, etc.
+ ///
+ /// # Panics
+ ///
+ /// This method panics if `values_builder` rejects `data_type`.
+ pub fn with_data_type(mut self, data_type: arrow_schema::DataType) -> Self
{
+ self.values_builder = self.values_builder.with_data_type(data_type);
+ self
+ }
}
impl<R, V> ArrayBuilder for PrimitiveRunBuilder<R, V>
@@ -260,10 +274,12 @@ where
#[cfg(test)]
mod tests {
+ use arrow_schema::DataType;
+
use crate::builder::PrimitiveRunBuilder;
use crate::cast::AsArray;
- use crate::types::{Int16Type, UInt32Type};
- use crate::{Array, UInt32Array};
+ use crate::types::{Decimal128Type, Int16Type, TimestampMicrosecondType,
UInt32Type};
+ use crate::{Array, Decimal128Array, TimestampMicrosecondArray,
UInt32Array};
#[test]
fn test_primitive_ree_array_builder() {
@@ -312,6 +328,51 @@ mod tests {
);
}
+ #[test]
+ #[should_panic]
+ fn test_override_data_type_invalid() {
+ PrimitiveRunBuilder::<Int16Type,
UInt32Type>::new().with_data_type(DataType::UInt64);
+ }
+
+ #[test]
+ fn test_override_data_type() {
+ // Noop.
+ PrimitiveRunBuilder::<Int16Type,
UInt32Type>::new().with_data_type(DataType::UInt32);
+
+ // Setting scale & precision.
+ let mut builder = PrimitiveRunBuilder::<Int16Type,
Decimal128Type>::new()
+ .with_data_type(DataType::Decimal128(1, 2));
+ builder.append_value(123);
+ let array = builder.finish();
+ let array = array.downcast::<Decimal128Array>().unwrap();
+ let values = array.values();
+ assert_eq!(values.precision(), 1);
+ assert_eq!(values.scale(), 2);
+
+ // Setting timezone.
+ let mut builder = PrimitiveRunBuilder::<Int16Type,
TimestampMicrosecondType>::new()
+ .with_data_type(DataType::Timestamp(
+ arrow_schema::TimeUnit::Microsecond,
+ Some("Europe/Paris".into()),
+ ));
+ builder.append_value(1);
+ let array = builder.finish();
+ let array = array.downcast::<TimestampMicrosecondArray>().unwrap();
+ let values = array.values();
+ assert_eq!(values.timezone(), Some("Europe/Paris"));
+ }
+
+ #[test]
+ fn test_finish_cloned_keeps_data_type() {
+ let data_type = DataType::Decimal128(5, 2);
+ let mut builder = PrimitiveRunBuilder::<Int16Type,
Decimal128Type>::new()
+ .with_data_type(data_type.clone());
+ // Leave the run open, so `finish_cloned` has to add it to a copy of
the values
+ builder.append_value(123);
+ let array = builder.finish_cloned();
+ assert_eq!(array.values().data_type(), &data_type);
+ }
+
#[test]
fn test_reuse_after_finish() {
let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
