[arrow-rs] branch master updated: Fix for primitive and boolean take kernel for nullable indices with an offset (#509)

[al...@apache.org]

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

alamb 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 b63c407  Fix for primitive and boolean take kernel for nullable indices with an offset (#509)
b63c407 is described below

commit b63c407b3ff4ff7bb7767091adcf553180c8e723
Author: Jörn Horstmann <gi...@jhorstmann.net>
AuthorDate: Wed Jun 30 22:41:24 2021 +0200

    Fix for primitive and boolean take kernel for nullable indices with an offset (#509)
    
    * Fix for take kernel with nullable indices and nonnull values
    
    * Fix for boolean take kernel when indices have an offset
    
    * Use different values for data so they cannot be confused with the indices
---
 arrow/src/compute/kernels/take.rs | 91 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 88 insertions(+), 3 deletions(-)

diff --git a/arrow/src/compute/kernels/take.rs b/arrow/src/compute/kernels/take.rs
index 66bfd61..bf9d1df 100644
--- a/arrow/src/compute/kernels/take.rs
+++ b/arrow/src/compute/kernels/take.rs
@@ -357,7 +357,13 @@ where
     // Soundness: `slice.map` is `TrustedLen`.
     let buffer = unsafe { Buffer::try_from_trusted_len_iter(values)? };
 
-    Ok((buffer, indices.data_ref().null_buffer().cloned()))
+    Ok((
+        buffer,
+        indices
+            .data_ref()
+            .null_buffer()
+            .map(|b| b.bit_slice(indices.offset(), indices.len())),
+    ))
 }
 
 // take implementation when both values and indices contain nulls
@@ -516,7 +522,7 @@ where
         nulls = match indices.data_ref().null_buffer() {
             Some(buffer) => Some(buffer_bin_and(
                 buffer,
-                0,
+                indices.offset(),
                 &null_buf.into(),
                 0,
                 indices.len(),
@@ -805,6 +811,24 @@ mod tests {
         Ok(())
     }
 
+    fn test_take_primitive_arrays_non_null<T>(
+        data: Vec<T::Native>,
+        index: &UInt32Array,
+        options: Option<TakeOptions>,
+        expected_data: Vec<Option<T::Native>>,
+    ) -> Result<()>
+    where
+        T: ArrowPrimitiveType,
+        PrimitiveArray<T>: From<Vec<T::Native>>,
+        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
+    {
+        let output = PrimitiveArray::<T>::from(data);
+        let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
+        let output = take(&output, index, options)?;
+        assert_eq!(&output, &expected);
+        Ok(())
+    }
+
     fn test_take_impl_primitive_arrays<T, I>(
         data: Vec<Option<T::Native>>,
         index: &PrimitiveArray<I>,
@@ -877,6 +901,48 @@ mod tests {
     }
 
     #[test]
+    fn test_take_primitive_nullable_indices_non_null_values_with_offset() {
+        let index =
+            UInt32Array::from(vec![Some(0), Some(1), Some(2), Some(3), None, None]);
+        let index = index.slice(2, 4);
+        let index = index.as_any().downcast_ref::<UInt32Array>().unwrap();
+
+        assert_eq!(
+            index,
+            &UInt32Array::from(vec![Some(2), Some(3), None, None])
+        );
+
+        test_take_primitive_arrays_non_null::<Int64Type>(
+            vec![0, 10, 20, 30, 40, 50],
+            &index,
+            None,
+            vec![Some(20), Some(30), None, None],
+        )
+        .unwrap();
+    }
+
+    #[test]
+    fn test_take_primitive_nullable_indices_nullable_values_with_offset() {
+        let index =
+            UInt32Array::from(vec![Some(0), Some(1), Some(2), Some(3), None, None]);
+        let index = index.slice(2, 4);
+        let index = index.as_any().downcast_ref::<UInt32Array>().unwrap();
+
+        assert_eq!(
+            index,
+            &UInt32Array::from(vec![Some(2), Some(3), None, None])
+        );
+
+        test_take_primitive_arrays::<Int64Type>(
+            vec![None, None, Some(20), Some(30), Some(40), Some(50)],
+            &index,
+            None,
+            vec![Some(20), Some(30), None, None],
+        )
+        .unwrap();
+    }
+
+    #[test]
     fn test_take_primitive() {
         let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
 
@@ -1100,7 +1166,7 @@ mod tests {
     }
 
     #[test]
-    fn test_take_primitive_bool() {
+    fn test_take_bool() {
         let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
         // boolean
         test_take_boolean_arrays(
@@ -1111,6 +1177,25 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_take_bool_with_offset() {
+        let index =
+            UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2), None]);
+        let index = index.slice(2, 4);
+        let index = index
+            .as_any()
+            .downcast_ref::<PrimitiveArray<UInt32Type>>()
+            .unwrap();
+
+        // boolean
+        test_take_boolean_arrays(
+            vec![Some(false), None, Some(true), Some(false), None],
+            &index,
+            None,
+            vec![None, Some(false), Some(true), None],
+        );
+    }
+
     fn _test_take_string<'a, K: 'static>()
     where
         K: Array + PartialEq + From<Vec<Option<&'a str>>>,