diff --git a/src/compute/take/primitive.rs b/src/compute/take/primitive.rs index dd015dfdab5..eade1fcd133 100644 --- a/src/compute/take/primitive.rs +++ b/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( values: &PrimitiveArray, indices: &[I], ) -> (Buffer, Option) { - 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