1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
// 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. //! Defines kernels suitable to perform operations to primitive arrays. use crate::array::{Array, ArrayData, PrimitiveArray}; use crate::buffer::Buffer; use crate::datatypes::ArrowPrimitiveType; #[inline] fn into_primitive_array_data<I: ArrowPrimitiveType, O: ArrowPrimitiveType>( array: &PrimitiveArray<I>, buffer: Buffer, ) -> ArrayData { ArrayData::new( O::DATA_TYPE, array.len(), None, array.data_ref().null_buffer().cloned(), 0, vec![buffer], vec![], ) } /// Applies an unary and infalible function to a primitive array. /// This is the fastest way to perform an operation on a primitive array when /// the benefits of a vectorized operation outweights the cost of branching nulls and non-nulls. /// # Implementation /// This will apply the function for all values, including those on null slots. /// This implies that the operation must be infalible for any value of the corresponding type /// or this function may panic. /// # Example /// ```rust /// # use arrow::array::Int32Array; /// # use arrow::datatypes::Int32Type; /// # use arrow::compute::kernels::arity::unary; /// # fn main() { /// let array = Int32Array::from(vec![Some(5), Some(7), None]); /// let c = unary::<_, _, Int32Type>(&array, |x| x * 2 + 1); /// assert_eq!(c, Int32Array::from(vec![Some(11), Some(15), None])); /// # } /// ``` pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> PrimitiveArray<O> where I: ArrowPrimitiveType, O: ArrowPrimitiveType, F: Fn(I::Native) -> O::Native, { let values = array.values().iter().map(|v| op(*v)); // JUSTIFICATION // Benefit // ~60% speedup // Soundness // `values` is an iterator with a known size because arrays are sized. let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; let data = into_primitive_array_data::<_, O>(array, buffer); PrimitiveArray::<O>::from(data) }