(arrow-rs) branch master updated: fix: ensure take_fixed_size_list can handle null indices (#5170)

[wj...@apache.org]

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

wjones127 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 f4bad6871f fix: ensure take_fixed_size_list can handle null indices (#5170)
f4bad6871f is described below

commit f4bad6871f9f64f2b7c0367ae198aff32c3c61b1
Author: Weston Pace <we...@gmail.com>
AuthorDate: Tue Dec 5 16:07:12 2023 -0800

    fix: ensure take_fixed_size_list can handle null indices (#5170)
    
    * fix: ensure take_fixed_size_list can handle null indices
    
    * chore: apply clippy suggestion
    
    * Apply suggestions from code review
    
    Co-authored-by: Will Jones <wi...@gmail.com>
    
    * Applying suggetions from review
    
    * Using a builder, per review suggestion
    
    * Apply suggestions from code review
    
    Co-authored-by: Will Jones <wi...@gmail.com>
    
    * Cast length to usize to avoid compile error
    
    ---------
    
    Co-authored-by: Will Jones <wi...@gmail.com>
---
 arrow-select/src/take.rs | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index d47b884ae3..44269e3875 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -19,7 +19,7 @@
 
 use std::sync::Arc;
 
-use arrow_array::builder::BufferBuilder;
+use arrow_array::builder::{BufferBuilder, UInt32Builder};
 use arrow_array::cast::AsArray;
 use arrow_array::types::*;
 use arrow_array::*;
@@ -689,7 +689,7 @@ fn take_value_indices_from_fixed_size_list<IndexType>(
 where
     IndexType: ArrowPrimitiveType,
 {
-    let mut values = vec![];
+    let mut values = UInt32Builder::with_capacity(length as usize * indices.len());
 
     for i in 0..indices.len() {
         if indices.is_valid(i) {
@@ -699,11 +699,16 @@ where
                 .ok_or_else(|| ArrowError::ComputeError("Cast to usize failed".to_string()))?;
             let start = list.value_offset(index) as <UInt32Type as ArrowPrimitiveType>::Native;
 
-            values.extend(start..start + length);
+            // Safety: Range always has known length.
+            unsafe {
+                values.append_trusted_len_iter(start..start + length);
+            }
+        } else {
+            values.append_nulls(length as usize);
         }
     }
 
-    Ok(PrimitiveArray::<UInt32Type>::from(values))
+    Ok(values.finish())
 }
 
 /// To avoid generating take implementations for every index type, instead we
@@ -1985,6 +1990,23 @@ mod tests {
         assert_eq!(&values, &[Some(23), Some(4), None, None])
     }
 
+    #[test]
+    fn test_take_fixed_size_list_null_indices() {
+        let indices = Int32Array::from_iter([Some(0), None]);
+        let values = Arc::new(Int32Array::from(vec![0, 1, 2, 3]));
+        let arr_field = Arc::new(Field::new("item", values.data_type().clone(), true));
+        let values = FixedSizeListArray::try_new(arr_field, 2, values, None).unwrap();
+
+        let r = take(&values, &indices, None).unwrap();
+        let values = r
+            .as_fixed_size_list()
+            .values()
+            .as_primitive::<Int32Type>()
+            .into_iter()
+            .collect::<Vec<_>>();
+        assert_eq!(values, &[Some(0), Some(1), None, None])
+    }
+
     #[test]
     fn test_take_bytes_null_indices() {
         let indices = Int32Array::new(