Faster take with null values.

src/compute/take/primitive.rs

@@ -1,19 +1,3 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements.  See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership.  The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License.  You may obtain a copy of the License at
-//
-//   http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied.  See the License for the
-// specific language governing permissions and limitations
-// under the License.
 use crate::{
     array::{Array, PrimitiveArray},
     bitmap::{Bitmap, MutableBitmap},
@@ -39,24 +23,22 @@ fn take_values_validity<T: NativeType, I: Index>(
     values: &PrimitiveArray<T>,
     indices: &[I],
 ) -> (Buffer<T>, Option<Bitmap>) {
-    let mut null = MutableBitmap::with_capacity(indices.len());
+    let values_validity = values.validity().unwrap();
 
-    let null_values = values.validity().unwrap();
+    let validity = indices.iter().map(|index| {
+        let index = index.to_usize();
+        values_validity.get_bit(index)
+    });
+    let mut validity = MutableBitmap::from_trusted_len_iter(validity);
 
     let values_values = values.values();
-
     let values = indices.iter().map(|index| {
         let index = index.to_usize();
-        if null_values.get_bit(index) {
-            null.push(true);
-        } else {
-            null.push(false);
-        }
         values_values[index]
     });
     let buffer = MutableBuffer::from_trusted_len_iter(values);
 
-    (buffer.into(), null.into())
+    (buffer.into(), validity.into())
 }
 
 // take implementation when only indices contain nulls