fix(rust): Centralize FixedSizeList outer-to-child layout arithmetic and fix under-reservation

crates/polars-arrow/src/array/fixed_size_list/builder.rs

@@ -1,6 +1,6 @@
 use polars_utils::IdxSize;
 
-use super::FixedSizeListArray;
+use super::{FixedSizeListArray, child_length, child_offset};
 use crate::array::builder::{ArrayBuilder, ShareStrategy, StaticArrayBuilder};
 use crate::bitmap::OptBitmapBuilder;
 use crate::datatypes::ArrowDataType;
@@ -32,7 +32,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
     }
 
     fn reserve(&mut self, additional: usize) {
-        self.inner_builder.reserve(additional * self.size);
+        self.inner_builder
+            .reserve(child_length(additional, self.size));
         self.validity.reserve(additional);
     }
 
@@ -55,7 +56,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
     }
 
     fn extend_nulls(&mut self, length: usize) {
-        self.inner_builder.extend_nulls(length * self.size);
+        self.inner_builder
+            .extend_nulls(child_length(length, self.size));
         self.validity.extend_constant(length, false);
         self.length += length;
     }
@@ -69,8 +71,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
     ) {
         self.inner_builder.subslice_extend(
             &**other.values(),
-            start * self.size,
-            length * self.size,
+            child_offset(start, self.size),
+            child_length(length, self.size),
             share,
         );
         self.validity
@@ -87,11 +89,12 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
         share: ShareStrategy,
     ) {
         let other_values = &**other.values();
-        self.inner_builder.reserve(repeats * length * self.size);
+        self.inner_builder
+            .reserve(child_length(repeats * length, self.size));
         for outer_idx in start..start + length {
             self.inner_builder.subslice_extend_repeated(
                 other_values,
-                outer_idx * self.size,
+                child_offset(outer_idx, self.size),
                 self.size,
                 repeats,
                 share,
@@ -114,7 +117,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
         share: ShareStrategy,
     ) {
         let other_values = &**other.values();
-        self.inner_builder.reserve(idxs.len() * self.size);
+        self.inner_builder
+            .reserve(child_length(idxs.len(), self.size));
 
         // Group consecutive indices into larger copies.
         let mut group_start = 0;
@@ -128,8 +132,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
             }
             self.inner_builder.subslice_extend(
                 other_values,
-                start_idx * self.size,
-                group_len * self.size,
+                child_offset(start_idx, self.size),
+                child_length(group_len, self.size),
                 share,
             );
             group_start += group_len;
@@ -147,7 +151,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
         share: ShareStrategy,
     ) {
         let other_values = &**other.values();
-        self.inner_builder.reserve(idxs.len() * self.size);
+        self.inner_builder
+            .reserve(child_length(idxs.len(), self.size));
 
         // Group consecutive indices into larger copies.
         let mut group_start = 0;
@@ -166,8 +171,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
 
                 self.inner_builder.subslice_extend(
                     other_values,
-                    start_idx * self.size,
-                    group_len * self.size,
+                    child_offset(start_idx, self.size),
+                    child_length(group_len, self.size),
                     share,
                 );
             } else {
@@ -177,7 +182,8 @@ impl<B: ArrayBuilder> StaticArrayBuilder for FixedSizeListArrayBuilder<B> {
                     group_len += 1;
                 }
 
-                self.inner_builder.extend_nulls(group_len * self.size);
+                self.inner_builder
+                    .extend_nulls(child_length(group_len, self.size));
             }
             group_start += group_len;
         }

crates/polars-arrow/src/array/fixed_size_list/mod.rs

@@ -18,6 +18,16 @@ pub mod proptest;
 
 use crate::datatypes::reshape::{Dimension, ReshapeDimension};
 
+#[inline(always)]
+pub(crate) const fn child_offset(outer_offset: usize, size: usize) -> usize {
+    outer_offset * size
+}
+
+#[inline(always)]
+pub(crate) const fn child_length(outer_length: usize, size: usize) -> usize {
+    outer_length * size
+}
+
 /// The Arrow's equivalent to an immutable `Vec<Option<[T; size]>>` where `T` is an Arrow type.
 /// Cloning and slicing this struct is `O(1)`.
 #[derive(Clone)]
@@ -124,7 +134,7 @@ impl FixedSizeListArray {
     pub fn new_null(dtype: ArrowDataType, length: usize) -> Self {
         let (field, size) = Self::get_child_and_size(&dtype);
 
-        let values = new_null_array(field.dtype().clone(), length * size);
+        let values = new_null_array(field.dtype().clone(), child_length(length, size));
         Self::new(dtype, length, values, Some(Bitmap::new_zeroed(length)))
     }
 
@@ -259,8 +269,10 @@ impl FixedSizeListArray {
             .take()
             .map(|bitmap| bitmap.sliced_unchecked(offset, length))
             .filter(|bitmap| bitmap.unset_bits() > 0);
-        self.values
-            .slice_unchecked(offset * self.size, length * self.size);
+        self.values.slice_unchecked(
+            child_offset(offset, self.size),
+            child_length(length, self.size),
+        );
         self.length = length;
     }
 
@@ -294,7 +306,7 @@ impl FixedSizeListArray {
     /// panics iff `i >= self.len()`
     #[inline]
     pub fn value(&self, i: usize) -> Box<dyn Array> {
-        self.values.sliced(i * self.size, self.size)
+        self.values.sliced(child_offset(i, self.size), self.size)
     }
 
     /// Returns the `Vec<T>` at position `i`.
@@ -303,7 +315,8 @@ impl FixedSizeListArray {
     /// Caller must ensure that `i < self.len()`
     #[inline]
     pub unsafe fn value_unchecked(&self, i: usize) -> Box<dyn Array> {
-        self.values.sliced_unchecked(i * self.size, self.size)
+        self.values
+            .sliced_unchecked(child_offset(i, self.size), self.size)
     }
 
     /// Returns the element at index `i` or `None` if it is null
@@ -358,8 +371,10 @@ impl Splitable for FixedSizeListArray {
     }
 
     unsafe fn _split_at_unchecked(&self, offset: usize) -> (Self, Self) {
-        let (lhs_values, rhs_values) =
-            unsafe { self.values.split_at_boxed_unchecked(offset * self.size) };
+        let (lhs_values, rhs_values) = unsafe {
+            self.values
+                .split_at_boxed_unchecked(child_offset(offset, self.size))
+        };
         let (lhs_validity, rhs_validity) = unsafe { self.validity.split_at_unchecked(offset) };
 
         let size = self.size;

crates/polars-arrow/src/array/fixed_size_list/mutable.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use polars_error::{PolarsResult, polars_bail};
 use polars_utils::pl_str::PlSmallStr;
 
-use super::FixedSizeListArray;
+use super::{FixedSizeListArray, child_length};
 use crate::array::physical_binary::extend_validity;
 use crate::array::{Array, MutableArray, PushUnchecked, TryExtend, TryExtendFromSelf, TryPush};
 use crate::bitmap::MutableBitmap;
@@ -136,7 +136,7 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
 
     /// Reserves `additional` slots.
     pub fn reserve(&mut self, additional: usize) {
-        self.values.reserve(additional * self.size);
+        self.values.reserve(child_length(additional, self.size));
         if let Some(x) = self.validity.as_mut() {
             x.reserve(additional)
         }

crates/polars/tests/it/arrow/array/fixed_size_list/mod.rs

@@ -127,3 +127,34 @@ fn wrong_dtype() {
         .is_err()
     );
 }
+
+#[test]
+fn builder() {
+    use arrow::array::builder::{ShareStrategy, make_builder};
+
+    let dtype = ArrowDataType::FixedSizeList(
+        Box::new(Field::new("a".into(), ArrowDataType::Int32, true)),
+        2,
+    );
+    let mut builder = make_builder(&dtype);
+
+    builder.reserve(10);
+    builder.extend_nulls(2);
+
+    let arr = data();
+    builder.subslice_extend(&arr, 0, 2, ShareStrategy::Never);
+
+    let result = builder.freeze_reset();
+    let result = result
+        .as_any()
+        .downcast_ref::<FixedSizeListArray>()
+        .unwrap();
+    assert_eq!(result.len(), 4);
+    assert_eq!(result.size(), 2);
+    assert_eq!(
+        result.validity().cloned(),
+        Some([false, false, true, false].into())
+    );
+    assert_eq!(result.value(2).as_ref(), Int32Array::from_slice([10, 20]));
+    assert_eq!(result.value(3).as_ref(), Int32Array::from_slice([0, 0]));
+}

crates/polars/tests/it/arrow/array/fixed_size_list/mutable.rs

@@ -90,3 +90,10 @@ fn extend_from_self() {
 
     assert_eq!(a, b);
 }
+
+#[test]
+fn reserve() {
+    let mut list = MutableFixedSizeListArray::new(MutablePrimitiveArray::<i32>::new(), 3);
+    list.reserve(10);
+    assert!(list.values().capacity() >= 30);
+}