| // 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. |
| |
| #pragma once |
| |
| #include <cstdint> |
| #include <utility> |
| |
| #include "arrow/array/util.h" |
| #include "arrow/buffer.h" |
| #include "arrow/compute/kernels/codegen_internal.h" |
| #include "arrow/compute/type_fwd.h" |
| #include "arrow/util/bit_run_reader.h" |
| |
| namespace arrow { |
| namespace compute { |
| namespace internal { |
| |
| // An internal data structure for unpacking a primitive argument to pass to a |
| // kernel implementation |
| struct PrimitiveArg { |
| const uint8_t* is_valid; |
| // If the bit_width is a multiple of 8 (i.e. not boolean), then "data" should |
| // be shifted by offset * (bit_width / 8). For bit-packed data, the offset |
| // must be used when indexing. |
| const uint8_t* data; |
| int bit_width; |
| int64_t length; |
| int64_t offset; |
| // This may be kUnknownNullCount if the null_count has not yet been computed, |
| // so use null_count != 0 to determine "may have nulls". |
| int64_t null_count; |
| }; |
| |
| // Get validity bitmap data or return nullptr if there is no validity buffer |
| const uint8_t* GetValidityBitmap(const ArrayData& data); |
| |
| int GetBitWidth(const DataType& type); |
| |
| // Reduce code size by dealing with the unboxing of the kernel inputs once |
| // rather than duplicating compiled code to do all these in each kernel. |
| PrimitiveArg GetPrimitiveArg(const ArrayData& arr); |
| |
| // Augment a unary ArrayKernelExec which supports only array-like inputs with support for |
| // scalar inputs. Scalars will be transformed to 1-long arrays with the scalar's value (or |
| // null if the scalar is null) as its only element. This 1-long array will be passed to |
| // the original exec, then the only element of the resulting array will be extracted as |
| // the output scalar. This could be far more efficient, but instead of optimizing this |
| // it'd be better to support scalar inputs "upstream" in original exec. |
| ArrayKernelExec TrivialScalarUnaryAsArraysExec( |
| ArrayKernelExec exec, NullHandling::type null_handling = NullHandling::INTERSECTION); |
| |
| // Return (min, max) of a numerical array, ignore nulls. |
| // For empty array, return the maximal number limit as 'min', and minimal limit as 'max'. |
| template <typename T> |
| ARROW_NOINLINE std::pair<T, T> GetMinMax(const ArrayData& data) { |
| T min = std::numeric_limits<T>::max(); |
| T max = std::numeric_limits<T>::lowest(); |
| |
| const T* values = data.GetValues<T>(1); |
| arrow::internal::VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length, |
| [&](int64_t pos, int64_t len) { |
| for (int64_t i = 0; i < len; ++i) { |
| min = std::min(min, values[pos + i]); |
| max = std::max(max, values[pos + i]); |
| } |
| }); |
| |
| return std::make_pair(min, max); |
| } |
| |
| template <typename T> |
| std::pair<T, T> GetMinMax(const Datum& datum) { |
| T min = std::numeric_limits<T>::max(); |
| T max = std::numeric_limits<T>::lowest(); |
| |
| for (const auto& array : datum.chunks()) { |
| T local_min, local_max; |
| std::tie(local_min, local_max) = GetMinMax<T>(*array->data()); |
| min = std::min(min, local_min); |
| max = std::max(max, local_max); |
| } |
| |
| return std::make_pair(min, max); |
| } |
| |
| // Count value occurrences of an array, ignore nulls. |
| // 'counts' must be zeroed and with enough size. |
| template <typename T> |
| ARROW_NOINLINE int64_t CountValues(uint64_t* counts, const ArrayData& data, T min) { |
| const int64_t n = data.length - data.GetNullCount(); |
| if (n > 0) { |
| const T* values = data.GetValues<T>(1); |
| arrow::internal::VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length, |
| [&](int64_t pos, int64_t len) { |
| for (int64_t i = 0; i < len; ++i) { |
| ++counts[values[pos + i] - min]; |
| } |
| }); |
| } |
| return n; |
| } |
| |
| template <typename T> |
| int64_t CountValues(uint64_t* counts, const Datum& datum, T min) { |
| int64_t n = 0; |
| for (const auto& array : datum.chunks()) { |
| n += CountValues<T>(counts, *array->data(), min); |
| } |
| return n; |
| } |
| |
| // Copy numerical array values to a buffer, ignore nulls. |
| template <size_t SizeOfCType> |
| ARROW_NOINLINE int64_t CopyNonNullValues(const ArrayData& data, void* out) { |
| uint8_t* u8_buffer = reinterpret_cast<uint8_t*>(out); |
| const int64_t n = data.length - data.GetNullCount(); |
| if (n > 0) { |
| int64_t index = 0; |
| const uint8_t* u8_values = data.GetValues<uint8_t>(1); |
| arrow::internal::VisitSetBitRunsVoid( |
| data.buffers[0], data.offset, data.length, [&](int64_t pos, int64_t len) { |
| memcpy(u8_buffer + index * SizeOfCType, u8_values + pos * SizeOfCType, |
| len * SizeOfCType); |
| index += len; |
| }); |
| } |
| return n; |
| } |
| |
| template <size_t SizeOfCType> |
| int64_t CopyNonNullValues(const Datum& datum, void* out) { |
| uint8_t* u8_buffer = reinterpret_cast<uint8_t*>(out); |
| int64_t n = 0; |
| for (const auto& array : datum.chunks()) { |
| n += CopyNonNullValues<SizeOfCType>(*array->data(), u8_buffer + n * SizeOfCType); |
| } |
| return n; |
| } |
| |
| } // namespace internal |
| } // namespace compute |
| } // namespace arrow |