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 ae21c899ac Replace `ArrayData` with direct `Array` construction in
`arrow-row` (#10229)
ae21c899ac is described below
commit ae21c899ace6e8cca364194abc790583710b9690
Author: Jeffrey Vo <[email protected]>
AuthorDate: Tue Jun 30 21:28:55 2026 +0900
Replace `ArrayData` with direct `Array` construction in `arrow-row` (#10229)
# 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 #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.
-->
Avoid indirection via `ArrayData` construction for potentially some
small performance benefits.
# 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.
-->
Switch all usages of `ArrayData` to direct `Array` construction in
`arrow-row`.
# 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-row/src/fixed.rs | 99 +++++++++++++----------------------------------
arrow-row/src/lib.rs | 23 ++++++-----
arrow-row/src/list.rs | 78 +++++++++++++++----------------------
arrow-row/src/variable.rs | 61 ++++++++++++-----------------
4 files changed, 94 insertions(+), 167 deletions(-)
diff --git a/arrow-row/src/fixed.rs b/arrow-row/src/fixed.rs
index 422c0c1a32..4268c42ace 100644
--- a/arrow-row/src/fixed.rs
+++ b/arrow-row/src/fixed.rs
@@ -17,13 +17,10 @@
use crate::array::PrimitiveArray;
use crate::null_sentinel;
-use arrow_array::builder::BufferBuilder;
use arrow_array::{ArrowPrimitiveType, BooleanArray, FixedSizeBinaryArray};
use arrow_buffer::{
- ArrowNativeType, BooleanBuffer, Buffer, IntervalDayTime,
IntervalMonthDayNano, MutableBuffer,
- NullBuffer, bit_util, i256,
+ BooleanBuffer, IntervalDayTime, IntervalMonthDayNano, MutableBuffer,
NullBuffer, bit_util, i256,
};
-use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{DataType, SortOptions};
use half::f16;
@@ -364,7 +361,6 @@ pub fn decode_bool(rows: &mut [&[u8]], options:
SortOptions) -> BooleanArray {
let len = rows.len();
- let mut null_count = 0;
let mut nulls = MutableBuffer::new(bit_util::ceil(len, 64) * 8);
let mut values = MutableBuffer::new(bit_util::ceil(len, 64) * 8);
@@ -377,7 +373,6 @@ pub fn decode_bool(rows: &mut [&[u8]], options:
SortOptions) -> BooleanArray {
for bit_idx in 0..64 {
let i = split_off(&mut rows[bit_idx + chunk * 64], 2);
let (null, value) = (i[0] == 1, i[1] == true_val);
- null_count += !null as usize;
null_packed |= (null as u64) << bit_idx;
values_packed |= (value as u64) << bit_idx;
}
@@ -393,7 +388,6 @@ pub fn decode_bool(rows: &mut [&[u8]], options:
SortOptions) -> BooleanArray {
for bit_idx in 0..remainder {
let i = split_off(&mut rows[bit_idx + chunks * 64], 2);
let (null, value) = (i[0] == 1, i[1] == true_val);
- null_count += !null as usize;
null_packed |= (null as u64) << bit_idx;
values_packed |= (value as u64) << bit_idx;
}
@@ -402,61 +396,19 @@ pub fn decode_bool(rows: &mut [&[u8]], options:
SortOptions) -> BooleanArray {
values.push(values_packed);
}
- let builder = ArrayDataBuilder::new(DataType::Boolean)
- .len(rows.len())
- .null_count(null_count)
- .add_buffer(values.into())
- .null_bit_buffer(Some(nulls.into()));
+ let nulls = NullBuffer::new(BooleanBuffer::new(nulls.into(), 0, len));
+ let nulls = (nulls.null_count() > 0).then_some(nulls);
- // SAFETY:
- // Buffers are the correct length
- unsafe { BooleanArray::from(builder.build_unchecked()) }
+ BooleanArray::new(BooleanBuffer::new(values.into(), 0, len), nulls)
}
/// Decodes a single byte from each row, interpreting `0x01` as a valid value
-/// and all other values as a null
-///
-/// Returns the null count and null buffer
-pub fn decode_nulls(rows: &[&[u8]]) -> (usize, Buffer) {
- let mut null_count = 0;
- let buffer = MutableBuffer::collect_bool(rows.len(), |idx| {
- let valid = rows[idx][0] == 1;
- null_count += !valid as usize;
- valid
- })
- .into();
- (null_count, buffer)
-}
-
-/// Decodes a `ArrayData` from rows based on the provided
`FixedLengthEncoding` `T`
-///
-/// # Safety
-///
-/// `data_type` must be appropriate native type for `T`
-unsafe fn decode_fixed<T: FixedLengthEncoding + ArrowNativeType>(
- rows: &mut [&[u8]],
- data_type: DataType,
- options: SortOptions,
-) -> ArrayData {
- let len = rows.len();
-
- let mut values = BufferBuilder::<T>::new(len);
- let (null_count, nulls) = decode_nulls(rows);
-
- for row in rows {
- let i = split_off(row, T::ENCODED_LEN);
- let value = T::Encoded::from_slice(&i[1..], options.descending);
- values.append(T::decode(value));
- }
-
- let builder = ArrayDataBuilder::new(data_type)
- .len(len)
- .null_count(null_count)
- .add_buffer(values.finish())
- .null_bit_buffer(Some(nulls));
-
- // SAFETY: Buffers correct length
- unsafe { builder.build_unchecked() }
+/// and all other values as a null. Optionally returns `None` if there are no
+/// null values.
+pub fn decode_nulls(rows: &[&[u8]]) -> Option<NullBuffer> {
+ let nulls = BooleanBuffer::collect_bool(rows.len(), |idx| rows[idx][0] ==
1);
+ let nulls = NullBuffer::new(nulls);
+ (nulls.null_count() > 0).then_some(nulls)
}
/// Decodes a `PrimitiveArray` from rows
@@ -469,9 +421,21 @@ where
T::Native: FixedLengthEncoding,
{
assert!(PrimitiveArray::<T>::is_compatible(&data_type));
- // SAFETY:
- // Validated data type above
- unsafe { decode_fixed::<T::Native>(rows, data_type, options).into() }
+
+ let nulls = decode_nulls(rows);
+ let values = rows
+ .iter_mut()
+ .map(|row| {
+ let i = split_off(row, T::Native::ENCODED_LEN);
+ let value = <T::Native as
FixedLengthEncoding>::Encoded::from_slice(
+ &i[1..],
+ options.descending,
+ );
+ T::Native::decode(value)
+ })
+ .collect::<Vec<_>>();
+
+ PrimitiveArray::new(values.into(), nulls).with_data_type(data_type)
}
/// Decodes a `FixedLengthBinary` from rows
@@ -483,13 +447,11 @@ pub fn decode_fixed_size_binary(
size: i32,
options: SortOptions,
) -> FixedSizeBinaryArray {
- let len = rows.len();
-
if size < 0 {
panic!("cannot decode FixedSizeBinary({size})");
}
let mut values = MutableBuffer::new(size as usize * rows.len());
- let (null_count, nulls) = decode_nulls(rows);
+ let nulls = decode_nulls(rows);
let encoded_len = size as usize + 1;
@@ -504,12 +466,5 @@ pub fn decode_fixed_size_binary(
}
}
- let builder = ArrayDataBuilder::new(DataType::FixedSizeBinary(size))
- .len(len)
- .null_count(null_count)
- .add_buffer(values.into())
- .null_bit_buffer(Some(nulls));
-
- // SAFETY: Buffers correct length
- unsafe { builder.build_unchecked().into() }
+ FixedSizeBinaryArray::new(size, values.into(), nulls)
}
diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs
index 379c7c2425..182481f42b 100644
--- a/arrow-row/src/lib.rs
+++ b/arrow-row/src/lib.rs
@@ -169,7 +169,6 @@ use arrow_array::cast::*;
use arrow_array::types::{ArrowDictionaryKeyType, ByteArrayType, ByteViewType};
use arrow_array::*;
use arrow_buffer::{ArrowNativeType, Buffer, OffsetBuffer, ScalarBuffer};
-use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::*;
use variable::{decode_binary_view, decode_string_view};
@@ -2119,17 +2118,16 @@ unsafe fn decode_column(
cols.into_iter().next().unwrap()
}
Codec::Struct(converter, _) => {
- let (null_count, nulls) = fixed::decode_nulls(rows);
+ let nulls = fixed::decode_nulls(rows);
rows.iter_mut().for_each(|row| *row = &row[1..]);
let children = unsafe { converter.convert_raw(rows, validate_utf8)
}?;
- let child_data: Vec<ArrayData> = children.iter().map(|c|
c.to_data()).collect();
// Since RowConverter flattens certain data types (i.e.
Dictionary),
// we need to use updated data type instead of original field
let corrected_fields: Vec<Field> = match &field.data_type {
DataType::Struct(struct_fields) => struct_fields
.iter()
- .zip(child_data.iter())
+ .zip(children.iter())
.map(|(orig_field, child_array)| {
orig_field
.as_ref()
@@ -2139,14 +2137,15 @@ unsafe fn decode_column(
.collect(),
_ => unreachable!("Only Struct types should be corrected
here"),
};
- let corrected_struct_type =
DataType::Struct(corrected_fields.into());
- let builder = ArrayDataBuilder::new(corrected_struct_type)
- .len(rows.len())
- .null_count(null_count)
- .null_bit_buffer(Some(nulls))
- .child_data(child_data);
-
- Arc::new(StructArray::from(unsafe { builder.build_unchecked() }))
+
+ Arc::new(unsafe {
+ StructArray::new_unchecked_with_length(
+ corrected_fields.into(),
+ children,
+ nulls,
+ rows.len(),
+ )
+ })
}
Codec::List(converter) => match &field.data_type {
DataType::List(_) => {
diff --git a/arrow-row/src/list.rs b/arrow-row/src/list.rs
index 843101673f..7f76e88cec 100644
--- a/arrow-row/src/list.rs
+++ b/arrow-row/src/list.rs
@@ -21,9 +21,8 @@ use arrow_array::{
new_null_array,
};
use arrow_buffer::{
- ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer,
ScalarBuffer,
+ ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer,
ScalarBuffer,
};
-use arrow_data::ArrayDataBuilder;
use arrow_schema::{ArrowError, DataType, SortOptions};
use std::{ops::Range, sync::Arc};
@@ -129,12 +128,7 @@ pub unsafe fn decode<O: OffsetSizeTrait>(
}
O::from_usize(offset).expect("overflow");
- let mut null_count = 0;
- let nulls = MutableBuffer::collect_bool(rows.len(), |x| {
- let valid = rows[x][0] != null_sentinel(opts);
- null_count += !valid as usize;
- valid
- });
+ let nulls = crate::variable::decode_nulls_sentinel(rows, opts);
let mut values_offsets = Vec::with_capacity(offset);
let mut values_bytes = Vec::with_capacity(values_bytes);
@@ -167,37 +161,34 @@ pub unsafe fn decode<O: OffsetSizeTrait>(
})
.collect();
- let child = unsafe { converter.convert_raw(&mut child_rows, validate_utf8)
}?;
+ let mut child = unsafe { converter.convert_raw(&mut child_rows,
validate_utf8) }?;
assert_eq!(child.len(), 1);
-
- let child_data = child[0].to_data();
+ let child = child.pop().unwrap();
// Since RowConverter flattens certain data types (i.e. Dictionary),
// we need to use updated data type instead of original field
- let corrected_type = match &field.data_type {
- DataType::List(inner_field) => DataType::List(Arc::new(
+ let corrected_inner_field = match &field.data_type {
+ DataType::List(inner_field) => Arc::new(
inner_field
.as_ref()
.clone()
- .with_data_type(child_data.data_type().clone()),
- )),
- DataType::LargeList(inner_field) => DataType::LargeList(Arc::new(
+ .with_data_type(child.data_type().clone()),
+ ),
+ DataType::LargeList(inner_field) => Arc::new(
inner_field
.as_ref()
.clone()
- .with_data_type(child_data.data_type().clone()),
- )),
+ .with_data_type(child.data_type().clone()),
+ ),
_ => unreachable!(),
};
- let builder = ArrayDataBuilder::new(corrected_type)
- .len(rows.len())
- .null_count(null_count)
- .null_bit_buffer(Some(nulls.into()))
- .add_buffer(Buffer::from_vec(offsets))
- .add_child_data(child_data);
-
- Ok(GenericListArray::from(unsafe { builder.build_unchecked() }))
+ GenericListArray::try_new(
+ corrected_inner_field,
+ OffsetBuffer::new(offsets.into()),
+ child,
+ nulls,
+ )
}
pub fn compute_lengths_fixed_size_list(
@@ -267,19 +258,16 @@ pub unsafe fn decode_fixed_size_list(
value_length: usize,
) -> Result<FixedSizeListArray, ArrowError> {
let list_type = &field.data_type;
- let element_type = match list_type {
- DataType::FixedSizeList(element_field, _) => element_field.data_type(),
- _ => {
- return Err(ArrowError::InvalidArgumentError(format!(
- "Expected FixedSizeListArray, found: {list_type}",
- )));
- }
+ let DataType::FixedSizeList(element_field, size) = list_type else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Expected FixedSizeListArray, found: {list_type}",
+ )));
};
- let len = rows.len();
- let (null_count, nulls) = fixed::decode_nulls(rows);
+ let nulls = fixed::decode_nulls(rows);
- let null_element_encoded =
converter.convert_columns(&[new_null_array(element_type, 1)])?;
+ let null_element_encoded =
+ converter.convert_columns(&[new_null_array(element_field.data_type(),
1)])?;
let null_element_encoded = null_element_encoded.row(0);
let null_element_slice = null_element_encoded.as_ref();
@@ -304,17 +292,15 @@ pub unsafe fn decode_fixed_size_list(
*row = &row[row_offset..]; // Update row for the next decoder
}
- let children = unsafe { converter.convert_raw(&mut child_rows,
validate_utf8) }?;
- let child_data = children.iter().map(|c| c.to_data()).collect();
- let builder = ArrayDataBuilder::new(list_type.clone())
- .len(len)
- .null_count(null_count)
- .null_bit_buffer(Some(nulls))
- .child_data(child_data);
+ let mut children = unsafe { converter.convert_raw(&mut child_rows,
validate_utf8) }?;
+ assert_eq!(children.len(), 1);
- Ok(FixedSizeListArray::from(unsafe {
- builder.build_unchecked()
- }))
+ Ok(FixedSizeListArray::new(
+ Arc::clone(element_field),
+ *size,
+ children.pop().unwrap(),
+ nulls,
+ ))
}
/// Computes the encoded length for a single list element given its child rows.
diff --git a/arrow-row/src/variable.rs b/arrow-row/src/variable.rs
index 9d557c5746..aeb7f888b0 100644
--- a/arrow-row/src/variable.rs
+++ b/arrow-row/src/variable.rs
@@ -20,9 +20,11 @@ use arrow_array::builder::BufferBuilder;
use arrow_array::types::ByteArrayType;
use arrow_array::*;
use arrow_buffer::bit_util::ceil;
-use arrow_buffer::{ArrowNativeType, MutableBuffer};
-use arrow_data::{ArrayDataBuilder, MAX_INLINE_VIEW_LEN};
-use arrow_schema::{DataType, SortOptions};
+use arrow_buffer::{
+ ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer,
ScalarBuffer,
+};
+use arrow_data::MAX_INLINE_VIEW_LEN;
+use arrow_schema::SortOptions;
use builder::make_view;
/// The block size of the variable length encoding
@@ -69,6 +71,15 @@ pub(crate) fn non_null_padded_length(len: usize) -> usize {
}
}
+/// Decodes a single byte from each row, where a valid value is determined via
+/// a configurable sentinel. Optionally returns `None` if there are no null
values.
+pub(crate) fn decode_nulls_sentinel(rows: &[&[u8]], options: SortOptions) ->
Option<NullBuffer> {
+ let null_sentinel = null_sentinel(options);
+ let nulls = BooleanBuffer::collect_bool(rows.len(), |x| rows[x][0] !=
null_sentinel);
+ let nulls = NullBuffer::new(nulls);
+ (nulls.null_count() > 0).then_some(nulls)
+}
+
/// Variable length values are encoded as
///
/// - single `0_u8` if null
@@ -270,12 +281,7 @@ pub fn decode_binary<I: OffsetSizeTrait>(
options: SortOptions,
) -> GenericBinaryArray<I> {
let len = rows.len();
- let mut null_count = 0;
- let nulls = MutableBuffer::collect_bool(len, |x| {
- let valid = rows[x][0] != null_sentinel(options);
- null_count += !valid as usize;
- valid
- });
+ let nulls = decode_nulls_sentinel(rows, options);
let values_capacity = rows.iter().map(|row| decoded_len(row,
options)).sum();
let mut offsets = BufferBuilder::<I>::new(len + 1);
@@ -292,21 +298,15 @@ pub fn decode_binary<I: OffsetSizeTrait>(
values.as_slice_mut().iter_mut().for_each(|o| *o = !*o)
}
- let d = match I::IS_LARGE {
- true => DataType::LargeBinary,
- false => DataType::Binary,
- };
-
- let builder = ArrayDataBuilder::new(d)
- .len(len)
- .null_count(null_count)
- .null_bit_buffer(Some(nulls.into()))
- .add_buffer(offsets.finish())
- .add_buffer(values.into());
-
// SAFETY:
// Valid by construction above
- unsafe { GenericBinaryArray::from(builder.build_unchecked()) }
+ unsafe {
+ GenericBinaryArray::new_unchecked(
+ OffsetBuffer::new(ScalarBuffer::from(offsets)),
+ values.into(),
+ nulls,
+ )
+ }
}
fn decode_binary_view_inner(
@@ -317,13 +317,7 @@ fn decode_binary_view_inner(
let len = rows.len();
let inline_str_max_len = MAX_INLINE_VIEW_LEN as usize;
- let mut null_count = 0;
-
- let nulls = MutableBuffer::collect_bool(len, |x| {
- let valid = rows[x][0] != null_sentinel(options);
- null_count += !valid as usize;
- valid
- });
+ let nulls = decode_nulls_sentinel(rows, options);
// If we are validating UTF-8, decode all string values (including short
strings)
// into the values buffer and validate UTF-8 once. If not validating,
@@ -381,16 +375,9 @@ fn decode_binary_view_inner(
std::str::from_utf8(values.as_slice()).unwrap();
}
- let builder = ArrayDataBuilder::new(DataType::BinaryView)
- .len(len)
- .null_count(null_count)
- .null_bit_buffer(Some(nulls.into()))
- .add_buffer(views.finish())
- .add_buffer(values.into());
-
// SAFETY:
// Valid by construction above
- unsafe { BinaryViewArray::from(builder.build_unchecked()) }
+ unsafe { BinaryViewArray::new_unchecked(views.into(), [values.into()],
nulls) }
}
/// Decodes a binary view array from `rows` with the provided `options`
