GH-47377: Extract dense selection execution

Author: zanmato1984

cpp/src/arrow/CMakeLists.txt

@@ -722,6 +722,7 @@ set(ARROW_COMPUTE_SRCS
     compute/api_vector.cc
     compute/cast.cc
     compute/exec.cc
+    compute/exec_selection.cc
     compute/expression.cc
     compute/function.cc
     compute/function_internal.cc

cpp/src/arrow/compute/exec.cc

@@ -477,9 +477,8 @@ bool ExecSpanIterator::Next(ExecSpan* span, SelectionSpan* selection_span) {
   // Then the selection span
   if (selection_vector_) {
     DCHECK_NE(selection_span, nullptr);
-    const uint64_t chunk_start = static_cast<uint64_t>(position_);
-    const uint64_t chunk_end =
-        static_cast<uint64_t>(position_) + static_cast<uint64_t>(iteration_size);
+    const int64_t chunk_start = position_;
+    const int64_t chunk_end = position_ + iteration_size;
 
     const int64_t consumed = selection_vector_->GetSpanForChunk(
         chunk_start, chunk_end, selection_position_, selection_span);
@@ -892,12 +891,6 @@ class ScalarExecutor : public KernelExecutorImpl<ScalarKernel> {
       return EmitResult(result->data(), listener);
     }
 
-    if (batch.selection_vector && !kernel_->selective_exec) {
-      // If the batch contains a selection vector but the kernel does not support
-      // selective execution, we need to execute the batch in a "dense" manner.
-      return ExecuteSelectiveDense(batch, listener);
-    }
-
     return ExecuteBatch(batch, listener);
   }
 
@@ -940,39 +933,6 @@ class ScalarExecutor : public KernelExecutorImpl<ScalarKernel> {
     }
   }
 
-  // Execute a single batch with a selection vector "densely" for a kernel that doesn't
-  // support selective execution. "Densely" here means that we first gather the rows
-  // indicated by the selection vector into a contiguous ExecBatch, execute that, and
-  // then scatter the result back to the original row positions in the output.
-  Status ExecuteSelectiveDense(const ExecBatch& batch, ExecListener* listener) {
-    DCHECK(batch.selection_vector && !kernel_->selective_exec);
-
-    if (CheckIfAllScalar(batch)) {
-      // For all-scalar batch, we can skip the gather/scatter steps as if there is no
-      // selection vector - the result is a scalar anyway.
-      ExecBatch input = batch;
-      input.selection_vector = nullptr;
-      return ExecuteBatch(input, listener);
-    }
-
-    ARROW_ASSIGN_OR_RAISE(
-        std::vector<Datum> values,
-        batch.selection_vector->MakeDenseValues(batch.values, exec_context()));
-    ARROW_ASSIGN_OR_RAISE(
-        ExecBatch input,
-        ExecBatch::Make(std::move(values), batch.selection_vector->length()));
-
-    DatumAccumulator dense_listener;
-    RETURN_NOT_OK(ExecuteBatch(input, &dense_listener));
-    Datum dense_result = WrapResults(input.values, dense_listener.values());
-
-    ARROW_ASSIGN_OR_RAISE(
-        auto result,
-        batch.selection_vector->ScatterDenseResult(dense_result, batch.length,
-                                                   exec_context()));
-    return listener->OnResult(std::move(result));
-  }
-
   Status EmitResult(std::shared_ptr<ArrayData> out, ExecListener* listener) {
     if (span_iterator_.have_all_scalars()) {
       // ARROW-16757 We boxed scalar inputs as ArraySpan, so now we have to
@@ -1569,7 +1529,7 @@ void PropagateNullsSpans(const ExecSpan& batch, ArraySpan* out) {
 }
 
 std::unique_ptr<KernelExecutor> KernelExecutor::MakeScalar() {
-  return std::make_unique<detail::ScalarExecutor>();
+  return MakeDenseSelectionExecutor(std::make_unique<detail::ScalarExecutor>());
 }
 
 std::unique_ptr<KernelExecutor> KernelExecutor::MakeVector() {
@@ -1708,26 +1668,27 @@ class IndexSelectionVector final : public SelectionVector {
     return Scatter(dense_result, indices, ScatterOptions{/*max_index=*/output_length - 1});
   }
 
-  int64_t GetSpanForChunk(uint64_t chunk_start, uint64_t chunk_end,
+  int64_t GetSpanForChunk(int64_t chunk_start, int64_t chunk_end,
                           int64_t selection_position,
                           SelectionSpan* out) const override {
     DCHECK_NE(out, nullptr);
+    DCHECK_GE(chunk_start, 0);
     DCHECK_LE(chunk_start, chunk_end);
 
     const int32_t* indices_begin = indices_ + selection_position;
     const int32_t* indices_end = indices_ + length();
     DCHECK_LE(indices_begin, indices_end);
 
-    const int32_t chunk_end_i32 = static_cast<int32_t>(chunk_end);
-    const int32_t* indices_limit =
-        std::lower_bound(indices_begin, indices_end, chunk_end_i32);
+    const int32_t* indices_limit = std::lower_bound(
+        indices_begin, indices_end, chunk_end,
+        [](int32_t index, int64_t end) { return static_cast<int64_t>(index) < end; });
     const int64_t num_indices = indices_limit - indices_begin;
 
     if (num_indices > 0) {
       const int32_t first = indices_begin[0];
       const int32_t last = indices_begin[num_indices - 1];
-      DCHECK_GE(static_cast<uint64_t>(first), chunk_start);
-      DCHECK_LT(static_cast<uint64_t>(last), chunk_end);
+      DCHECK_GE(static_cast<int64_t>(first), chunk_start);
+      DCHECK_LT(static_cast<int64_t>(last), chunk_end);
 
       // If the discrete indices form a contiguous run, represent them as such.
       // Since Validate enforces strict increasing order, checking
@@ -1737,10 +1698,10 @@ class IndexSelectionVector final : public SelectionVector {
                                   static_cast<int64_t>(chunk_start),
                               num_indices};
       } else {
-        *out = DiscreteSpan{indices_begin, num_indices, static_cast<int32_t>(chunk_start)};
+        *out = DiscreteSpan{indices_begin, num_indices, chunk_start};
       }
     } else {
-      *out = DiscreteSpan{indices_begin, /*length=*/0, static_cast<int32_t>(chunk_start)};
+      *out = DiscreteSpan{indices_begin, /*length=*/0, chunk_start};
     }
 
     return num_indices;

cpp/src/arrow/compute/exec.h

@@ -155,10 +155,10 @@ struct ARROW_EXPORT FilteredSpan {
 struct ARROW_EXPORT DiscreteSpan {
   const int32_t* indices = NULLPTR;
   int64_t length = 0;
-  int32_t index_back_shift = 0;
+  int64_t index_back_shift = 0;
 
   int64_t operator[](int64_t i) const {
-    return static_cast<int64_t>(indices[i] - index_back_shift);
+    return static_cast<int64_t>(indices[i]) - index_back_shift;
   }
 };
 
@@ -224,7 +224,7 @@ class ARROW_EXPORT SelectionVector {
   /// in earlier chunks (rank).
   /// \param[out] out Selection span relative to chunk_start.
   /// \return Number of selected indices consumed from this chunk.
-  virtual int64_t GetSpanForChunk(uint64_t chunk_start, uint64_t chunk_end,
+  virtual int64_t GetSpanForChunk(int64_t chunk_start, int64_t chunk_end,
                                   int64_t selection_position,
                                   SelectionSpan* out) const = 0;
 };

cpp/src/arrow/compute/exec_internal.h

@@ -156,6 +156,14 @@ class ARROW_EXPORT KernelExecutor {
   static std::unique_ptr<KernelExecutor> MakeScalarAggregate();
 };
 
+/// \brief Wrap a scalar kernel executor with dense selection fallback.
+///
+/// Batches without a selection, and selections handled natively by the kernel, are
+/// delegated unchanged. For kernels without selective execution, selected inputs are
+/// gathered before delegation and the dense result is scattered back afterwards.
+std::unique_ptr<KernelExecutor> MakeDenseSelectionExecutor(
+    std::unique_ptr<KernelExecutor> executor);
+
 ARROW_EXPORT
 int64_t InferBatchLength(const std::vector<Datum>& values, bool* all_same);
 

cpp/src/arrow/compute/exec_selection.cc

@@ -0,0 +1,112 @@
+// 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/exec_internal.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "arrow/compute/kernel.h"
+#include "arrow/datum.h"
+#include "arrow/status.h"
+#include "arrow/util/logging_internal.h"
+
+namespace arrow::compute::detail {
+
+namespace {
+
+bool CheckIfAllScalar(const ExecBatch& batch) {
+  for (const Datum& value : batch.values) {
+    if (!value.is_scalar()) {
+      return false;
+    }
+  }
+  return batch.num_values() > 0;
+}
+
+class DenseSelectionExecutor : public KernelExecutor {
+ public:
+  explicit DenseSelectionExecutor(std::unique_ptr<KernelExecutor> executor)
+      : executor_(std::move(executor)) {
+    DCHECK_NE(executor_, nullptr);
+  }
+
+  Status Init(KernelContext* kernel_ctx, KernelInitArgs args) override {
+    kernel_ = static_cast<const ScalarKernel*>(args.kernel);
+    exec_context_ = kernel_ctx->exec_context();
+    return executor_->Init(kernel_ctx, args);
+  }
+
+  Status Execute(const ExecBatch& batch, ExecListener* listener) override {
+    DCHECK_NE(kernel_, nullptr);
+
+    // Preserve the scalar executor's zero-length handling. A selection cannot make a
+    // zero-length input any denser.
+    if (batch.length == 0 || !batch.selection_vector || kernel_->selective_exec) {
+      return executor_->Execute(batch, listener);
+    }
+
+    if (CheckIfAllScalar(batch)) {
+      // The result is scalar regardless of the selection, so gathering and scattering
+      // would only box an otherwise scalar result.
+      ExecBatch input = batch;
+      input.selection_vector = nullptr;
+      return executor_->Execute(input, listener);
+    }
+
+    ARROW_ASSIGN_OR_RAISE(
+        std::vector<Datum> values,
+        batch.selection_vector->MakeDenseValues(batch.values, exec_context_));
+    ARROW_ASSIGN_OR_RAISE(
+        ExecBatch input,
+        ExecBatch::Make(std::move(values), batch.selection_vector->length()));
+
+    DatumAccumulator dense_listener;
+    RETURN_NOT_OK(executor_->Execute(input, &dense_listener));
+    Datum dense_result = executor_->WrapResults(input.values, dense_listener.values());
+
+    ARROW_ASSIGN_OR_RAISE(
+        Datum result,
+        batch.selection_vector->ScatterDenseResult(dense_result, batch.length,
+                                                   exec_context_));
+    return listener->OnResult(std::move(result));
+  }
+
+  Datum WrapResults(const std::vector<Datum>& args,
+                    const std::vector<Datum>& outputs) override {
+    return executor_->WrapResults(args, outputs);
+  }
+
+  Status CheckResultType(const Datum& out, const char* function_name) override {
+    return executor_->CheckResultType(out, function_name);
+  }
+
+ private:
+  std::unique_ptr<KernelExecutor> executor_;
+  const ScalarKernel* kernel_ = nullptr;
+  ExecContext* exec_context_ = nullptr;
+};
+
+}  // namespace
+
+std::unique_ptr<KernelExecutor> MakeDenseSelectionExecutor(
+    std::unique_ptr<KernelExecutor> executor) {
+  return std::make_unique<DenseSelectionExecutor>(std::move(executor));
+}
+
+}  // namespace arrow::compute::detail

cpp/src/arrow/compute/exec_test.cc

@@ -194,6 +194,22 @@ TEST(SelectionVector, Validate) {
   }
 }
 
+TEST(SelectionVector, GetSpanForChunkAtInt32Limit) {
+  auto selection =
+      SelectionVectorFromJSON("[" + std::to_string(std::numeric_limits<int32_t>::max()) +
+                              "]");
+  SelectionSpan span;
+  const int64_t chunk_end =
+      static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
+
+  ASSERT_EQ(1, selection->GetSpanForChunk(/*chunk_start=*/0, chunk_end,
+                                          /*selection_position=*/0, &span));
+  const auto* contiguous = std::get_if<ContiguousSpan>(&span);
+  ASSERT_NE(contiguous, nullptr);
+  ASSERT_EQ(std::numeric_limits<int32_t>::max(), contiguous->start_offset);
+  ASSERT_EQ(1, contiguous->length);
+}
+
 TEST(DiscreteSpan, Basics) {
   auto indices = ArrayFromJSON(int32(), "[0, 3, 7]");
   const int32_t* idx = indices->data()->GetValues<int32_t>(1);

cpp/src/arrow/meson.build

@@ -52,6 +52,7 @@ arrow_components = {
             'compute/api_vector.cc',
             'compute/cast.cc',
             'compute/exec.cc',
+            'compute/exec_selection.cc',
             'compute/expression.cc',
             'compute/function.cc',
             'compute/function_internal.cc',