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