ARROW-16824: [C++] Migrate VectorKernels to use ExecSpan, split out ChunkedArray execution

cpp/src/arrow/array/array_primitive.cc

@@ -58,18 +58,9 @@ int64_t BooleanArray::false_count() const {
 int64_t BooleanArray::true_count() const {
   if (data_->null_count.load() != 0) {
     DCHECK(data_->buffers[0]);
-    internal::BinaryBitBlockCounter bit_counter(data_->buffers[0]->data(), data_->offset,
-                                                data_->buffers[1]->data(), data_->offset,
-                                                data_->length);
-    int64_t count = 0;
-    while (true) {
-      internal::BitBlockCount block = bit_counter.NextAndWord();
-      if (block.length == 0) {
-        break;
-      }
-      count += block.popcount;
-    }
-    return count;
+    return internal::CountAndSetBits(data_->buffers[0]->data(), data_->offset,
+                                     data_->buffers[1]->data(), data_->offset,
+                                     data_->length);
   } else {
     return internal::CountSetBits(data_->buffers[1]->data(), data_->offset,
                                   data_->length);

cpp/src/arrow/array/data.cc

@@ -151,6 +151,13 @@ void ArraySpan::SetMembers(const ArrayData& data) {
     }
   }
 
+  Type::type type_id = this->type->id();
+  if (data.buffers[0] == nullptr && type_id != Type::NA &&
+      type_id != Type::SPARSE_UNION && type_id != Type::DENSE_UNION) {
+    // This should already be zero but we make for sure
+    this->null_count = 0;
+  }
+
   // Makes sure any other buffers are seen as null / non-existent
   for (int i = static_cast<int>(data.buffers.size()); i < 3; ++i) {
     ClearBuffer(i);
@@ -208,7 +215,6 @@ int64_t ArraySpan::GetNullCount() const {
 int GetNumBuffers(const DataType& type) {
   switch (type.id()) {
     case Type::NA:
-      return 0;
     case Type::STRUCT:
     case Type::FIXED_SIZE_LIST:
       return 1;
@@ -232,7 +238,7 @@ int ArraySpan::num_buffers() const { return GetNumBuffers(*this->type); }
 
 std::shared_ptr<ArrayData> ArraySpan::ToArrayData() const {
   auto result = std::make_shared<ArrayData>(this->type->Copy(), this->length,
-                                            kUnknownNullCount, this->offset);
+                                            this->null_count, this->offset);
 
   for (int i = 0; i < this->num_buffers(); ++i) {
     if (this->buffers[i].owner) {

cpp/src/arrow/builder.cc

@@ -170,7 +170,7 @@ struct DictionaryBuilderCase {
       out->reset(new internal::DictionaryBuilderBase<TypeErasedIntBuilder, ValueType>(
           index_type, value_type, pool));
     } else {
-      auto start_int_size = internal::GetByteWidth(*index_type);
+      auto start_int_size = index_type->byte_width();
       out->reset(new AdaptiveBuilderType(start_int_size, value_type, pool));
     }
     return Status::OK();

cpp/src/arrow/compare.cc

@@ -1022,10 +1022,10 @@ bool IntegerTensorEquals(const Tensor& left, const Tensor& right) {
     if (!(left_row_major_p && right_row_major_p) &&
         !(left_column_major_p && right_column_major_p)) {
       const auto& type = checked_cast<const FixedWidthType&>(*left.type());
-      are_equal = StridedIntegerTensorContentEquals(0, 0, 0, internal::GetByteWidth(type),
-                                                    left, right);
+      are_equal =
+          StridedIntegerTensorContentEquals(0, 0, 0, type.byte_width(), left, right);
     } else {
-      const int byte_width = internal::GetByteWidth(*left.type());
+      const int byte_width = left.type()->byte_width();
       DCHECK_GT(byte_width, 0);
 
       const uint8_t* left_data = left.data()->data();
@@ -1195,7 +1195,7 @@ struct SparseTensorEqualsImpl<SparseIndexType, SparseIndexType> {
       return false;
     }
 
-    const int byte_width = internal::GetByteWidth(*left.type());
+    const int byte_width = left.type()->byte_width();
     DCHECK_GT(byte_width, 0);
 
     const uint8_t* left_data = left.data()->data();

cpp/src/arrow/compute/api_vector.cc

@@ -17,6 +17,7 @@
 
 #include "arrow/compute/api_vector.h"
 
+#include <algorithm>
 #include <memory>
 #include <sstream>
 #include <utility>
@@ -26,6 +27,7 @@
 #include "arrow/array/builder_primitive.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/function_internal.h"
+#include "arrow/compute/kernels/vector_sort_internal.h"
 #include "arrow/compute/registry.h"
 #include "arrow/datum.h"
 #include "arrow/record_batch.h"
@@ -305,10 +307,7 @@ Result<std::shared_ptr<Array>> SortIndices(const ChunkedArray& chunked_array,
 
 Result<std::shared_ptr<Array>> SortIndices(const ChunkedArray& chunked_array,
                                            SortOrder order, ExecContext* ctx) {
-  SortOptions options({SortKey("not-used", order)});
-  ARROW_ASSIGN_OR_RAISE(
-      Datum result, CallFunction("sort_indices", {Datum(chunked_array)}, &options, ctx));
-  return result.make_array();
+  return SortIndices(chunked_array, ArraySortOptions(order), ctx);
 }
 
 Result<std::shared_ptr<Array>> SortIndices(const Datum& datum, const SortOptions& options,

cpp/src/arrow/compute/api_vector.h

@@ -275,14 +275,14 @@ namespace internal {
 
 /// \brief Return the number of selected indices in the boolean filter
 ARROW_EXPORT
-int64_t GetFilterOutputSize(const ArrayData& filter,
+int64_t GetFilterOutputSize(const ArraySpan& filter,
                             FilterOptions::NullSelectionBehavior null_selection);
 
 /// \brief Compute uint64 selection indices for use with Take given a boolean
 /// filter
 ARROW_EXPORT
 Result<std::shared_ptr<ArrayData>> GetTakeIndices(
-    const ArrayData& filter, FilterOptions::NullSelectionBehavior null_selection,
+    const ArraySpan& filter, FilterOptions::NullSelectionBehavior null_selection,
     MemoryPool* memory_pool = default_memory_pool());
 
 }  // namespace internal

cpp/src/arrow/compute/exec.cc

@@ -332,10 +332,14 @@ Status ExecSpanIterator::Init(const ExecBatch& batch, ValueDescr::Shape output_s
                               int64_t max_chunksize) {
   if (batch.num_values() > 0) {
     // Validate arguments
-    ARROW_ASSIGN_OR_RAISE(int64_t inferred_length, InferBatchLength(batch.values));
+    bool all_args_same_length = false;
+    int64_t inferred_length = InferBatchLength(batch.values, &all_args_same_length);
     if (inferred_length != batch.length) {
       return Status::Invalid("Value lengths differed from ExecBatch length");
     }
+    if (!all_args_same_length) {
+      return Status::Invalid("Array arguments must all be the same length");
+    }
   }
   args_ = &batch.values;
   initialized_ = have_chunked_arrays_ = false;
@@ -991,43 +995,62 @@ class ScalarExecutor : public KernelExecutorImpl<ScalarKernel> {
   ExecSpanIterator span_iterator_;
 };
 
-Status PackBatchNoChunks(const std::vector<Datum>& args, ExecBatch* out) {
-  int64_t length = 0;
-  for (const auto& arg : args) {
-    switch (arg.kind()) {
-      case Datum::SCALAR:
-      case Datum::ARRAY:
-      case Datum::CHUNKED_ARRAY:
-        length = std::max(arg.length(), length);
-        break;
-      default:
-        DCHECK(false);
-        break;
-    }
-  }
-  out->length = length;
-  out->values = args;
-  return Status::OK();
-}
-
 class VectorExecutor : public KernelExecutorImpl<VectorKernel> {
  public:
-  Status ExecuteImpl(const std::vector<Datum>& args, ExecListener* listener) {
-    RETURN_NOT_OK(PrepareExecute(args));
-    ExecBatch batch;
+  Status Execute(const ExecBatch& batch, ExecListener* listener) override {
+    // TODO(wesm): remove in ARROW-16577
+    if (output_descr_.shape == ValueDescr::SCALAR) {
+      return Status::Invalid("VectorExecutor only supports array output types");
+    }
+
+    // Some vector kernels have a separate code path for handling
+    // chunked arrays (VectorKernel::exec_chunked) so we check if we
+    // have any chunked arrays. If we do and an exec_chunked function
+    // is defined then we call that.
+    bool have_chunked_arrays = false;
+    for (const Datum& arg : batch.values) {
+      if (arg.is_chunked_array()) have_chunked_arrays = true;
+    }
+
+    output_num_buffers_ = static_cast<int>(output_descr_.type->layout().buffers.size());
+
+    // Decide if we need to preallocate memory for this kernel
+    validity_preallocated_ =
+        (kernel_->null_handling != NullHandling::COMPUTED_NO_PREALLOCATE &&
+         kernel_->null_handling != NullHandling::OUTPUT_NOT_NULL);
+    if (kernel_->mem_allocation == MemAllocation::PREALLOCATE) {
+      ComputeDataPreallocate(*output_descr_.type, &data_preallocated_);
+    }
+
     if (kernel_->can_execute_chunkwise) {
-      while (batch_iterator_->Next(&batch)) {
-        RETURN_NOT_OK(ExecuteBatch(batch, listener));
+      RETURN_NOT_OK(span_iterator_.Init(batch, output_descr_.shape,
+                                        exec_context()->exec_chunksize()));
+      ExecSpan span;
+      while (span_iterator_.Next(&span)) {
+        RETURN_NOT_OK(Exec(span, listener));
       }
     } else {
-      RETURN_NOT_OK(PackBatchNoChunks(args, &batch));
-      RETURN_NOT_OK(ExecuteBatch(batch, listener));
+      // Kernel cannot execute chunkwise. If we have any chunked
+      // arrays, then VectorKernel::exec_chunked must be defined
+      // otherwise we raise an error
+      if (have_chunked_arrays) {
+        RETURN_NOT_OK(ExecChunked(batch, listener));
+      } else {
+        // No chunked arrays. We pack the args into an ExecSpan and
+        // call the regular exec code path
+        RETURN_NOT_OK(Exec(ExecSpan(batch), listener));
+      }
     }
-    return Finalize(listener);
-  }
 
-  Status Execute(const ExecBatch& batch, ExecListener* listener) override {
-    return ExecuteImpl(batch.values, listener);
+    if (kernel_->finalize) {
+      // Intermediate results require post-processing after the execution is
+      // completed (possibly involving some accumulated state)
+      RETURN_NOT_OK(kernel_->finalize(kernel_ctx_, &results_));
+      for (const auto& result : results_) {
+        RETURN_NOT_OK(listener->OnResult(result));
+      }
+    }
+    return Status::OK();
   }
 
   Datum WrapResults(const std::vector<Datum>& inputs,
@@ -1047,59 +1070,54 @@ class VectorExecutor : public KernelExecutorImpl<VectorKernel> {
   }
 
  protected:
-  Status ExecuteBatch(const ExecBatch& batch, ExecListener* listener) {
-    Datum out;
-    if (output_descr_.shape == ValueDescr::ARRAY) {
-      // We preallocate (maybe) only for the output of processing the current
-      // batch
-      ARROW_ASSIGN_OR_RAISE(out.value, PrepareOutput(batch.length));
-    }
+  Status Exec(const ExecSpan& span, ExecListener* listener) {
+    ExecResult out;
 
-    if (kernel_->null_handling == NullHandling::INTERSECTION &&
-        output_descr_.shape == ValueDescr::ARRAY) {
-      RETURN_NOT_OK(PropagateNulls(kernel_ctx_, ExecSpan(batch), out.mutable_array()));
+    // We preallocate (maybe) only for the output of processing the current
+    // batch, but create an output ArrayData instance regardless
+    ARROW_ASSIGN_OR_RAISE(out.value, PrepareOutput(span.length));
+
+    if (kernel_->null_handling == NullHandling::INTERSECTION) {
+      RETURN_NOT_OK(PropagateNulls(kernel_ctx_, span, out.array_data().get()));
     }
-    RETURN_NOT_OK(kernel_->exec(kernel_ctx_, batch, &out));
+    RETURN_NOT_OK(kernel_->exec(kernel_ctx_, span, &out));
     if (!kernel_->finalize) {
       // If there is no result finalizer (e.g. for hash-based functions, we can
       // emit the processed batch right away rather than waiting
-      RETURN_NOT_OK(listener->OnResult(std::move(out)));
+      RETURN_NOT_OK(listener->OnResult(out.array_data()));
     } else {
-      results_.emplace_back(std::move(out));
+      results_.emplace_back(out.array_data());
     }
     return Status::OK();
   }
 
-  Status Finalize(ExecListener* listener) {
-    if (kernel_->finalize) {
-      // Intermediate results require post-processing after the execution is
-      // completed (possibly involving some accumulated state)
-      RETURN_NOT_OK(kernel_->finalize(kernel_ctx_, &results_));
-      for (const auto& result : results_) {
-        RETURN_NOT_OK(listener->OnResult(result));
-      }
+  Status ExecChunked(const ExecBatch& batch, ExecListener* listener) {
+    if (kernel_->exec_chunked == nullptr) {
+      return Status::Invalid(
+          "Vector kernel cannot execute chunkwise and no "
+          "chunked exec function was defined");
     }
-    return Status::OK();
-  }
 
-  Status PrepareExecute(const std::vector<Datum>& args) {
-    if (kernel_->can_execute_chunkwise) {
-      ARROW_ASSIGN_OR_RAISE(batch_iterator_, ExecBatchIterator::Make(
-                                                 args, exec_context()->exec_chunksize()));
+    if (kernel_->null_handling == NullHandling::INTERSECTION) {
+      return Status::Invalid(
+          "Null pre-propagation is unsupported for ChunkedArray "
+          "execution in vector kernels");
     }
-    output_num_buffers_ = static_cast<int>(output_descr_.type->layout().buffers.size());
 
-    // Decide if we need to preallocate memory for this kernel
-    validity_preallocated_ =
-        (kernel_->null_handling != NullHandling::COMPUTED_NO_PREALLOCATE &&
-         kernel_->null_handling != NullHandling::OUTPUT_NOT_NULL);
-    if (kernel_->mem_allocation == MemAllocation::PREALLOCATE) {
-      ComputeDataPreallocate(*output_descr_.type, &data_preallocated_);
+    Datum out;
+    ARROW_ASSIGN_OR_RAISE(out.value, PrepareOutput(batch.length));
+    RETURN_NOT_OK(kernel_->exec_chunked(kernel_ctx_, batch, &out));
+    if (!kernel_->finalize) {
+      // If there is no result finalizer (e.g. for hash-based functions, we can
+      // emit the processed batch right away rather than waiting
+      RETURN_NOT_OK(listener->OnResult(std::move(out)));
+    } else {
+      results_.emplace_back(std::move(out));
     }
     return Status::OK();
   }
 
-  std::unique_ptr<ExecBatchIterator> batch_iterator_;
+  ExecSpanIterator span_iterator_;
   std::vector<Datum> results_;
 };
 
@@ -1270,7 +1288,7 @@ std::unique_ptr<KernelExecutor> KernelExecutor::MakeScalarAggregate() {
   return ::arrow::internal::make_unique<detail::ScalarAggExecutor>();
 }
 
-Result<int64_t> InferBatchLength(const std::vector<Datum>& values) {
+int64_t InferBatchLength(const std::vector<Datum>& values, bool* all_same) {
   int64_t length = -1;
   bool are_all_scalar = true;
   for (const Datum& arg : values) {
@@ -1280,7 +1298,8 @@ Result<int64_t> InferBatchLength(const std::vector<Datum>& values) {
         length = arg_length;
       } else {
         if (length != arg_length) {
-          return Status::Invalid("Array arguments must all be the same length");
+          *all_same = false;
+          return length;
         }
       }
       are_all_scalar = false;
@@ -1290,7 +1309,8 @@ Result<int64_t> InferBatchLength(const std::vector<Datum>& values) {
         length = arg_length;
       } else {
         if (length != arg_length) {
-          return Status::Invalid("Array arguments must all be the same length");
+          *all_same = false;
+          return length;
         }
       }
       are_all_scalar = false;
@@ -1302,6 +1322,7 @@ Result<int64_t> InferBatchLength(const std::vector<Datum>& values) {
   } else if (length < 0) {
     length = 0;
   }
+  *all_same = true;
   return length;
 }
 

cpp/src/arrow/compute/exec.h

@@ -20,7 +20,6 @@
 
 #pragma once
 
-#include <algorithm>
 #include <atomic>
 #include <cstdint>
 #include <limits>
@@ -397,8 +396,12 @@ struct ARROW_EXPORT ExecSpan {
   }
 
   bool is_all_scalar() const {
-    return std::all_of(this->values.begin(), this->values.end(),
-                       [](const ExecValue& v) { return v.is_scalar(); });
+    for (const ExecValue& value : this->values) {
+      if (value.is_array()) {
+        return false;
+      }
+    }
+    return true;
   }
 
   /// \brief Return the value at the i-th index