This is an automated email from the ASF dual-hosted git repository.
alenka pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new d83af8f749 GH-38770: [C++][Python] RecordBatch.filter() segfaults if
passed a ChunkedArray (#40971)
d83af8f749 is described below
commit d83af8f749ee560c0b04d986ba2912e696e1cd68
Author: Alenka Frim <[email protected]>
AuthorDate: Wed May 8 12:57:10 2024 +0200
GH-38770: [C++][Python] RecordBatch.filter() segfaults if passed a
ChunkedArray (#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: #38770
Authored-by: AlenkaF <[email protected]>
Signed-off-by: AlenkaF <[email protected]>
---
.../kernels/vector_selection_filter_internal.cc | 26 +++++++++++++++++-----
python/pyarrow/tests/test_compute.py | 5 +++++
2 files changed, 25 insertions(+), 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
index d5e5e5ad28..8d43c65668 100644
--- 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
index 17cc546f83..d7dee1ad05 100644
--- 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)
