Implement SelectionVector::FromMask

zanmato1984 · Oct 21, 2024 · 8e2368c · 8e2368c
1 parent 7825c12
commit 8e2368c
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 8 deletions.
diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc
@@ -27,6 +27,7 @@
 
 #include "arrow/array/array_base.h"
 #include "arrow/array/array_primitive.h"
+#include "arrow/array/builder_primitive.h"
 #include "arrow/array/data.h"
 #include "arrow/array/util.h"
 #include "arrow/buffer.h"
@@ -50,6 +51,7 @@
 #include "arrow/util/logging.h"
 #include "arrow/util/thread_pool.h"
 #include "arrow/util/vector.h"
+#include "arrow/visit_data_inline.h"
 
 namespace arrow {
 
@@ -1355,8 +1357,27 @@ SelectionVector::SelectionVector(const Array& arr) : SelectionVector(arr.data())
 int32_t SelectionVector::length() const { return static_cast<int32_t>(data_->length); }
 
 Result<std::shared_ptr<SelectionVector>> SelectionVector::FromMask(
-    const BooleanArray& arr) {
-  return Status::NotImplemented("FromMask");
+    const BooleanArray& arr, MemoryPool* pool) {
+  auto length = arr.null_count();
+  Int32Builder builder(pool);
+  RETURN_NOT_OK(builder.Reserve(length));
+  ArraySpan span(*arr.data());
+  int32_t i = 0;
+  RETURN_NOT_OK(VisitArraySpanInline<Int32Type>(
+      span,
+      [&](bool mask) -> Status {
+        if (mask) {
+          RETURN_NOT_OK(builder.Append(i));
+        }
+        ++i;
+        return Status::OK();
+      },
+      [&]() {
+        ++i;
+        return Status::OK();
+      }));
+  ARROW_ASSIGN_OR_RAISE(auto indices, builder.Finish());
+  return std::make_shared<SelectionVector>(indices->data());
 }
 
 Result<Datum> CallFunction(const std::string& func_name, const std::vector<Datum>& args,

diff --git a/cpp/src/arrow/compute/exec.h b/cpp/src/arrow/compute/exec.h
@@ -141,7 +141,8 @@ class ARROW_EXPORT SelectionVector {
   explicit SelectionVector(const Array& arr);
 
   /// \brief Create SelectionVector from boolean mask
-  static Result<std::shared_ptr<SelectionVector>> FromMask(const BooleanArray& arr);
+  static Result<std::shared_ptr<SelectionVector>> FromMask(
+      const BooleanArray& arr, MemoryPool* pool = default_memory_pool());
 
   const int32_t* indices() const { return indices_; }
   int32_t length() const;

diff --git a/cpp/src/arrow/compute/special_form.cc b/cpp/src/arrow/compute/special_form.cc
@@ -115,16 +115,16 @@ Result<Datum> IfElseSpecialForm::Execute(const std::vector<Expression>& argument
   if (IsSelectionVectorAwarePathAvailable({if_true_expr, if_false_expr}, input,
                                           exec_context)) {
     DCHECK(cond.is_array());
-    ARROW_ASSIGN_OR_RAISE(auto sel_true, SelectionVector::FromMask(cond.array()));
+    auto boolean_cond = cond.array_as<BooleanArray>();
+    ARROW_ASSIGN_OR_RAISE(auto sel_true, SelectionVector::FromMask(*boolean_cond));
     ARROW_ASSIGN_OR_RAISE(auto cond_inverted,
                           CallFunction("invert", {cond}, exec_context));
     DCHECK(cond_inverted.is_array());
-    ARROW_ASSIGN_OR_RAISE(auto sel_false,
-                          SelectionVector::FromMask(cond_inverted.array()));
+    ARROW_ASSIGN_OR_RAISE(auto sel_false, SelectionVector::FromMask(*boolean_cond));
     ExecBatch input_true = input;
-    input_true.selection_vector = std::make_shared<SelectionVector>(sel_true.array());
+    input_true.selection_vector = sel_true;
     ExecBatch input_false = input;
-    input_false.selection_vector = std::make_shared<SelectionVector>(sel_false.array());
+    input_false.selection_vector = sel_false;
     ARROW_ASSIGN_OR_RAISE(
         auto if_true, ExecuteScalarExpression(if_true_expr, input_true, exec_context));
     ARROW_ASSIGN_OR_RAISE(