Fix take string on sliced indices (#2960)

apache · Oct 28, 2022 · 87ac05b · 87ac05b
1 parent 0a115d6
commit 87ac05b
Showing 1 changed file with 19 additions and 9 deletions.
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
@@ -21,9 +21,7 @@ use std::{ops::AddAssign, sync::Arc};
 
 use arrow_array::types::*;
 use arrow_array::*;
-use arrow_buffer::{
-    bit_util, buffer::buffer_bin_and, ArrowNativeType, Buffer, MutableBuffer,
-};
+use arrow_buffer::{bit_util, ArrowNativeType, Buffer, MutableBuffer};
 use arrow_data::{ArrayData, ArrayDataBuilder};
 use arrow_schema::{ArrowError, DataType, Field};
 
@@ -675,12 +673,7 @@ where
             *offset = length_so_far;
         }
 
-        nulls = match indices.data_ref().null_buffer() {
-            Some(buffer) => {
-                Some(buffer_bin_and(buffer, 0, &null_buf.into(), 0, data_len))
-            }
-            None => Some(null_buf.into()),
-        };
+        nulls = Some(null_buf.into())
     }
 
     let array_data = ArrayData::builder(GenericStringArray::<OffsetSize>::DATA_TYPE)
@@ -1547,6 +1540,23 @@ mod tests {
         _test_take_string::<LargeStringArray>()
     }
 
+    #[test]
+    fn test_take_slice_string() {
+        let strings =
+            StringArray::from(vec![Some("hello"), None, Some("world"), None, Some("hi")]);
+        let indices = Int32Array::from(vec![Some(0), Some(1), None, Some(0), Some(2)]);
+        let indices_slice = indices.slice(1, 4);
+        let indices_slice = indices_slice
+            .as_ref()
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+
+        let expected = StringArray::from(vec![None, None, Some("hello"), Some("world")]);
+        let result = take(&strings, indices_slice, None).unwrap();
+        assert_eq!(result.as_ref(), &expected);
+    }
+
     macro_rules! test_take_list {
         ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
             // Construct a value array, [[0,0,0], [-1,-2,-1], [2,3]]