You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2022/10/18 18:50:00 UTC

[arrow-rs] branch master updated: Simplify ListArray::from_iter_primitive (#2886)

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new bf64ab528 Simplify ListArray::from_iter_primitive (#2886)
bf64ab528 is described below

commit bf64ab528234bf27c3e09985d35cda0e2718fd45
Author: Raphael Taylor-Davies <17...@users.noreply.github.com>
AuthorDate: Wed Oct 19 07:49:54 2022 +1300

    Simplify ListArray::from_iter_primitive (#2886)
---
 arrow-array/src/array/list_array.rs | 63 +++++++++++++------------------------
 1 file changed, 21 insertions(+), 42 deletions(-)

diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs
index cdc7531d9..0db40a796 100644
--- a/arrow-array/src/array/list_array.rs
+++ b/arrow-array/src/array/list_array.rs
@@ -16,12 +16,12 @@
 // under the License.
 
 use crate::array::make_array;
+use crate::builder::{GenericListBuilder, PrimitiveBuilder};
 use crate::{
-    builder::BooleanBufferBuilder, iterator::GenericListArrayIter, print_long_array,
-    raw_pointer::RawPtrBox, Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType,
-    PrimitiveArray,
+    iterator::GenericListArrayIter, print_long_array, raw_pointer::RawPtrBox, Array,
+    ArrayAccessor, ArrayRef, ArrowPrimitiveType,
 };
-use arrow_buffer::{ArrowNativeType, MutableBuffer};
+use arrow_buffer::ArrowNativeType;
 use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field};
 use num::Integer;
@@ -157,47 +157,26 @@ impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
     pub fn from_iter_primitive<T, P, I>(iter: I) -> Self
     where
         T: ArrowPrimitiveType,
-        P: AsRef<[Option<<T as ArrowPrimitiveType>::Native>]>
-            + IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
+        P: IntoIterator<Item = Option<<T as ArrowPrimitiveType>::Native>>,
         I: IntoIterator<Item = Option<P>>,
     {
-        let iterator = iter.into_iter();
-        let (lower, _) = iterator.size_hint();
-
-        let mut offsets =
-            MutableBuffer::new((lower + 1) * std::mem::size_of::<OffsetSize>());
-        let mut length_so_far = OffsetSize::zero();
-        offsets.push(length_so_far);
-
-        let mut null_buf = BooleanBufferBuilder::new(lower);
-
-        let values: PrimitiveArray<T> = iterator
-            .filter_map(|maybe_slice| {
-                // regardless of whether the item is Some, the offsets and null buffers must be updated.
-                match &maybe_slice {
-                    Some(x) => {
-                        length_so_far +=
-                            OffsetSize::from_usize(x.as_ref().len()).unwrap();
-                        null_buf.append(true);
+        let iter = iter.into_iter();
+        let size_hint = iter.size_hint().0;
+        let mut builder =
+            GenericListBuilder::with_capacity(PrimitiveBuilder::<T>::new(), size_hint);
+
+        for i in iter {
+            match i {
+                Some(p) => {
+                    for t in p {
+                        builder.values().append_option(t);
                     }
-                    None => null_buf.append(false),
-                };
-                offsets.push(length_so_far);
-                maybe_slice
-            })
-            .flatten()
-            .collect();
-
-        let field = Box::new(Field::new("item", T::DATA_TYPE, true));
-        let data_type = Self::DATA_TYPE_CONSTRUCTOR(field);
-        let array_data = ArrayData::builder(data_type)
-            .len(null_buf.len())
-            .add_buffer(offsets.into())
-            .add_child_data(values.into_data())
-            .null_bit_buffer(Some(null_buf.into()));
-        let array_data = unsafe { array_data.build_unchecked() };
-
-        Self::from(array_data)
+                    builder.append(true);
+                }
+                None => builder.append(false),
+            }
+        }
+        builder.finish()
     }
 }