wesm commented on a change in pull request #7240: URL: https://github.com/apache/arrow/pull/7240#discussion_r428740035
########## File path: cpp/src/arrow/compute/kernel.cc ########## @@ -0,0 +1,308 @@ +// 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. + +#include "arrow/compute/kernel.h" + +#include <cstddef> +#include <memory> +#include <sstream> +#include <string> + +#include "arrow/buffer.h" +#include "arrow/compute/exec.h" +#include "arrow/result.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/hashing.h" Review comment: This is only one .cc file, I only need `hash_combine` so perhaps we need `arrow/util/hash_util.h` ########## File path: cpp/src/arrow/compute/kernel.h ########## @@ -15,295 +15,517 @@ // specific language governing permissions and limitations // under the License. +// NOTE: API is EXPERIMENTAL and will change without going through a +// deprecation cycle + #pragma once +#include <cstdint> +#include <functional> #include <memory> +#include <string> #include <utility> #include <vector> -#include "arrow/array.h" -#include "arrow/record_batch.h" -#include "arrow/scalar.h" -#include "arrow/table.h" -#include "arrow/util/macros.h" -#include "arrow/util/memory.h" -#include "arrow/util/variant.h" // IWYU pragma: export +#include "arrow/compute/exec.h" +#include "arrow/datum.h" +#include "arrow/result.h" +#include "arrow/status.h" +#include "arrow/type.h" #include "arrow/util/visibility.h" namespace arrow { + +class Buffer; +struct Datum; + namespace compute { -class FunctionContext; +struct FunctionOptions; -/// \class OpKernel -/// \brief Base class for operator kernels -/// -/// Note to implementors: -/// Operator kernels are intended to be the lowest level of an analytics/compute -/// engine. They will generally not be exposed directly to end-users. Instead -/// they will be wrapped by higher level constructs (e.g. top-level functions -/// or physical execution plan nodes). These higher level constructs are -/// responsible for user input validation and returning the appropriate -/// error Status. -/// -/// Due to this design, implementations of Call (the execution -/// method on subclasses) should use assertions (i.e. DCHECK) to double-check -/// parameter arguments when in higher level components returning an -/// InvalidArgument error might be more appropriate. -/// -class ARROW_EXPORT OpKernel { +/// \brief Base class for opaque kernel-specific state. For example, if there +/// is some kind of initialization required +struct KernelState { + virtual ~KernelState() = default; +}; + +/// \brief Context/state for the execution of a particular kernel +class ARROW_EXPORT KernelContext { public: - virtual ~OpKernel() = default; - /// \brief EXPERIMENTAL The output data type of the kernel - /// \return the output type - virtual std::shared_ptr<DataType> out_type() const = 0; + explicit KernelContext(ExecContext* exec_ctx) : exec_ctx_(exec_ctx) {} + + /// \brief Allocate buffer from the context's memory pool + Result<std::shared_ptr<Buffer>> Allocate(int64_t nbytes); + + /// \brief Allocate buffer for bitmap from the context's memory pool + Result<std::shared_ptr<Buffer>> AllocateBitmap(int64_t num_bits); + + /// \brief Indicate that an error has occurred, to be checked by a exec caller + /// \param[in] status a Status instance + /// + /// \note Will not overwrite a prior set Status, so we will have the first + /// error that occurred until ExecContext::ResetStatus is called + void SetStatus(const Status& status); + + /// \brief Clear any error status + void ResetStatus(); + + /// \brief Return true if an error has occurred + bool HasError() const { return !status_.ok(); } + + /// \brief Return the current status of the context + const Status& status() const { return status_; } + + // For passing kernel state to + void SetState(KernelState* state) { state_ = state; } + + KernelState* state() { return state_; } + + /// \brief Common state related to function execution + ExecContext* exec_context() { return exec_ctx_; } + + MemoryPool* memory_pool() { return exec_ctx_->memory_pool(); } + + private: + ExecContext* exec_ctx_; + Status status_; + KernelState* state_; }; -struct Datum; -static inline bool CollectionEquals(const std::vector<Datum>& left, - const std::vector<Datum>& right); - -// Datums variants may have a length. This special value indicate that the -// current variant does not have a length. -constexpr int64_t kUnknownLength = -1; - -/// \class Datum -/// \brief Variant type for various Arrow C++ data structures -struct ARROW_EXPORT Datum { - enum type { NONE, SCALAR, ARRAY, CHUNKED_ARRAY, RECORD_BATCH, TABLE, COLLECTION }; - - util::variant<decltype(NULLPTR), std::shared_ptr<Scalar>, std::shared_ptr<ArrayData>, - std::shared_ptr<ChunkedArray>, std::shared_ptr<RecordBatch>, - std::shared_ptr<Table>, std::vector<Datum>> - value; - - /// \brief Empty datum, to be populated elsewhere - Datum() : value(NULLPTR) {} - - Datum(const std::shared_ptr<Scalar>& value) // NOLINT implicit conversion - : value(value) {} - Datum(const std::shared_ptr<ArrayData>& value) // NOLINT implicit conversion - : value(value) {} - - Datum(const std::shared_ptr<Array>& value) // NOLINT implicit conversion - : Datum(value ? value->data() : NULLPTR) {} - - Datum(const std::shared_ptr<ChunkedArray>& value) // NOLINT implicit conversion - : value(value) {} - Datum(const std::shared_ptr<RecordBatch>& value) // NOLINT implicit conversion - : value(value) {} - Datum(const std::shared_ptr<Table>& value) // NOLINT implicit conversion - : value(value) {} - Datum(const std::vector<Datum>& value) // NOLINT implicit conversion - : value(value) {} - - // Cast from subtypes of Array to Datum - template <typename T, typename = enable_if_t<std::is_base_of<Array, T>::value>> - Datum(const std::shared_ptr<T>& value) // NOLINT implicit conversion - : Datum(std::shared_ptr<Array>(value)) {} - - // Convenience constructors - explicit Datum(bool value) : value(std::make_shared<BooleanScalar>(value)) {} - explicit Datum(int8_t value) : value(std::make_shared<Int8Scalar>(value)) {} - explicit Datum(uint8_t value) : value(std::make_shared<UInt8Scalar>(value)) {} - explicit Datum(int16_t value) : value(std::make_shared<Int16Scalar>(value)) {} - explicit Datum(uint16_t value) : value(std::make_shared<UInt16Scalar>(value)) {} - explicit Datum(int32_t value) : value(std::make_shared<Int32Scalar>(value)) {} - explicit Datum(uint32_t value) : value(std::make_shared<UInt32Scalar>(value)) {} - explicit Datum(int64_t value) : value(std::make_shared<Int64Scalar>(value)) {} - explicit Datum(uint64_t value) : value(std::make_shared<UInt64Scalar>(value)) {} - explicit Datum(float value) : value(std::make_shared<FloatScalar>(value)) {} - explicit Datum(double value) : value(std::make_shared<DoubleScalar>(value)) {} - - ~Datum() {} - - Datum(const Datum& other) noexcept { this->value = other.value; } - - Datum& operator=(const Datum& other) noexcept { - value = other.value; - return *this; - } +#define ARROW_CTX_RETURN_IF_ERROR(CTX) \ + do { \ + if (ARROW_PREDICT_FALSE((CTX)->HasError())) { \ + Status s = (CTX)->status(); \ + (CTX)->ResetStatus(); \ + return s; \ + } \ + } while (0) + +/// A standard function taking zero or more Array/Scalar values and returning +/// Array/Scalar output. May be used for SCALAR and VECTOR kernel kinds. Should +/// write into pre-allocated memory except in cases when a builder +/// (e.g. StringBuilder) must be employed +using ArrayKernelExec = std::function<void(KernelContext*, const ExecBatch&, Datum*)>; + +/// \brief A container to express what kernel argument input types are accepted +class ARROW_EXPORT InputType { + public: + enum Kind { + /// Accept any value type + ANY_TYPE, - // Define move constructor and move assignment, for better performance - Datum(Datum&& other) noexcept : value(std::move(other.value)) {} + /// A fixed arrow::DataType and will only exact match having this exact + /// type (e.g. same TimestampType unit, same decimal scale and precision, + /// or same nested child types + EXACT_TYPE, - Datum& operator=(Datum&& other) noexcept { - value = std::move(other.value); - return *this; - } + /// Any type having the indicated Type::type id. For example, accept + /// any Type::LIST or any Type::TIMESTAMP + SAME_TYPE_ID, + }; - Datum::type kind() const { - switch (this->value.index()) { - case 0: - return Datum::NONE; - case 1: - return Datum::SCALAR; - case 2: - return Datum::ARRAY; - case 3: - return Datum::CHUNKED_ARRAY; - case 4: - return Datum::RECORD_BATCH; - case 5: - return Datum::TABLE; - case 6: - return Datum::COLLECTION; - default: - return Datum::NONE; - } - } + InputType(ValueDescr::Shape shape = ValueDescr::ANY) // NOLINT implicit construction + : kind_(ANY_TYPE), shape_(shape) {} - std::shared_ptr<ArrayData> array() const { - return util::get<std::shared_ptr<ArrayData>>(this->value); - } + InputType(std::shared_ptr<DataType> type, + ValueDescr::Shape shape = ValueDescr::ANY) // NOLINT implicit construction + : kind_(EXACT_TYPE), shape_(shape), type_(std::move(type)), type_id_(type_->id()) {} - std::shared_ptr<Array> make_array() const { - return MakeArray(util::get<std::shared_ptr<ArrayData>>(this->value)); - } + InputType(const ValueDescr& descr) // NOLINT implicit construction + : InputType(descr.type, descr.shape) {} - std::shared_ptr<ChunkedArray> chunked_array() const { - return util::get<std::shared_ptr<ChunkedArray>>(this->value); - } + InputType(Type::type type_id, + ValueDescr::Shape shape = ValueDescr::ANY) // NOLINT implicit construction + : kind_(SAME_TYPE_ID), shape_(shape), type_id_(type_id) {} - std::shared_ptr<RecordBatch> record_batch() const { - return util::get<std::shared_ptr<RecordBatch>>(this->value); - } + InputType(const InputType& other) { CopyInto(other); } - std::shared_ptr<Table> table() const { - return util::get<std::shared_ptr<Table>>(this->value); + // Convenience ctors + static InputType Array(std::shared_ptr<DataType> type) { + return InputType(std::move(type), ValueDescr::ARRAY); } - const std::vector<Datum> collection() const { - return util::get<std::vector<Datum>>(this->value); + static InputType Scalar(std::shared_ptr<DataType> type) { + return InputType(std::move(type), ValueDescr::SCALAR); } - std::shared_ptr<Scalar> scalar() const { - return util::get<std::shared_ptr<Scalar>>(this->value); - } + static InputType Array(Type::type id) { return InputType(id, ValueDescr::ARRAY); } - bool is_array() const { return this->kind() == Datum::ARRAY; } + static InputType Scalar(Type::type id) { return InputType(id, ValueDescr::SCALAR); } - bool is_arraylike() const { - return this->kind() == Datum::ARRAY || this->kind() == Datum::CHUNKED_ARRAY; - } + void operator=(const InputType& other) { CopyInto(other); } - bool is_scalar() const { return this->kind() == Datum::SCALAR; } + InputType(InputType&& other) { MoveInto(std::forward<InputType>(other)); } - bool is_collection() const { return this->kind() == Datum::COLLECTION; } + void operator=(InputType&& other) { MoveInto(std::forward<InputType>(other)); } - /// \brief The value type of the variant, if any - /// - /// \return nullptr if no type - std::shared_ptr<DataType> type() const { - if (this->kind() == Datum::ARRAY) { - return util::get<std::shared_ptr<ArrayData>>(this->value)->type; - } else if (this->kind() == Datum::CHUNKED_ARRAY) { - return util::get<std::shared_ptr<ChunkedArray>>(this->value)->type(); - } else if (this->kind() == Datum::SCALAR) { - return util::get<std::shared_ptr<Scalar>>(this->value)->type; - } - return NULLPTR; + /// \brief Return true if this type exactly matches another + bool Equals(const InputType& other) const; + + bool operator==(const InputType& other) const { return this->Equals(other); } + + bool operator!=(const InputType& other) const { return !(*this == other); } + + /// \brief Return hash code + uint64_t Hash() const; + + /// \brief Render a human-readable string representation + std::string ToString() const; + + /// \brief Return true if the value matches this argument kind in type + /// and shape + bool Matches(const Datum& value) const; + + /// \brief Return true if the value descriptor matches this argument kind in + /// type and shape + bool Matches(const ValueDescr& value) const; + + /// \brief The type matching rule that this InputType uses + Kind kind() const { return kind_; } + + ValueDescr::Shape shape() const { return shape_; } + + /// \brief For ArgKind::EXACT_TYPE, the exact type that this InputType must + /// match. Otherwise this function should not be used + const std::shared_ptr<DataType>& type() const; + + /// \brief For ArgKind::SAME_TYPE_ID, the Type::type that this InputType must + /// match, Otherwise this function should not be used + Type::type type_id() const; + + private: + void CopyInto(const InputType& other) { + this->kind_ = other.kind_; + this->shape_ = other.shape_; + this->type_ = other.type_; + this->type_id_ = other.type_id_; } - /// \brief The value length of the variant, if any - /// - /// \return kUnknownLength if no type - int64_t length() const { - if (this->kind() == Datum::ARRAY) { - return util::get<std::shared_ptr<ArrayData>>(this->value)->length; - } else if (this->kind() == Datum::CHUNKED_ARRAY) { - return util::get<std::shared_ptr<ChunkedArray>>(this->value)->length(); - } else if (this->kind() == Datum::SCALAR) { - return 1; - } - return kUnknownLength; + void MoveInto(InputType&& other) { + this->kind_ = other.kind_; + this->shape_ = other.shape_; + this->type_ = std::move(other.type_); + this->type_id_ = other.type_id_; } - /// \brief The array chunks of the variant, if any - /// - /// \return empty if not arraylike - ArrayVector chunks() const { - if (!this->is_arraylike()) { - return {}; - } - if (this->is_array()) { - return {this->make_array()}; - } - return this->chunked_array()->chunks(); + Kind kind_; + + ValueDescr::Shape shape_ = ValueDescr::ANY; + + // For EXACT_TYPE ArgKind + std::shared_ptr<DataType> type_; + + // For SAME_TYPE_ID ArgKind + Type::type type_id_ = Type::NA; +}; + +/// \brief Container to capture both exact and input-dependent output types +/// +/// The value shape returned by Resolve will be determined by broadcasting the +/// shapes of the input arguments, otherwise this is handled by the +/// user-defined resolver function +/// +/// * Any ARRAY shape -> output shape is ARRAY +/// * All SCALAR shapes -> output shape is SCALAR +class ARROW_EXPORT OutputType { + public: + /// \brief An enum indicating whether the value type is an invariant fixed + /// value or one that's computed by a kernel-defined resolver function + enum ResolveKind { FIXED, COMPUTED }; + + /// Type resolution function. Given input types and shapes, return output + /// type and shape. This function SHOULD _not_ be used to check for arity, + /// that SHOULD be performed one or more layers above. May make use of kernel + /// state to know what type to output + using Resolver = + std::function<Result<ValueDescr>(KernelContext*, const std::vector<ValueDescr>&)>; + + OutputType(std::shared_ptr<DataType> type) // NOLINT implicit construction + : kind_(FIXED), type_(std::move(type)) {} + + /// For outputting a particular type and shape + OutputType(ValueDescr descr); // NOLINT implicit construction + + explicit OutputType(Resolver resolver) : kind_(COMPUTED), resolver_(resolver) {} + + OutputType(const OutputType& other) { + this->kind_ = other.kind_; + this->shape_ = other.shape_; + this->type_ = other.type_; + this->resolver_ = other.resolver_; } - bool Equals(const Datum& other) const { - if (this->kind() != other.kind()) return false; - - switch (this->kind()) { - case Datum::NONE: - return true; - case Datum::SCALAR: - return internal::SharedPtrEquals(this->scalar(), other.scalar()); - case Datum::ARRAY: - return internal::SharedPtrEquals(this->make_array(), other.make_array()); - case Datum::CHUNKED_ARRAY: - return internal::SharedPtrEquals(this->chunked_array(), other.chunked_array()); - case Datum::RECORD_BATCH: - return internal::SharedPtrEquals(this->record_batch(), other.record_batch()); - case Datum::TABLE: - return internal::SharedPtrEquals(this->table(), other.table()); - case Datum::COLLECTION: - return CollectionEquals(this->collection(), other.collection()); - default: - return false; - } + OutputType(OutputType&& other) { + this->kind_ = other.kind_; + this->type_ = std::move(other.type_); + this->shape_ = other.shape_; + this->resolver_ = other.resolver_; } + + /// \brief Return the shape and type of the expected output value of the + /// kernel given the value descriptors (shapes and types). The resolver may + /// make use of state information kept in the KernelContext + Result<ValueDescr> Resolve(KernelContext* ctx, + const std::vector<ValueDescr>& args) const; + + /// \brief The value type for the FIXED kind rule + const std::shared_ptr<DataType>& type() const; + + /// \brief For use with COMPUTED resolution strategy, the output type depends + /// on the input type. It may be more convenient to invoke this with + /// OutputType::Resolve returned from this method + const Resolver& resolver() const; + + /// \brief Render a human-readable string representation + std::string ToString() const; + + /// \brief Return the kind of type resolution of this output type, whether + /// fixed/invariant or computed by a "user"-defined resolver + ResolveKind kind() const { return kind_; } + + /// \brief If the shape is ANY, then Resolve will compute the shape based on + /// the input arguments + ValueDescr::Shape shape() const { return shape_; } + + private: + ResolveKind kind_; + + // For FIXED resolution + std::shared_ptr<DataType> type_; + + ValueDescr::Shape shape_ = ValueDescr::ANY; + + // For COMPUTED resolution + Resolver resolver_; }; -/// \class UnaryKernel -/// \brief An array-valued function of a single input argument. +/// \brief Holds the input types and output type of the kernel /// -/// Note to implementors: Try to avoid making kernels that allocate memory if -/// the output size is a deterministic function of the Input Datum's metadata. -/// Instead separate the logic of the kernel and allocations necessary into -/// two different kernels. Some reusable kernels that allocate buffers -/// and delegate computation to another kernel are available in util-internal.h. -class ARROW_EXPORT UnaryKernel : public OpKernel { +/// Varargs functions should pass a single input type to be used to validate +/// the the input types of a function invocation +class ARROW_EXPORT KernelSignature { public: - /// \brief Executes the kernel. - /// - /// \param[in] ctx The function context for the kernel - /// \param[in] input The kernel input data - /// \param[out] out The output of the function. Each implementation of this - /// function might assume different things about the existing contents of out - /// (e.g. which buffers are preallocated). In the future it is expected that - /// there will be a more generic mechanism for understanding the necessary - /// contracts. - virtual Status Call(FunctionContext* ctx, const Datum& input, Datum* out) = 0; + KernelSignature(std::vector<InputType> in_types, OutputType out_type, + bool is_varargs = false); + + /// \brief Convenience ctor since make_shared can be awkward + static std::shared_ptr<KernelSignature> Make(std::vector<InputType> in_types, + OutputType out_type, + bool is_varargs = false); + + /// \brief Return true if the signature if compatible with the list of input + /// value descriptors + bool MatchesInputs(const std::vector<ValueDescr>& descriptors) const; + + /// \brief Returns true if the input types of each signature are + /// equal. Well-formed functions should have a deterministic output type + /// given input types, but currently it is the responsibility of the + /// developer to ensure this + bool Equals(const KernelSignature& other) const; + + bool operator==(const KernelSignature& other) const { return this->Equals(other); } + + bool operator!=(const KernelSignature& other) const { return !(*this == other); } + + /// \brief Compute a hash code for the signature + int64_t Hash() const; + + const std::vector<InputType>& in_types() const { return in_types_; } + + const OutputType& out_type() const { return out_type_; } + + /// \brief Render a human-readable string representation + std::string ToString() const; + + bool is_varargs() const { return is_varargs_; } + + private: + std::vector<InputType> in_types_; + OutputType out_type_; + bool is_varargs_; + + // For caching the hash code after it's computed the first time + mutable int64_t hash_code_; }; -/// \class BinaryKernel -/// \brief An array-valued function of a two input arguments -class ARROW_EXPORT BinaryKernel : public OpKernel { - public: - virtual Status Call(FunctionContext* ctx, const Datum& left, const Datum& right, - Datum* out) = 0; +struct SimdLevel { + enum type { NONE, SSE4_2, AVX, AVX2, AVX512, NEON }; }; -// TODO doxygen 1.8.16 does not like the following code -///@cond INTERNAL +struct NullHandling { + enum type { + /// Compute the output validity bitmap by intersecting the validity bitmaps + /// of the arguments. Kernel does not do anything with the bitmap + INTERSECTION, -static inline bool CollectionEquals(const std::vector<Datum>& left, - const std::vector<Datum>& right) { - if (left.size() != right.size()) { - return false; - } + /// Kernel expects a pre-allocated buffer to write the result bitmap into + COMPUTED_PREALLOCATE, - for (size_t i = 0; i < left.size(); i++) { - if (!left[i].Equals(right[i])) { - return false; - } - } - return true; -} + /// Kernel allocates and populates the validity bitmap of the output + COMPUTED_NO_PREALLOCATE, + + /// Output is never null + OUTPUT_NOT_NULL + }; +}; + +struct MemAllocation { + enum type { + // For data types that support pre-allocation (fixed-type), the kernel + // expects to be provided pre-allocated memory to write + // into. Non-fixed-width must always allocate their own memory but perhaps + // not their validity bitmaps. The allocation made for the same length as + // the execution batch, so vector kernels yielding differently sized output + // should not use this + PREALLOCATE, + + // The kernel does its own memory allocation + NO_PREALLOCATE + }; +}; + +struct Kernel; + +struct KernelInitArgs { + const Kernel* kernel; + const std::vector<ValueDescr>& inputs; + const FunctionOptions* options; +}; + +// Kernel initializer (context, argument descriptors, options) +using KernelInit = + std::function<std::unique_ptr<KernelState>(KernelContext*, const KernelInitArgs&)>; + +/// \brief Base type for kernels. Contains the function signature and +/// optionally the state initialization function, along with some common +/// attributes +struct Kernel { + Kernel() {} + + Kernel(std::shared_ptr<KernelSignature> sig, KernelInit init) + : signature(std::move(sig)), init(init) {} + + Kernel(std::vector<InputType> in_types, OutputType out_type, KernelInit init) + : Kernel(KernelSignature::Make(std::move(in_types), out_type), init) {} + + std::shared_ptr<KernelSignature> signature; + + /// \brief Create a new KernelState for invocations of this kernel, e.g. to + /// set up any options or state relevant for execution. May be nullptr + KernelInit init; -///@endcond + // Does execution benefit from parallelization (splitting large chunks into + // smaller chunks and using multiple threads). Some vector kernels may + // require single-threaded execution. + bool parallelizable = true; + + SimdLevel::type simd_level = SimdLevel::NONE; Review comment: It will be ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org
