1 // Licensed to the Apache Software Foundation (ASF) under one 2 // or more contributor license agreements. See the NOTICE file 3 // distributed with this work for additional information 4 // regarding copyright ownership. The ASF licenses this file 5 // to you under the Apache License, Version 2.0 (the 6 // "License"); you may not use this file except in compliance 7 // with the License. You may obtain a copy of the License at 8 // 9 // http://www.apache.org/licenses/LICENSE-2.0 10 // 11 // Unless required by applicable law or agreed to in writing, 12 // software distributed under the License is distributed on an 13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14 // KIND, either express or implied. See the License for the 15 // specific language governing permissions and limitations 16 // under the License. 17 18 // NOTE: API is EXPERIMENTAL and will change without going through a 19 // deprecation cycle 20 21 #pragma once 22 23 #include <cstdint> 24 #include <functional> 25 #include <memory> 26 #include <string> 27 #include <utility> 28 #include <vector> 29 30 #include "arrow/compute/exec.h" 31 #include "arrow/datum.h" 32 #include "arrow/result.h" 33 #include "arrow/status.h" 34 #include "arrow/type.h" 35 #include "arrow/util/macros.h" 36 #include "arrow/util/visibility.h" 37 38 namespace arrow { 39 40 class Buffer; 41 class MemoryPool; 42 43 namespace compute { 44 45 struct FunctionOptions; 46 47 /// \brief Base class for opaque kernel-specific state. For example, if there 48 /// is some kind of initialization required. 49 struct ARROW_EXPORT KernelState { 50 virtual ~KernelState() = default; 51 }; 52 53 /// \brief Context/state for the execution of a particular kernel. 54 class ARROW_EXPORT KernelContext { 55 public: KernelContext(ExecContext * exec_ctx)56 explicit KernelContext(ExecContext* exec_ctx) : exec_ctx_(exec_ctx) {} 57 58 /// \brief Allocate buffer from the context's memory pool. The contents are 59 /// not uninitialized. 60 Result<std::shared_ptr<Buffer>> Allocate(int64_t nbytes); 61 62 /// \brief Allocate buffer for bitmap from the context's memory pool. Like 63 /// Allocate, the contents of the buffer are not initialized but the last 64 /// byte is preemptively zeroed to help avoid ASAN or valgrind issues. 65 Result<std::shared_ptr<Buffer>> AllocateBitmap(int64_t num_bits); 66 67 /// \brief Indicate that an error has occurred, to be checked by a exec caller 68 /// \param[in] status a Status instance. 69 /// 70 /// \note Will not overwrite a prior set Status, so we will have the first 71 /// error that occurred until ExecContext::ResetStatus is called. 72 void SetStatus(const Status& status); 73 74 /// \brief Clear any error status. 75 void ResetStatus(); 76 77 /// \brief Return true if an error has occurred. HasError()78 bool HasError() const { return !status_.ok(); } 79 80 /// \brief Return the current status of the context. status()81 const Status& status() const { return status_; } 82 83 /// \brief Assign the active KernelState to be utilized for each stage of 84 /// kernel execution. Ownership and memory lifetime of the KernelState must 85 /// be minded separately. SetState(KernelState * state)86 void SetState(KernelState* state) { state_ = state; } 87 state()88 KernelState* state() { return state_; } 89 90 /// \brief Configuration related to function execution that is to be shared 91 /// across multiple kernels. exec_context()92 ExecContext* exec_context() { return exec_ctx_; } 93 94 /// \brief The memory pool to use for allocations. For now, it uses the 95 /// MemoryPool contained in the ExecContext used to create the KernelContext. memory_pool()96 MemoryPool* memory_pool() { return exec_ctx_->memory_pool(); } 97 98 private: 99 ExecContext* exec_ctx_; 100 Status status_; 101 KernelState* state_; 102 }; 103 104 // A macro to invoke for error control flow after invoking functions (such as 105 // kernel init or exec functions) that propagate errors via KernelContext. 106 #define ARROW_CTX_RETURN_IF_ERROR(CTX) \ 107 do { \ 108 if (ARROW_PREDICT_FALSE((CTX)->HasError())) { \ 109 Status s = (CTX)->status(); \ 110 (CTX)->ResetStatus(); \ 111 return s; \ 112 } \ 113 } while (0) 114 115 /// \brief The standard kernel execution API that must be implemented for 116 /// SCALAR and VECTOR kernel types. This includes both stateless and stateful 117 /// kernels. Kernels depending on some execution state access that state via 118 /// subclasses of KernelState set on the KernelContext object. May be used for 119 /// SCALAR and VECTOR kernel kinds. Implementations should endeavor to write 120 /// into pre-allocated memory if they are able, though for some kernels 121 /// (e.g. in cases when a builder like StringBuilder) must be employed this may 122 /// not be possible. 123 using ArrayKernelExec = std::function<void(KernelContext*, const ExecBatch&, Datum*)>; 124 125 /// \brief An type-checking interface to permit customizable validation rules 126 /// for use with InputType and KernelSignature. This is for scenarios where the 127 /// acceptance is not an exact type instance, such as a TIMESTAMP type for a 128 /// specific TimeUnit, but permitting any time zone. 129 struct ARROW_EXPORT TypeMatcher { 130 virtual ~TypeMatcher() = default; 131 132 /// \brief Return true if this matcher accepts the data type. 133 virtual bool Matches(const DataType& type) const = 0; 134 135 /// \brief A human-interpretable string representation of what the type 136 /// matcher checks for, usable when printing KernelSignature or formatting 137 /// error messages. 138 virtual std::string ToString() const = 0; 139 140 /// \brief Return true if this TypeMatcher contains the same matching rule as 141 /// the other. Currently depends on RTTI. 142 virtual bool Equals(const TypeMatcher& other) const = 0; 143 }; 144 145 namespace match { 146 147 /// \brief Match any DataType instance having the same DataType::id. 148 ARROW_EXPORT std::shared_ptr<TypeMatcher> SameTypeId(Type::type type_id); 149 150 /// \brief Match any TimestampType instance having the same unit, but the time 151 /// zones can be different. 152 ARROW_EXPORT std::shared_ptr<TypeMatcher> TimestampUnit(TimeUnit::type unit); 153 154 } // namespace match 155 156 /// \brief An object used for type- and shape-checking arguments to be passed 157 /// to a kernel and stored in a KernelSignature. Distinguishes between ARRAY 158 /// and SCALAR arguments using ValueDescr::Shape. The type-checking rule can be 159 /// supplied either with an exact DataType instance or a custom TypeMatcher. 160 class ARROW_EXPORT InputType { 161 public: 162 /// \brief The kind of type-checking rule that the InputType contains. 163 enum Kind { 164 /// \brief Accept any value type. 165 ANY_TYPE, 166 167 /// \brief A fixed arrow::DataType and will only exact match having this 168 /// exact type (e.g. same TimestampType unit, same decimal scale and 169 /// precision, or same nested child types). 170 EXACT_TYPE, 171 172 /// \brief Uses a TypeMatcher implementation to check the type. 173 USE_TYPE_MATCHER 174 }; 175 176 /// \brief Accept any value type but with a specific shape (e.g. any Array or 177 /// any Scalar). 178 InputType(ValueDescr::Shape shape = ValueDescr::ANY) // NOLINT implicit construction kind_(ANY_TYPE)179 : kind_(ANY_TYPE), shape_(shape) {} 180 181 /// \brief Accept an exact value type. 182 InputType(std::shared_ptr<DataType> type, 183 ValueDescr::Shape shape = ValueDescr::ANY) // NOLINT implicit construction kind_(EXACT_TYPE)184 : kind_(EXACT_TYPE), shape_(shape), type_(std::move(type)) {} 185 186 /// \brief Accept an exact value type and shape provided by a ValueDescr. InputType(const ValueDescr & descr)187 InputType(const ValueDescr& descr) // NOLINT implicit construction 188 : InputType(descr.type, descr.shape) {} 189 190 /// \brief Use the passed TypeMatcher to type check. 191 InputType(std::shared_ptr<TypeMatcher> type_matcher, 192 ValueDescr::Shape shape = ValueDescr::ANY) kind_(USE_TYPE_MATCHER)193 : kind_(USE_TYPE_MATCHER), shape_(shape), type_matcher_(std::move(type_matcher)) {} 194 195 /// \brief Match any type with the given Type::type. Uses a TypeMatcher for 196 /// its implementation. 197 explicit InputType(Type::type type_id, ValueDescr::Shape shape = ValueDescr::ANY) InputType(match::SameTypeId (type_id),shape)198 : InputType(match::SameTypeId(type_id), shape) {} 199 InputType(const InputType & other)200 InputType(const InputType& other) { CopyInto(other); } 201 202 void operator=(const InputType& other) { CopyInto(other); } 203 InputType(InputType && other)204 InputType(InputType&& other) { MoveInto(std::forward<InputType>(other)); } 205 206 void operator=(InputType&& other) { MoveInto(std::forward<InputType>(other)); } 207 208 // \brief Match an array with the given exact type. Convenience constructor. Array(std::shared_ptr<DataType> type)209 static InputType Array(std::shared_ptr<DataType> type) { 210 return InputType(std::move(type), ValueDescr::ARRAY); 211 } 212 213 // \brief Match a scalar with the given exact type. Convenience constructor. Scalar(std::shared_ptr<DataType> type)214 static InputType Scalar(std::shared_ptr<DataType> type) { 215 return InputType(std::move(type), ValueDescr::SCALAR); 216 } 217 218 // \brief Match an array with the given Type::type id. Convenience 219 // constructor. Array(Type::type id)220 static InputType Array(Type::type id) { return InputType(id, ValueDescr::ARRAY); } 221 222 // \brief Match a scalar with the given Type::type id. Convenience 223 // constructor. Scalar(Type::type id)224 static InputType Scalar(Type::type id) { return InputType(id, ValueDescr::SCALAR); } 225 226 /// \brief Return true if this input type matches the same type cases as the 227 /// other. 228 bool Equals(const InputType& other) const; 229 230 bool operator==(const InputType& other) const { return this->Equals(other); } 231 232 bool operator!=(const InputType& other) const { return !(*this == other); } 233 234 /// \brief Return hash code. 235 size_t Hash() const; 236 237 /// \brief Render a human-readable string representation. 238 std::string ToString() const; 239 240 /// \brief Return true if the value matches this argument kind in type 241 /// and shape. 242 bool Matches(const Datum& value) const; 243 244 /// \brief Return true if the value descriptor matches this argument kind in 245 /// type and shape. 246 bool Matches(const ValueDescr& value) const; 247 248 /// \brief The type matching rule that this InputType uses. kind()249 Kind kind() const { return kind_; } 250 251 /// \brief Indicates whether this InputType matches Array (ValueDescr::ARRAY), 252 /// Scalar (ValueDescr::SCALAR) values, or both (ValueDescr::ANY). shape()253 ValueDescr::Shape shape() const { return shape_; } 254 255 /// \brief For InputType::EXACT_TYPE kind, the exact type that this InputType 256 /// must match. Otherwise this function should not be used and will assert in 257 /// debug builds. 258 const std::shared_ptr<DataType>& type() const; 259 260 /// \brief For InputType::USE_TYPE_MATCHER, the TypeMatcher to be used for 261 /// checking the type of a value. Otherwise this function should not be used 262 /// and will assert in debug builds. 263 const TypeMatcher& type_matcher() const; 264 265 private: CopyInto(const InputType & other)266 void CopyInto(const InputType& other) { 267 this->kind_ = other.kind_; 268 this->shape_ = other.shape_; 269 this->type_ = other.type_; 270 this->type_matcher_ = other.type_matcher_; 271 } 272 MoveInto(InputType && other)273 void MoveInto(InputType&& other) { 274 this->kind_ = other.kind_; 275 this->shape_ = other.shape_; 276 this->type_ = std::move(other.type_); 277 this->type_matcher_ = std::move(other.type_matcher_); 278 } 279 280 Kind kind_; 281 282 ValueDescr::Shape shape_ = ValueDescr::ANY; 283 284 // For EXACT_TYPE Kind 285 std::shared_ptr<DataType> type_; 286 287 // For USE_TYPE_MATCHER Kind 288 std::shared_ptr<TypeMatcher> type_matcher_; 289 }; 290 291 /// \brief Container to capture both exact and input-dependent output types. 292 /// 293 /// The value shape returned by Resolve will be determined by broadcasting the 294 /// shapes of the input arguments, otherwise this is handled by the 295 /// user-defined resolver function: 296 /// 297 /// * Any ARRAY shape -> output shape is ARRAY 298 /// * All SCALAR shapes -> output shape is SCALAR 299 class ARROW_EXPORT OutputType { 300 public: 301 /// \brief An enum indicating whether the value type is an invariant fixed 302 /// value or one that's computed by a kernel-defined resolver function. 303 enum ResolveKind { FIXED, COMPUTED }; 304 305 /// Type resolution function. Given input types and shapes, return output 306 /// type and shape. This function SHOULD _not_ be used to check for arity, 307 /// that is to be performed one or more layers above. May make use of kernel 308 /// state to know what type to output in some cases. 309 using Resolver = 310 std::function<Result<ValueDescr>(KernelContext*, const std::vector<ValueDescr>&)>; 311 312 /// \brief Output an exact type, but with shape determined by promoting the 313 /// shapes of the inputs (any ARRAY argument yields ARRAY). OutputType(std::shared_ptr<DataType> type)314 OutputType(std::shared_ptr<DataType> type) // NOLINT implicit construction 315 : kind_(FIXED), type_(std::move(type)) {} 316 317 /// \brief Output the exact type and shape provided by a ValueDescr 318 OutputType(ValueDescr descr); // NOLINT implicit construction 319 OutputType(Resolver resolver)320 explicit OutputType(Resolver resolver) : kind_(COMPUTED), resolver_(resolver) {} 321 OutputType(const OutputType & other)322 OutputType(const OutputType& other) { 323 this->kind_ = other.kind_; 324 this->shape_ = other.shape_; 325 this->type_ = other.type_; 326 this->resolver_ = other.resolver_; 327 } 328 OutputType(OutputType && other)329 OutputType(OutputType&& other) { 330 this->kind_ = other.kind_; 331 this->type_ = std::move(other.type_); 332 this->shape_ = other.shape_; 333 this->resolver_ = other.resolver_; 334 } 335 336 /// \brief Return the shape and type of the expected output value of the 337 /// kernel given the value descriptors (shapes and types) of the input 338 /// arguments. The resolver may make use of state information kept in the 339 /// KernelContext. 340 Result<ValueDescr> Resolve(KernelContext* ctx, 341 const std::vector<ValueDescr>& args) const; 342 343 /// \brief The exact output value type for the FIXED kind. 344 const std::shared_ptr<DataType>& type() const; 345 346 /// \brief For use with COMPUTED resolution strategy. It may be more 347 /// convenient to invoke this with OutputType::Resolve returned from this 348 /// method. 349 const Resolver& resolver() const; 350 351 /// \brief Render a human-readable string representation. 352 std::string ToString() const; 353 354 /// \brief Return the kind of type resolution of this output type, whether 355 /// fixed/invariant or computed by a resolver. kind()356 ResolveKind kind() const { return kind_; } 357 358 /// \brief If the shape is ANY, then Resolve will compute the shape based on 359 /// the input arguments. shape()360 ValueDescr::Shape shape() const { return shape_; } 361 362 private: 363 ResolveKind kind_; 364 365 // For FIXED resolution 366 std::shared_ptr<DataType> type_; 367 368 /// \brief The shape of the output type to return when using Resolve. If ANY 369 /// will promote the input shapes. 370 ValueDescr::Shape shape_ = ValueDescr::ANY; 371 372 // For COMPUTED resolution 373 Resolver resolver_; 374 }; 375 376 /// \brief Holds the input types and output type of the kernel. 377 /// 378 /// VarArgs functions should pass a single input type to be used to validate 379 /// the input types of a function invocation. 380 class ARROW_EXPORT KernelSignature { 381 public: 382 KernelSignature(std::vector<InputType> in_types, OutputType out_type, 383 bool is_varargs = false); 384 385 /// \brief Convenience ctor since make_shared can be awkward 386 static std::shared_ptr<KernelSignature> Make(std::vector<InputType> in_types, 387 OutputType out_type, 388 bool is_varargs = false); 389 390 /// \brief Return true if the signature if compatible with the list of input 391 /// value descriptors. 392 bool MatchesInputs(const std::vector<ValueDescr>& descriptors) const; 393 394 /// \brief Returns true if the input types of each signature are 395 /// equal. Well-formed functions should have a deterministic output type 396 /// given input types, but currently it is the responsibility of the 397 /// developer to ensure this. 398 bool Equals(const KernelSignature& other) const; 399 400 bool operator==(const KernelSignature& other) const { return this->Equals(other); } 401 402 bool operator!=(const KernelSignature& other) const { return !(*this == other); } 403 404 /// \brief Compute a hash code for the signature 405 size_t Hash() const; 406 407 /// \brief The input types for the kernel. For VarArgs functions, this should 408 /// generally contain a single validator to use for validating all of the 409 /// function arguments. in_types()410 const std::vector<InputType>& in_types() const { return in_types_; } 411 412 /// \brief The output type for the kernel. Use Resolve to return the exact 413 /// output given input argument ValueDescrs, since many kernels' output types 414 /// depend on their input types (or their type metadata). out_type()415 const OutputType& out_type() const { return out_type_; } 416 417 /// \brief Render a human-readable string representation 418 std::string ToString() const; 419 is_varargs()420 bool is_varargs() const { return is_varargs_; } 421 422 private: 423 std::vector<InputType> in_types_; 424 OutputType out_type_; 425 bool is_varargs_; 426 427 // For caching the hash code after it's computed the first time 428 mutable uint64_t hash_code_; 429 }; 430 431 /// \brief A function may contain multiple variants of a kernel for a given 432 /// type combination for different SIMD levels. Based on the active system's 433 /// CPU info or the user's preferences, we can elect to use one over the other. 434 struct SimdLevel { 435 enum type { NONE, SSE4_2, AVX, AVX2, AVX512, NEON }; 436 }; 437 438 /// \brief The strategy to use for propagating or otherwise populating the 439 /// validity bitmap of a kernel output. 440 struct NullHandling { 441 enum type { 442 /// Compute the output validity bitmap by intersecting the validity bitmaps 443 /// of the arguments using bitwise-and operations. This means that values 444 /// in the output are valid/non-null only if the corresponding values in 445 /// all input arguments were valid/non-null. Kernel generally need not 446 /// touch the bitmap thereafter, but a kernel's exec function is permitted 447 /// to alter the bitmap after the null intersection is computed if it needs 448 /// to. 449 INTERSECTION, 450 451 /// Kernel expects a pre-allocated buffer to write the result bitmap 452 /// into. The preallocated memory is not zeroed (except for the last byte), 453 /// so the kernel should ensure to completely populate the bitmap. 454 COMPUTED_PREALLOCATE, 455 456 /// Kernel allocates and sets the validity bitmap of the output. 457 COMPUTED_NO_PREALLOCATE, 458 459 /// Kernel output is never null and a validity bitmap does not need to be 460 /// allocated. 461 OUTPUT_NOT_NULL 462 }; 463 }; 464 465 /// \brief The preference for memory preallocation of fixed-width type outputs 466 /// in kernel execution. 467 struct MemAllocation { 468 enum type { 469 // For data types that support pre-allocation (i.e. fixed-width), the 470 // kernel expects to be provided a pre-allocated data buffer to write 471 // into. Non-fixed-width types must always allocate their own data 472 // buffers. The allocation made for the same length as the execution batch, 473 // so vector kernels yielding differently sized output should not use this. 474 // 475 // It is valid for the data to not be preallocated but the validity bitmap 476 // is (or is computed using the intersection/bitwise-and method). 477 // 478 // For variable-size output types like BinaryType or StringType, or for 479 // nested types, this option has no effect. 480 PREALLOCATE, 481 482 // The kernel is responsible for allocating its own data buffer for 483 // fixed-width type outputs. 484 NO_PREALLOCATE 485 }; 486 }; 487 488 struct Kernel; 489 490 /// \brief Arguments to pass to a KernelInit function. A struct is used to help 491 /// avoid API breakage should the arguments passed need to be expanded. 492 struct KernelInitArgs { 493 /// \brief A pointer to the kernel being initialized. The init function may 494 /// depend on the kernel's KernelSignature or other data contained there. 495 const Kernel* kernel; 496 497 /// \brief The types and shapes of the input arguments that the kernel is 498 /// about to be executed against. 499 /// 500 /// TODO: should this be const std::vector<ValueDescr>*? const-ref is being 501 /// used to avoid the cost of copying the struct into the args struct. 502 const std::vector<ValueDescr>& inputs; 503 504 /// \brief Opaque options specific to this kernel. Is nullptr for functions 505 /// that do not require options. 506 const FunctionOptions* options; 507 }; 508 509 /// \brief Common initializer function for all kernel types. 510 using KernelInit = 511 std::function<std::unique_ptr<KernelState>(KernelContext*, const KernelInitArgs&)>; 512 513 /// \brief Base type for kernels. Contains the function signature and 514 /// optionally the state initialization function, along with some common 515 /// attributes 516 struct Kernel { KernelKernel517 Kernel() {} 518 KernelKernel519 Kernel(std::shared_ptr<KernelSignature> sig, KernelInit init) 520 : signature(std::move(sig)), init(init) {} 521 KernelKernel522 Kernel(std::vector<InputType> in_types, OutputType out_type, KernelInit init) 523 : Kernel(KernelSignature::Make(std::move(in_types), out_type), init) {} 524 525 /// \brief The "signature" of the kernel containing the InputType input 526 /// argument validators and OutputType output type and shape resolver. 527 std::shared_ptr<KernelSignature> signature; 528 529 /// \brief Create a new KernelState for invocations of this kernel, e.g. to 530 /// set up any options or state relevant for execution. May be nullptr 531 KernelInit init; 532 533 /// \brief Indicates whether execution can benefit from parallelization 534 /// (splitting large chunks into smaller chunks and using multiple 535 /// threads). Some kernels may not support parallel execution at 536 /// all. Synchronization and concurrency-related issues are currently the 537 /// responsibility of the Kernel's implementation. 538 bool parallelizable = true; 539 540 /// \brief Indicates the level of SIMD instruction support in the host CPU is 541 /// required to use the function. Currently this is not used, but the 542 /// intention is for functions to be able to contain multiple kernels with 543 /// the same signature but different levels of SIMD, so that the most 544 /// optimized kernel supported on a host's processor can be chosen. 545 SimdLevel::type simd_level = SimdLevel::NONE; 546 }; 547 548 /// \brief Common kernel base data structure for ScalarKernel and 549 /// VectorKernel. It is called "ArrayKernel" in that the functions generally 550 /// output array values (as opposed to scalar values in the case of aggregate 551 /// functions). 552 struct ArrayKernel : public Kernel { ArrayKernelArrayKernel553 ArrayKernel() {} 554 555 ArrayKernel(std::shared_ptr<KernelSignature> sig, ArrayKernelExec exec, 556 KernelInit init = NULLPTR) KernelArrayKernel557 : Kernel(std::move(sig), init), exec(exec) {} 558 559 ArrayKernel(std::vector<InputType> in_types, OutputType out_type, ArrayKernelExec exec, 560 KernelInit init = NULLPTR) KernelArrayKernel561 : Kernel(std::move(in_types), std::move(out_type), init), exec(exec) {} 562 563 /// \brief Perform a single invocation of this kernel. Depending on the 564 /// implementation, it may only write into preallocated memory, while in some 565 /// cases it will allocate its own memory. Any required state is managed 566 /// through the KernelContext. 567 ArrayKernelExec exec; 568 569 /// \brief Writing execution results into larger contiguous allocations 570 /// requires that the kernel be able to write into sliced output ArrayData*, 571 /// including sliced output validity bitmaps. Some kernel implementations may 572 /// not be able to do this, so setting this to false disables this 573 /// functionality. 574 bool can_write_into_slices = true; 575 }; 576 577 /// \brief Kernel data structure for implementations of ScalarFunction. In 578 /// addition to the members found in ArrayKernel, contains the null handling 579 /// and memory pre-allocation preferences. 580 struct ScalarKernel : public ArrayKernel { 581 using ArrayKernel::ArrayKernel; 582 583 // For scalar functions preallocated data and intersecting arg validity 584 // bitmaps is a reasonable default 585 NullHandling::type null_handling = NullHandling::INTERSECTION; 586 MemAllocation::type mem_allocation = MemAllocation::PREALLOCATE; 587 }; 588 589 // ---------------------------------------------------------------------- 590 // VectorKernel (for VectorFunction) 591 592 /// \brief See VectorKernel::finalize member for usage 593 using VectorFinalize = std::function<void(KernelContext*, std::vector<Datum>*)>; 594 595 /// \brief Kernel data structure for implementations of VectorFunction. In 596 /// addition to the members found in ArrayKernel, contains an optional 597 /// finalizer function, the null handling and memory pre-allocation preferences 598 /// (which have different defaults from ScalarKernel), and some other 599 /// execution-related options. 600 struct VectorKernel : public ArrayKernel { VectorKernelVectorKernel601 VectorKernel() {} 602 VectorKernelVectorKernel603 VectorKernel(std::shared_ptr<KernelSignature> sig, ArrayKernelExec exec) 604 : ArrayKernel(std::move(sig), exec) {} 605 606 VectorKernel(std::vector<InputType> in_types, OutputType out_type, ArrayKernelExec exec, 607 KernelInit init = NULLPTR, VectorFinalize finalize = NULLPTR) ArrayKernelVectorKernel608 : ArrayKernel(std::move(in_types), out_type, exec, init), finalize(finalize) {} 609 610 VectorKernel(std::shared_ptr<KernelSignature> sig, ArrayKernelExec exec, 611 KernelInit init = NULLPTR, VectorFinalize finalize = NULLPTR) ArrayKernelVectorKernel612 : ArrayKernel(std::move(sig), exec, init), finalize(finalize) {} 613 614 /// \brief For VectorKernel, convert intermediate results into finalized 615 /// results. Mutates input argument. Some kernels may accumulate state 616 /// (example: hashing-related functions) through processing chunked inputs, and 617 /// then need to attach some accumulated state to each of the outputs of 618 /// processing each chunk of data. 619 VectorFinalize finalize; 620 621 /// Since vector kernels generally are implemented rather differently from 622 /// scalar/elementwise kernels (and they may not even yield arrays of the same 623 /// size), so we make the developer opt-in to any memory preallocation rather 624 /// than having to turn it off. 625 NullHandling::type null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; 626 MemAllocation::type mem_allocation = MemAllocation::NO_PREALLOCATE; 627 628 /// Some vector kernels can do chunkwise execution using ExecBatchIterator, 629 /// in some cases accumulating some state. Other kernels (like Take) need to 630 /// be passed whole arrays and don't work on ChunkedArray inputs 631 bool can_execute_chunkwise = true; 632 633 /// Some kernels (like unique and value_counts) yield non-chunked output from 634 /// chunked-array inputs. This option controls how the results are boxed when 635 /// returned from ExecVectorFunction 636 /// 637 /// true -> ChunkedArray 638 /// false -> Array 639 bool output_chunked = true; 640 }; 641 642 // ---------------------------------------------------------------------- 643 // ScalarAggregateKernel (for ScalarAggregateFunction) 644 645 using ScalarAggregateConsume = std::function<void(KernelContext*, const ExecBatch&)>; 646 647 using ScalarAggregateMerge = 648 std::function<void(KernelContext*, const KernelState&, KernelState*)>; 649 650 // Finalize returns Datum to permit multiple return values 651 using ScalarAggregateFinalize = std::function<void(KernelContext*, Datum*)>; 652 653 /// \brief Kernel data structure for implementations of 654 /// ScalarAggregateFunction. The four necessary components of an aggregation 655 /// kernel are the init, consume, merge, and finalize functions. 656 /// 657 /// * init: creates a new KernelState for a kernel. 658 /// * consume: processes an ExecBatch and updates the KernelState found in the 659 /// KernelContext. 660 /// * merge: combines one KernelState with another. 661 /// * finalize: produces the end result of the aggregation using the 662 /// KernelState in the KernelContext. 663 struct ScalarAggregateKernel : public Kernel { ScalarAggregateKernelScalarAggregateKernel664 ScalarAggregateKernel() {} 665 ScalarAggregateKernelScalarAggregateKernel666 ScalarAggregateKernel(std::shared_ptr<KernelSignature> sig, KernelInit init, 667 ScalarAggregateConsume consume, ScalarAggregateMerge merge, 668 ScalarAggregateFinalize finalize) 669 : Kernel(std::move(sig), init), 670 consume(consume), 671 merge(merge), 672 finalize(finalize) {} 673 ScalarAggregateKernelScalarAggregateKernel674 ScalarAggregateKernel(std::vector<InputType> in_types, OutputType out_type, 675 KernelInit init, ScalarAggregateConsume consume, 676 ScalarAggregateMerge merge, ScalarAggregateFinalize finalize) 677 : ScalarAggregateKernel(KernelSignature::Make(std::move(in_types), out_type), init, 678 consume, merge, finalize) {} 679 680 ScalarAggregateConsume consume; 681 ScalarAggregateMerge merge; 682 ScalarAggregateFinalize finalize; 683 }; 684 685 } // namespace compute 686 } // namespace arrow 687