GH-32107: [C++] Create Filter Kernels for REE values

cpp/src/arrow/CMakeLists.txt

@@ -461,6 +461,7 @@ if(ARROW_COMPUTE)
        compute/kernels/vector_rank.cc
        compute/kernels/vector_replace.cc
        compute/kernels/vector_run_end_encode.cc
+       compute/kernels/vector_run_end_selection.cc
        compute/kernels/vector_select_k.cc
        compute/kernels/vector_sort.cc)
 

cpp/src/arrow/compute/kernels/vector_run_end_encode.cc

@@ -245,10 +245,9 @@ ARROW_NOINLINE Status RunEndEncodeNullArray(const std::shared_ptr<DataType>& run
   RETURN_NOT_OK(ValidateRunEndType(run_end_type, input_array.length));
 
   ARROW_ASSIGN_OR_RAISE(
-      auto output_array_data,
+      output->value,
       ree_util::MakeNullREEArray(run_end_type, input_length, ctx->memory_pool()));
 
-  output->value = std::move(output_array_data);
   return Status::OK();
 }
 

cpp/src/arrow/compute/kernels/vector_run_end_selection.h

@@ -0,0 +1,40 @@
+// 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.
+
+#pragma once
+
+#include <cstdint>
+
+#include "arrow/array/data.h"
+#include "arrow/compute/api_vector.h"
+#include "arrow/result.h"
+#include "arrow/util/bit_util.h"
+#include "arrow/util/ree_util.h"
+
+// Filtering from and using run-end encoded filter arrays.
+//
+// Used by vector_selection.cc to implement the actual selection compute kernels.
+
+namespace arrow::compute::internal {
+
+ARROW_EXPORT Status REExREEFilterExec(KernelContext* ctx, const ExecSpan& span,
+                                      ExecResult* result);
+
+ARROW_EXPORT Status REExPlainFilterExec(KernelContext* ctx, const ExecSpan& span,
+                                        ExecResult* result);
+
+}  // namespace arrow::compute::internal

cpp/src/arrow/compute/kernels/vector_selection.cc

@@ -44,6 +44,7 @@
 #include "arrow/util/bit_util.h"
 #include "arrow/util/bitmap_ops.h"
 #include "arrow/util/bitmap_reader.h"
+#include "arrow/util/checked_cast.h"
 #include "arrow/util/int_util.h"
 
 namespace arrow {

cpp/src/arrow/compute/kernels/vector_selection_filter_internal.cc

@@ -29,6 +29,7 @@
 #include "arrow/compute/exec.h"
 #include "arrow/compute/kernel.h"
 #include "arrow/compute/kernels/codegen_internal.h"
+#include "arrow/compute/kernels/vector_run_end_selection.h"
 #include "arrow/compute/kernels/vector_selection_filter_internal.h"
 #include "arrow/compute/kernels/vector_selection_internal.h"
 #include "arrow/datum.h"
@@ -40,6 +41,7 @@
 #include "arrow/util/bit_run_reader.h"
 #include "arrow/util/bit_util.h"
 #include "arrow/util/bitmap_ops.h"
+#include "arrow/util/ree_util.h"
 
 namespace arrow {
 
@@ -1040,9 +1042,12 @@ std::unique_ptr<Function> MakeFilterMetaFunction() {
 void PopulateFilterKernels(std::vector<SelectionKernelData>* out) {
   auto plain_filter = InputType(Type::BOOL);
   auto ree_filter = InputType(match::RunEndEncoded(Type::BOOL));
+  auto ree_values = [](auto in) {
+    return InputType(match::RunEndEncoded(std::move(in)));
+  };
 
   *out = {
-      // * x Boolean
+      // Plain(*) x Plain(Boolean)
       {InputType(match::Primitive()), plain_filter, PrimitiveFilterExec},
       {InputType(match::BinaryLike()), plain_filter, BinaryFilterExec},
       {InputType(match::LargeBinaryLike()), plain_filter, BinaryFilterExec},
@@ -1060,7 +1065,7 @@ void PopulateFilterKernels(std::vector<SelectionKernelData>* out) {
       {InputType(Type::STRUCT), plain_filter, StructFilterExec},
       {InputType(Type::MAP), plain_filter, MapFilterExec},
 
-      // * x REE(Boolean)
+      // Plain(*) x REE(Boolean)
       {InputType(match::Primitive()), ree_filter, PrimitiveFilterExec},
       {InputType(match::BinaryLike()), ree_filter, BinaryFilterExec},
       {InputType(match::LargeBinaryLike()), ree_filter, BinaryFilterExec},
@@ -1077,6 +1082,24 @@ void PopulateFilterKernels(std::vector<SelectionKernelData>* out) {
       {InputType(Type::SPARSE_UNION), ree_filter, SparseUnionFilterExec},
       {InputType(Type::STRUCT), ree_filter, StructFilterExec},
       {InputType(Type::MAP), ree_filter, MapFilterExec},
+
+      // REE(*) x REE(Boolean)
+      {ree_values(match::Primitive()), ree_filter, REExREEFilterExec},
+      {ree_values(match::BinaryLike()), ree_filter, REExREEFilterExec},
+      {ree_values(match::LargeBinaryLike()), ree_filter, REExREEFilterExec},
+      {ree_values(Type::FIXED_SIZE_BINARY), ree_filter, REExREEFilterExec},
+      {ree_values(Type::NA), ree_filter, REExREEFilterExec},
+      {ree_values(Type::DECIMAL128), ree_filter, REExREEFilterExec},
+      {ree_values(Type::DECIMAL256), ree_filter, REExREEFilterExec},
+
+      // REE(*) x Plain(Boolean)
+      {ree_values(match::Primitive()), plain_filter, REExPlainFilterExec},
+      {ree_values(match::BinaryLike()), plain_filter, REExPlainFilterExec},
+      {ree_values(match::LargeBinaryLike()), plain_filter, REExPlainFilterExec},
+      {ree_values(Type::FIXED_SIZE_BINARY), plain_filter, REExPlainFilterExec},
+      {ree_values(Type::NA), plain_filter, REExPlainFilterExec},
+      {ree_values(Type::DECIMAL128), plain_filter, REExPlainFilterExec},
+      {ree_values(Type::DECIMAL256), plain_filter, REExPlainFilterExec},
   };
 }
 

cpp/src/arrow/compute/kernels/vector_selection_internal.cc

@@ -82,6 +82,37 @@ Status PreallocatePrimitiveArrayData(KernelContext* ctx, int64_t length, int bit
   return Status::OK();
 }
 
+Status PreallocateDataREE(KernelContext* ctx, int64_t physical_length, int bit_width,
+                          bool allocate_validity, ArrayData* out) {
+  // Preallocate memory
+  out->buffers = {NULLPTR};
+  out->child_data = {NULLPTR, NULLPTR};
+
+  auto& ree_type = checked_cast<RunEndEncodedType&>(*out->type);
+  auto values_array = std::make_shared<ArrayData>(ree_type.value_type(), physical_length);
+  values_array->buffers = {NULLPTR, NULLPTR};
+  auto run_ends_array = std::make_shared<ArrayData>(ree_type.run_end_type(),
+                                                    physical_length, /*null_count=*/0);
+  run_ends_array->buffers = {NULLPTR, NULLPTR};
+
+  if (allocate_validity) {
+    ARROW_ASSIGN_OR_RAISE(values_array->buffers[0], ctx->AllocateBitmap(physical_length));
+  }
+  if (bit_width == 1) {
+    ARROW_ASSIGN_OR_RAISE(values_array->buffers[1], ctx->AllocateBitmap(physical_length));
+  } else {
+    ARROW_ASSIGN_OR_RAISE(values_array->buffers[1],
+                          ctx->Allocate(physical_length * bit_width / 8));
+  }
+  ARROW_ASSIGN_OR_RAISE(
+      run_ends_array->buffers[1],
+      ctx->Allocate(physical_length * ree_type.run_end_type()->bit_width() / 8));
+
+  out->child_data[0] = std::move(run_ends_array);
+  out->child_data[1] = std::move(values_array);
+  return Status::OK();
+}
+
 namespace {
 
 /// \brief Iterate over a REE filter, emitting ranges of a plain values array that

cpp/src/arrow/compute/kernels/vector_selection_internal.h

@@ -51,6 +51,9 @@ void RegisterSelectionFunction(const std::string& name, FunctionDoc doc,
 Status PreallocatePrimitiveArrayData(KernelContext* ctx, int64_t length, int bit_width,
                                      bool allocate_validity, ArrayData* out);
 
+Status PreallocatePrimitiveArrayData(KernelContext* ctx, int64_t length, int bit_width,
+                                     bool allocate_validity, ArrayData* out);
+
 /// \brief Callback type for VisitPlainxREEFilterOutputSegments.
 ///
 /// position is the logical position in the values array relative to its offset.