This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] 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 0028aaa593 fix(arrow-select): fix `take` output capacity for sliced
lists (#11360)
0028aaa593 is described below
commit 0028aaa593bed532ce7893c5a0aca17b0a7139e9
Author: Neil Conway <[email protected]>
AuthorDate: Sun Oct 4 02:41:18 2026 +0000
fix(arrow-select): fix `take` output capacity for sliced lists (#11360)
# Which issue does this PR close?
- N/A
# Rationale for this change
`take` on a list or map uses the average input row length as part of its
estimate for sizing its output buffer. However, it calculated the
average row length starting from raw size of the input values buffer,
not the visible span. This could result in computing a very large
"average row length" for sliced lists and maps, and therefore
over-allocating output capacity.
# What changes are included in this PR?
* Use the visible span to compute average row length
* Refactor capacity calculations to avoid duplication
* Add unit test
# Are these changes tested?
Yes; existing tests pass, new test added.
# Are there any significant user-facing changes?
No.
# AI usage
Developed and revised with Claude Code (Opus 5.5), reviewed with Codex.
---
arrow-select/src/take.rs | 60 +++++++++++++++++++++++++++++++++++-------------
1 file changed, 44 insertions(+), 16 deletions(-)
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index ea505762e8..38d7945700 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -1027,6 +1027,12 @@ where
{
let src_offsets = values.value_offsets();
let child_data = values.values().to_data();
+ // The number of child values that the rows use, excluding values outside
a slice
+ let visible_child_len = (values.offsets().last() -
values.offsets().first()).as_usize();
+ // An estimate of the output's child length, from the average length of
the rows
+ let estimated_child_len = visible_child_len
+ .checked_div(values.len())
+ .map_or(0, |avg| avg.saturating_mul(indices.len()));
let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices);
let mut dst_offsets = Vec::with_capacity(indices.len() + 1);
@@ -1040,15 +1046,7 @@ where
let values_buf = &child_data.buffers()[0];
let child_buf_offset = child_data.offset() * bytes_per_value;
- let avg_row_len = child_data
- .len()
- .checked_div(values.len().max(1))
- .unwrap_or(0);
- let mut dst_buf = MutableBuffer::new(
- avg_row_len
- .saturating_mul(indices.len())
- .saturating_mul(bytes_per_value),
- );
+ let mut dst_buf =
MutableBuffer::new(estimated_child_len.saturating_mul(bytes_per_value));
let mut child_len = OffsetType::Native::zero();
@@ -1121,13 +1119,11 @@ where
return GenericListArray::<OffsetType::Native>::try_new(field, offsets,
child, nulls);
}
- let capacity = child_data
- .len()
- .checked_div(values.len())
- .map(|avg| avg * indices.len())
- .unwrap_or_default();
- let mut mutable =
- MutableArrayData::new(vec![&child_data], child_data.null_count() > 0,
capacity);
+ let mut mutable = MutableArrayData::new(
+ vec![&child_data],
+ child_data.null_count() > 0,
+ estimated_child_len,
+ );
match nulls.as_ref().filter(|n| n.null_count() > 0) {
None => {
@@ -3258,6 +3254,38 @@ mod tests {
test_take_sliced_list_with_value_nulls_generic::<i64>();
}
+ #[test]
+ fn test_take_sliced_list_memory_size() {
+ // Takes the same rows from a two-row slice of a large array and from
an unsliced
+ // copy of those two rows. The results should be equal and use the
same memory.
+ fn check(parent: ArrayRef, unsliced: ArrayRef) {
+ let sliced = parent.slice(500, 2);
+ let indices = UInt32Array::from_iter_values((0..100).map(|i| i %
2));
+ let actual = take(&sliced, &indices, None).unwrap();
+ let expected = take(&unsliced, &indices, None).unwrap();
+ assert_eq!(&actual, &expected);
+ assert_eq!(
+ actual.get_array_memory_size(),
+ expected.get_array_memory_size()
+ );
+ }
+
+ // Lists of 10 values per row
+ let list = |rows: usize, values: ArrayRef| -> ArrayRef {
+ let field =
Arc::new(Field::new_list_field(values.data_type().clone(), false));
+ let offsets = OffsetBuffer::from_repeated_length(10, rows);
+ Arc::new(ListArray::new(field, offsets, values, None))
+ };
+ check(
+ list(1000, Arc::new(Int64Array::from(vec![1; 10_000]))),
+ list(2, Arc::new(Int64Array::from(vec![1; 20]))),
+ );
+ check(
+ list(1000, Arc::new(StringArray::from(vec!["a"; 10_000]))),
+ list(2, Arc::new(StringArray::from(vec!["a"; 20]))),
+ );
+ }
+
#[test]
fn test_take_runs() {
let logical_array: Vec<i32> = vec![1_i32, 1, 2, 2, 1, 1, 1, 2, 2, 1,
1, 2, 2];
