Migrated other kernels.

jorgecarleitao · Nov 25, 2021 · 865ad04 · 865ad04
1 parent a0c16e3
commit 865ad04
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 64 deletions.
diff --git a/src/compute/take/boolean.rs b/src/compute/take/boolean.rs
@@ -1,20 +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, BooleanArray, PrimitiveArray},
     bitmap::{Bitmap, MutableBitmap},
@@ -35,21 +18,16 @@ fn take_values_validity<I: Index>(
     values: &BooleanArray,
     indices: &[I],
 ) -> (Bitmap, Option<Bitmap>) {
-    let mut validity = MutableBitmap::with_capacity(indices.len());
-
     let validity_values = values.validity().unwrap();
+    let validity = indices
+        .iter()
+        .map(|index| validity_values.get_bit(index.to_usize()));
+    let validity = Bitmap::from_trusted_len_iter(validity);
 
     let values_values = values.values();
-
-    let values = indices.iter().map(|index| {
-        let index = index.to_usize();
-        if validity_values.get_bit(index) {
-            validity.push(true);
-        } else {
-            validity.push(false);
-        }
-        values_values.get_bit(index)
-    });
+    let values = indices
+        .iter()
+        .map(|index| values_values.get_bit(index.to_usize()));
     let buffer = Bitmap::from_trusted_len_iter(values);
 
     (buffer, validity.into())

diff --git a/src/compute/take/generic_binary.rs b/src/compute/take/generic_binary.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::{GenericBinaryArray, Offset, PrimitiveArray},
     bitmap::{Bitmap, MutableBitmap},
@@ -66,25 +50,23 @@ pub fn take_values_validity<O: Offset, I: Index, A: GenericBinaryArray<O>>(
     values: &A,
     indices: &[I],
 ) -> (Buffer<O>, Buffer<u8>, Option<Bitmap>) {
+    let validity_values = values.validity().unwrap();
+    let validity = indices
+        .iter()
+        .map(|index| validity_values.get_bit(index.to_usize()));
+    let validity = Bitmap::from_trusted_len_iter(validity);
+
     let mut length = O::default();
-    let mut validity = MutableBitmap::with_capacity(indices.len());
 
-    let null_values = values.validity().unwrap();
     let offsets = values.offsets();
     let values_values = values.values();
 
     let mut starts = MutableBuffer::<O>::with_capacity(indices.len());
     let offsets = indices.iter().map(|index| {
         let index = index.to_usize();
-        if null_values.get_bit(index) {
-            validity.push(true);
-            let start = offsets[index];
-            length += offsets[index + 1] - start;
-            starts.push(start);
-        } else {
-            validity.push(false);
-            starts.push(O::default());
-        }
+        let start = offsets[index];
+        length += offsets[index + 1] - start;
+        starts.push(start);
         length
     });
     let offsets = std::iter::once(O::default()).chain(offsets);

diff --git a/src/compute/take/primitive.rs b/src/compute/take/primitive.rs
@@ -25,17 +25,13 @@ fn take_values_validity<T: NativeType, I: Index>(
 ) -> (Buffer<T>, Option<Bitmap>) {
     let values_validity = values.validity().unwrap();
 
-    let validity = indices.iter().map(|index| {
-        let index = index.to_usize();
-        values_validity.get_bit(index)
-    });
+    let validity = indices
+        .iter()
+        .map(|index| values_validity.get_bit(index.to_usize()));
     let validity = MutableBitmap::from_trusted_len_iter(validity);
 
     let values_values = values.values();
-    let values = indices.iter().map(|index| {
-        let index = index.to_usize();
-        values_values[index]
-    });
+    let values = indices.iter().map(|index| values_values[index.to_usize()]);
     let buffer = MutableBuffer::from_trusted_len_iter(values);
 
     (buffer.into(), validity.into())