Skip to content

GH-43851: [C++][Compute] Drop logical nulls from union values across containers #54725

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xgbtxy wants to merge 2 commits into
base: main
Choose a base branch
from
Open

GH-43851: [C++][Compute] Drop logical nulls from union values across containers #54725

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
03105c5
GH-43851: [C++][Compute] Drop logical nulls from union values across …
xgbtxy Oct 20, 2026
23fd210
Fix union drop_null with core-only kernels and extension children
xgbtxy Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 64 additions & 13 deletions cpp/src/arrow/compute/kernels/vector_selection.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
#include "arrow/array/array_nested.h"
#include "arrow/array/builder_primitive.h"
#include "arrow/array/concatenate.h"
#include "arrow/array/util.h"
#include "arrow/buffer_builder.h"
#include "arrow/chunked_array.h"
#include "arrow/compute/api_vector.h"
Expand All @@ -40,13 +41,15 @@
#include "arrow/result.h"
#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
#include "arrow/util/bit_block_counter.h"
#include "arrow/util/bit_run_reader.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/bitmap_ops.h"
#include "arrow/util/bitmap_reader.h"
#include "arrow/util/int_util.h"
#include "arrow/util/logging_internal.h"
#include "arrow/util/union_util.h"

namespace arrow {

Expand All @@ -70,33 +73,71 @@ using TakeState = OptionsWrapper<TakeOptions>;
// ----------------------------------------------------------------------
// DropNull Implementation

std::shared_ptr<arrow::BooleanArray> MakeDropNullFilter(const Array& values) {
int64_t DropNullCount(const Array& values) {
if (is_union(values.type_id())) {
return values.data()->MayHaveLogicalNulls() ? values.ComputeLogicalNullCount() : 0;
}
return values.null_count();
}

Result<std::shared_ptr<Array>> MakeDropNullEmptyArray(
const std::shared_ptr<DataType>& type, ExecContext* ctx) {
if (is_union(type->id())) {
// Unlike MakeEmptyArray, this doesn't require builders for extension children.
return MakeArrayOfNull(type, 0, ctx->memory_pool());
}
return MakeEmptyArray(type, ctx->memory_pool());
}

Result<std::shared_ptr<BooleanArray>> MakeDropNullFilter(const Array& values,
ExecContext* ctx) {
if (is_union(values.type_id())) {
ARROW_ASSIGN_OR_RAISE(auto validity,
AllocateEmptyBitmap(values.length(), ctx->memory_pool()));
const ArraySpan span(*values.data());
if (values.type_id() == Type::SPARSE_UNION) {
union_util::SetLogicalNullBitsSparse(span, validity->mutable_data(), 0,
/*set_on_null=*/false);
} else {
union_util::SetLogicalNullBitsDense(span, validity->mutable_data(), 0,
/*set_on_null=*/false);
}
return std::make_shared<BooleanArray>(values.length(), std::move(validity));
}
auto& bitmap_buffer = values.null_bitmap();
return std::make_shared<BooleanArray>(values.length(), bitmap_buffer, nullptr, 0,
values.offset());
}

Result<Datum> DropNullArray(const std::shared_ptr<Array>& values, ExecContext* ctx) {
if (values->null_count() == 0) {
const auto null_count = DropNullCount(*values);
if (null_count == 0) {
return values;
}
if (values->null_count() == values->length()) {
return MakeEmptyArray(values->type(), ctx->memory_pool());
if (null_count == values->length()) {

@dmitry-chirkov-dremio dmitry-chirkov-dremio Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Scenario: all values are null and the union has a bool8 extension child. MakeEmptyArray fails because it needs a builder for the extension, no? Same for chunked path.

Can we cover both dense and sparse unions with extension children?

@xgbtxy xgbtxy Oct 21, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed: the new bool8-child regression fails on the initial PR head with MakeBuilder: cannot construct builder for type extension<arrow.bool8>.

The same follow-up commit uses MakeArrayOfNull(type, 0) for empty unions across the array, chunked-array and record-batch paths; the table path uses the corrected batch handling. This avoids extension builders and retaining input buffers, while preserving the existing empty construction for non-union types.

Added dense/sparse coverage for all-null Array, ChunkedArray, RecordBatch and Table inputs, plus mixed values, slices and chunking. The full selection suite passes: 171 tests, one existing disabled test.

return MakeDropNullEmptyArray(values->type(), ctx);
}
if (values->type()->id() == Type::type::NA) {
return std::make_shared<NullArray>(0);
}
auto drop_null_filter = Datum{MakeDropNullFilter(*values)};
return Filter(values, drop_null_filter, FilterOptions::Defaults(), ctx);
ARROW_ASSIGN_OR_RAISE(auto drop_null_filter, MakeDropNullFilter(*values, ctx));
return Filter(values, Datum(drop_null_filter), FilterOptions::Defaults(), ctx);
}

Result<Datum> DropNullChunkedArray(const std::shared_ptr<ChunkedArray>& values,
ExecContext* ctx) {
if (values->null_count() == 0) {
auto null_count = values->null_count();
if (is_union(values->type()->id())) {
for (const auto& chunk : values->chunks()) {
null_count += DropNullCount(*chunk);
}
}
if (null_count == 0) {
return values;
}
if (values->null_count() == values->length()) {
return ChunkedArray::MakeEmpty(values->type(), ctx->memory_pool());
if (null_count == values->length()) {
ARROW_ASSIGN_OR_RAISE(auto empty, MakeDropNullEmptyArray(values->type(), ctx));
return std::make_shared<ChunkedArray>(std::move(empty));
}
std::vector<std::shared_ptr<Array>> new_chunks;
for (const auto& chunk : values->chunks()) {
Expand All @@ -113,7 +154,7 @@ Result<Datum> DropNullRecordBatch(const std::shared_ptr<RecordBatch>& batch,
// Compute an upper bound of the final null count
int64_t null_count = 0;
for (const auto& column : batch->columns()) {
null_count += column->null_count();
null_count += DropNullCount(*column);
}
if (null_count == 0) {
return batch;
Expand All @@ -126,15 +167,25 @@ Result<Datum> DropNullRecordBatch(const std::shared_ptr<RecordBatch>& batch,
bit_util::SetBitsTo(dst->mutable_data(), 0, batch->num_rows(), false);
break;
}
if (column->null_bitmap_data()) {
if (is_union(column->type_id())) {
ARROW_ASSIGN_OR_RAISE(auto validity, MakeDropNullFilter(*column, ctx));
::arrow::internal::BitmapAnd(validity->values()->data(), validity->offset(),
dst->data(), 0, column->length(), 0,
dst->mutable_data());
} else if (column->null_bitmap_data()) {
::arrow::internal::BitmapAnd(column->null_bitmap_data(), column->offset(),
dst->data(), 0, column->length(), 0,
dst->mutable_data());
}
}
auto drop_null_filter = std::make_shared<BooleanArray>(batch->num_rows(), dst);
if (drop_null_filter->true_count() == 0) {
return RecordBatch::MakeEmpty(batch->schema(), ctx->memory_pool());
ArrayVector empty_columns;
for (const auto& column : batch->columns()) {
ARROW_ASSIGN_OR_RAISE(auto empty, MakeDropNullEmptyArray(column->type(), ctx));
empty_columns.push_back(std::move(empty));
}
return RecordBatch::Make(batch->schema(), 0, std::move(empty_columns));
}
return Filter(Datum(batch), Datum(drop_null_filter), FilterOptions::Defaults(), ctx);
}
Expand All @@ -147,7 +198,7 @@ Result<Datum> DropNullTable(const std::shared_ptr<Table>& table, ExecContext* ct
int64_t null_count = 0;
for (const auto& col : table->columns()) {
for (const auto& column_chunk : col->chunks()) {
null_count += column_chunk->null_count();
null_count += DropNullCount(*column_chunk);
}
}
if (null_count == 0) {
Expand Down
173 changes: 172 additions & 1 deletion cpp/src/arrow/compute/kernels/vector_selection_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "arrow/chunked_array.h"
#include "arrow/compute/api.h"
#include "arrow/compute/kernels/test_util_internal.h"
#include "arrow/extension/bool8.h"
#include "arrow/scalar.h"
#include "arrow/table.h"
#include "arrow/testing/builder.h"
Expand Down Expand Up @@ -2597,6 +2598,25 @@ TEST_F(TestDropNullKernelWithStruct, DropNullStruct) {

class TestDropNullKernelWithUnion : public TestDropNullKernelTyped<UnionType> {};

namespace {

Result<std::shared_ptr<Array>> MakeBool8UnionArray(
bool is_dense, const std::shared_ptr<Array>& type_ids,
const std::shared_ptr<Array>& value_offsets,
const std::shared_ptr<Array>& bool8_storage,
const std::shared_ptr<Array>& other_values) {
auto bool8_values = ExtensionType::WrapArray(extension::bool8(), bool8_storage);
ArrayVector children{bool8_values, other_values};
if (is_dense) {
return DenseUnionArray::Make(*type_ids, *value_offsets, std::move(children),
std::vector<std::string>{"bool8", "other"});
}
return SparseUnionArray::Make(*type_ids, std::move(children),
std::vector<std::string>{"bool8", "other"});
}

} // namespace

TEST_F(TestDropNullKernelWithUnion, DropNullUnion) {
for (const auto& union_type :
{dense_union({field("a", int32()), field("b", utf8())}, {2, 5}),
Expand All @@ -2610,7 +2630,158 @@ TEST_F(TestDropNullKernelWithUnion, DropNullUnion) {
[2, 111],
[5, null]
])";
CheckDropNull(union_type, union_json, union_json);
auto expected_json = R"([
[2, 222],
[5, "hello"],
[5, "eh"],
[2, 111]
])";
CheckDropNull(union_type, union_json, expected_json);
CheckDropNull(union_type, expected_json, expected_json);
CheckDropNull(union_type, "[[2, null], [5, null]]", "[]");
CheckDropNull(union_type, "[]", "[]");

auto values = ArrayFromJSON(union_type, union_json);
AssertDropNullArrays(values->Slice(1, 5), ArrayFromJSON(union_type, expected_json));
}
}

TEST_F(TestDropNullKernelWithUnion, DropNullChunkedUnion) {
for (const auto& union_type :
{dense_union({field("a", int32()), field("b", utf8())}, {2, 5}),
sparse_union({field("a", int32()), field("b", utf8())}, {2, 5})}) {
auto values = ArrayFromJSON(union_type,
R"([[2, null], [2, 1], [5, null], [5, "x"], [2, null]])");
auto input = std::make_shared<ChunkedArray>(
ArrayVector{values->Slice(0, 1), values->Slice(1, 2), values->Slice(3, 2),
values->Slice(5, 0)});
ASSERT_OK_AND_ASSIGN(auto actual, DropNull(input));
ValidateOutput(actual);
AssertChunkedEqual(
*ChunkedArrayFromJSON(union_type, {R"([[2, 1]])", R"([[5, "x"]])"}),
*actual.chunked_array());

auto all_null =
ChunkedArrayFromJSON(union_type, {R"([[2, null]])", R"([[5, null]])"});
ASSERT_OK_AND_ASSIGN(actual, DropNull(all_null));
ValidateOutput(actual);
ASSERT_EQ(actual.length(), 0);
ASSERT_TRUE(actual.type()->Equals(union_type));
}
}

TEST_F(TestDropNullKernelWithUnion, DropNullBool8ExtensionChild) {
for (bool is_dense : {true, false}) {
auto all_null_type_ids = ArrayFromJSON(int8(), "[0, 0]");
auto all_null_offsets = ArrayFromJSON(int32(), "[0, 1]");
auto all_null_storage = ArrayFromJSON(int8(), "[null, null]");
auto all_null_other = ArrayFromJSON(int32(), is_dense ? "[]" : "[0, 0]");
ASSERT_OK_AND_ASSIGN(
auto all_null, MakeBool8UnionArray(is_dense, all_null_type_ids, all_null_offsets,
all_null_storage, all_null_other));

ASSERT_EQ(all_null->ComputeLogicalNullCount(), all_null->length());
ASSERT_OK_AND_ASSIGN(auto empty_array, DropNull(*all_null));
ValidateOutput(empty_array);
ASSERT_EQ(empty_array->length(), 0);
ASSERT_TRUE(empty_array->type()->Equals(all_null->type()));

auto all_null_chunks = std::make_shared<ChunkedArray>(
ArrayVector{all_null->Slice(0, 1), all_null->Slice(1, 1)});
ASSERT_OK_AND_ASSIGN(auto empty_chunked, DropNull(all_null_chunks));
ValidateOutput(empty_chunked);
ASSERT_EQ(empty_chunked.length(), 0);
ASSERT_TRUE(empty_chunked.type()->Equals(all_null->type()));

auto schm = schema({field("u", all_null->type())});
auto batch = RecordBatch::Make(schm, all_null->length(), {all_null});
ASSERT_OK_AND_ASSIGN(auto empty_batch, DropNull(batch));
ValidateOutput(empty_batch);
ASSERT_EQ(empty_batch.length(), 0);
ASSERT_TRUE(empty_batch.record_batch()->schema()->Equals(*schm));

auto table = Table::Make(schm, {all_null_chunks}, all_null->length());
ASSERT_OK_AND_ASSIGN(auto empty_table, DropNull(table));
ValidateOutput(empty_table);
ASSERT_EQ(empty_table.length(), 0);
ASSERT_TRUE(empty_table.table()->schema()->Equals(*schm));

// The first and last selected bool8 child slots are null. A sliced input
// checks that union offsets are respected while retaining the valid slot.
auto mixed_type_ids = ArrayFromJSON(int8(), "[1, 0, 0, 0]");
auto mixed_offsets = ArrayFromJSON(int32(), "[0, 0, 1, 2]");
auto mixed_storage =
ArrayFromJSON(int8(), is_dense ? "[null, 1, null]" : "[1, null, 1, null]");
auto mixed_other = ArrayFromJSON(int32(), is_dense ? "[42]" : "[42, 42, 42, 42]");
ASSERT_OK_AND_ASSIGN(auto mixed,
MakeBool8UnionArray(is_dense, mixed_type_ids, mixed_offsets,
mixed_storage, mixed_other));

ASSERT_OK_AND_ASSIGN(auto actual, DropNull(*mixed));
ValidateOutput(actual);
ASSERT_EQ(actual->length(), 2);
ASSERT_EQ(actual->ComputeLogicalNullCount(), 0);
const auto& actual_union = static_cast<const UnionArray&>(*actual);
ASSERT_EQ(actual_union.type_code(0), 1);
ASSERT_EQ(actual_union.type_code(1), 0);
AssertArraysEqual(*mixed->Slice(0, 1), *actual->Slice(0, 1),
/*verbose=*/true);
AssertArraysEqual(*mixed->Slice(2, 1), *actual->Slice(1, 1),
/*verbose=*/true);

ASSERT_OK_AND_ASSIGN(auto actual_slice, DropNull(*mixed->Slice(1, 3)));
ValidateOutput(actual_slice);
AssertArraysEqual(*mixed->Slice(2, 1), *actual_slice, /*verbose=*/true);

auto mixed_chunks = std::make_shared<ChunkedArray>(
ArrayVector{mixed->Slice(0, 1), mixed->Slice(1, 3)});
ASSERT_OK_AND_ASSIGN(auto actual_chunked, DropNull(mixed_chunks));
ValidateOutput(actual_chunked);
ASSERT_EQ(actual_chunked.length(), 2);
ASSERT_TRUE(actual_chunked.type()->Equals(mixed->type()));
}
}

TEST_F(TestDropNullKernelWithUnion, DropNullUnionRecordBatchAndTable) {
for (const auto& union_type :
{dense_union({field("a", int32()), field("b", utf8())}, {2, 5}),
sparse_union({field("a", int32()), field("b", utf8())}, {2, 5})}) {
auto schm = schema({field("u", union_type), field("n", int32())});
auto input = RecordBatchFromJSON(schm, R"([
{"u": [2, null], "n": 0},
{"u": [2, 222], "n": 1},
{"u": [5, "hello"], "n": null},
{"u": [5, "eh"], "n": 3},
{"u": [2, null], "n": 4},
{"u": [2, 111], "n": 5},
{"u": [5, null], "n": 6}
])");
auto expected = RecordBatchFromJSON(schm, R"([
{"u": [2, 222], "n": 1},
{"u": [5, "eh"], "n": 3},
{"u": [2, 111], "n": 5}
])");
ASSERT_OK_AND_ASSIGN(auto actual, DropNull(input));
ValidateOutput(actual);
ASSERT_BATCHES_EQUAL(*expected, *actual.record_batch());

// Different chunk boundaries exercise the sliced union columns produced by
// TableBatchReader while combining union and ordinary validity.
auto unions = std::make_shared<ChunkedArray>(
ArrayVector{input->column(0)->Slice(0, 2), input->column(0)->Slice(2)});
auto numbers = std::make_shared<ChunkedArray>(
ArrayVector{input->column(1)->Slice(0, 4), input->column(1)->Slice(4)});
auto table = Table::Make(schm, {unions, numbers});
ASSERT_OK_AND_ASSIGN(auto expected_table, Table::FromRecordBatches({expected}));
ASSERT_OK_AND_ASSIGN(actual, DropNull(table));
ValidateOutput(actual);
AssertTablesEqual(*expected_table, *actual.table(), /*same_chunk_layout=*/false);

// No physical validity bitmap is present when every column is a union.
auto union_only = Table::Make(schema({field("u", union_type)}), {unions});
ASSERT_OK_AND_ASSIGN(actual, DropNull(union_only));
ValidateOutput(actual);
ASSERT_EQ(actual.length(), 4);
}
}

Expand Down
Loading
Loading