apacheGH-38770: [C++][Python] RecordBatch.filter() segfaults if passe…

…d a ChunkedArray (apache#40971) ### Rationale for this change Filtering a record batch with a boolean mask in the form of a `ChunkedArray` results in a segmentation fault. ### What changes are included in this PR? In case chunked array is passed as a mask to filter record batch, the code path for `pa.Table.filter()` is taken resulting in a filtered table. ### Are these changes tested? Yes. ### Are there any user-facing changes? No. * GitHub Issue: apache#38770 Authored-by: AlenkaF <[email protected]> Signed-off-by: AlenkaF <[email protected]>
obin1000 · May 8, 2024 · d83af8f · d83af8f
1 parent 5385926
commit d83af8f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/cpp/src/arrow/compute/kernels/vector_selection_filter_internal.cc b/cpp/src/arrow/compute/kernels/vector_selection_filter_internal.cc
@@ -22,6 +22,7 @@
 #include <type_traits>
 #include <vector>
 
+#include "arrow/array/concatenate.h"
 #include "arrow/array/data.h"
 #include "arrow/buffer_builder.h"
 #include "arrow/chunked_array.h"
@@ -928,12 +929,26 @@ Result<std::shared_ptr<RecordBatch>> FilterRecordBatch(const RecordBatch& batch,
     return Status::Invalid("Filter inputs must all be the same length");
   }
 
-  // Convert filter to selection vector/indices and use Take
+  // Fetch filter
   const auto& filter_opts = *static_cast<const FilterOptions*>(options);
-  ARROW_ASSIGN_OR_RAISE(
-      std::shared_ptr<ArrayData> indices,
-      GetTakeIndices(*filter.array(), filter_opts.null_selection_behavior,
-                     ctx->memory_pool()));
+  ArrayData filter_array;
+  switch (filter.kind()) {
+    case Datum::ARRAY:
+      filter_array = *filter.array();
+      break;
+    case Datum::CHUNKED_ARRAY: {
+      ARROW_ASSIGN_OR_RAISE(auto combined, Concatenate(filter.chunked_array()->chunks()));
+      filter_array = *combined->data();
+      break;
+    }
+    default:
+      return Status::TypeError("Filter should be array-like");
+  }
+
+  // Convert filter to selection vector/indices and use Take
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrayData> indices,
+                        GetTakeIndices(filter_array, filter_opts.null_selection_behavior,
+                                       ctx->memory_pool()));
   std::vector<std::shared_ptr<Array>> columns(batch.num_columns());
   for (int i = 0; i < batch.num_columns(); ++i) {
     ARROW_ASSIGN_OR_RAISE(Datum out, Take(batch.column(i)->data(), Datum(indices),
@@ -1042,7 +1057,6 @@ class FilterMetaFunction : public MetaFunction {
     }
 
     if (args[0].kind() == Datum::RECORD_BATCH) {
-      auto values_batch = args[0].record_batch();
       ARROW_ASSIGN_OR_RAISE(
           std::shared_ptr<RecordBatch> out_batch,
           FilterRecordBatch(*args[0].record_batch(), args[1], options, ctx));

diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py
@@ -1345,6 +1345,11 @@ def test_filter_record_batch():
     expected = pa.record_batch([pa.array(["a", "e"])], names=["a'"])
     assert result.equals(expected)
 
+    # GH-38770: mask is chunked array
+    chunked_mask = pa.chunked_array([[True, False], [None], [False, True]])
+    result = batch.filter(chunked_mask)
+    assert result.equals(expected)
+
     result = batch.filter(mask, null_selection_behavior="emit_null")
     expected = pa.record_batch([pa.array(["a", None, "e"])], names=["a'"])
     assert result.equals(expected)