Dandandan commented on a change in pull request #8799:
URL: https://github.com/apache/arrow/pull/8799#discussion_r532203770



##########
File path: rust/arrow/src/compute/kernels/take.rs
##########
@@ -365,43 +365,81 @@ where
         .unwrap();
 
     let num_bytes = bit_util::ceil(data_len, 8);
-    let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, 
true);
-    let null_slice = null_buf.data_mut();
 
+    let null_count = array.null_count();
     let mut offsets = Vec::with_capacity(data_len + 1);
     let mut values = Vec::with_capacity(data_len);
     let mut length_so_far = OffsetSize::zero();
-
     offsets.push(length_so_far);
-    for i in 0..data_len {
-        let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
-            ArrowError::ComputeError("Cast to usize failed".to_string())
-        })?;
 
-        if array.is_valid(index) && indices.is_valid(i) {
+    let nulls;
+    if null_count == 0 && indices.null_count() == 0 {
+        for i in 0..data_len {
+            let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| 
{
+                ArrowError::ComputeError("Cast to usize failed".to_string())
+            })?;
+
             let s = array.value(index);
 
             length_so_far += OffsetSize::from_usize(s.len()).unwrap();
             values.extend_from_slice(s.as_bytes());
-        } else {
-            // set null bit
-            bit_util::unset_bit(null_slice, i);
+            offsets.push(length_so_far);
         }
-        offsets.push(length_so_far);
-    }
+        nulls = None
+    } else if null_count == 0 {
+        for i in 0..data_len {
+            let index = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| 
{

Review comment:
       Well spotted!




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to