From 03105c53722568c35f3ff2b03a6ee2aad3dd293e Mon Sep 17 00:00:00 2001 From: xgbtxy <115146404+xgbtxy@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:06:32 +0800 Subject: [PATCH 1/3] GH-43851: [C++][Compute] Drop logical nulls from union values across containers --- .../arrow/compute/kernels/vector_selection.cc | 46 ++++++++--- .../compute/kernels/vector_selection_test.cc | 81 ++++++++++++++++++- 2 files changed, 116 insertions(+), 11 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_selection.cc b/cpp/src/arrow/compute/kernels/vector_selection.cc index 687e0d933923..1e95f77af85e 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection.cc @@ -29,6 +29,7 @@ #include "arrow/array/concatenate.h" #include "arrow/buffer_builder.h" #include "arrow/chunked_array.h" +#include "arrow/compute/api_scalar.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/kernels/util_internal.h" @@ -40,6 +41,7 @@ #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" @@ -70,32 +72,51 @@ using TakeState = OptionsWrapper; // ---------------------------------------------------------------------- // DropNull Implementation -std::shared_ptr 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> MakeDropNullFilter(const Array& values, + ExecContext* ctx) { + if (is_union(values.type_id())) { + ARROW_ASSIGN_OR_RAISE(auto validity, IsValid(values, ctx)); + return checked_pointer_cast(validity.make_array()); + } auto& bitmap_buffer = values.null_bitmap(); return std::make_shared(values.length(), bitmap_buffer, nullptr, 0, values.offset()); } Result DropNullArray(const std::shared_ptr& 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()) { + if (null_count == values->length()) { return MakeEmptyArray(values->type(), ctx->memory_pool()); } if (values->type()->id() == Type::type::NA) { return std::make_shared(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 DropNullChunkedArray(const std::shared_ptr& 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()) { + if (null_count == values->length()) { return ChunkedArray::MakeEmpty(values->type(), ctx->memory_pool()); } std::vector> new_chunks; @@ -113,7 +134,7 @@ Result DropNullRecordBatch(const std::shared_ptr& 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; @@ -126,7 +147,12 @@ Result DropNullRecordBatch(const std::shared_ptr& 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()); @@ -147,7 +173,7 @@ Result DropNullTable(const std::shared_ptr& 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) { diff --git a/cpp/src/arrow/compute/kernels/vector_selection_test.cc b/cpp/src/arrow/compute/kernels/vector_selection_test.cc index 5e38ac15306f..eb5268dfb4b8 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_test.cc @@ -2610,7 +2610,86 @@ 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( + 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, 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( + ArrayVector{input->column(0)->Slice(0, 2), input->column(0)->Slice(2)}); + auto numbers = std::make_shared( + 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); } } From 23fd2109e000e1618ccbcb49a9118e0aa64d0a7d Mon Sep 17 00:00:00 2001 From: xgbtxy <115146404+xgbtxy@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:18:39 +0800 Subject: [PATCH 2/3] Fix union drop_null with core-only kernels and extension children --- .../arrow/compute/kernels/vector_selection.cc | 37 ++++++-- .../compute/kernels/vector_selection_test.cc | 92 +++++++++++++++++++ cpp/src/arrow/compute/registry_test.cc | 49 ++++++++++ 3 files changed, 172 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/vector_selection.cc b/cpp/src/arrow/compute/kernels/vector_selection.cc index 1e95f77af85e..11c4a40fd629 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection.cc @@ -27,9 +27,9 @@ #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_scalar.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/kernels/util_internal.h" @@ -49,6 +49,7 @@ #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 { @@ -79,11 +80,29 @@ int64_t DropNullCount(const Array& values) { return values.null_count(); } +Result> MakeDropNullEmptyArray( + const std::shared_ptr& 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> MakeDropNullFilter(const Array& values, ExecContext* ctx) { if (is_union(values.type_id())) { - ARROW_ASSIGN_OR_RAISE(auto validity, IsValid(values, ctx)); - return checked_pointer_cast(validity.make_array()); + 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(values.length(), std::move(validity)); } auto& bitmap_buffer = values.null_bitmap(); return std::make_shared(values.length(), bitmap_buffer, nullptr, 0, @@ -96,7 +115,7 @@ Result DropNullArray(const std::shared_ptr& values, ExecContext* c return values; } if (null_count == values->length()) { - return MakeEmptyArray(values->type(), ctx->memory_pool()); + return MakeDropNullEmptyArray(values->type(), ctx); } if (values->type()->id() == Type::type::NA) { return std::make_shared(0); @@ -117,7 +136,8 @@ Result DropNullChunkedArray(const std::shared_ptr& values, return values; } if (null_count == values->length()) { - return ChunkedArray::MakeEmpty(values->type(), ctx->memory_pool()); + ARROW_ASSIGN_OR_RAISE(auto empty, MakeDropNullEmptyArray(values->type(), ctx)); + return std::make_shared(std::move(empty)); } std::vector> new_chunks; for (const auto& chunk : values->chunks()) { @@ -160,7 +180,12 @@ Result DropNullRecordBatch(const std::shared_ptr& batch, } auto drop_null_filter = std::make_shared(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); } diff --git a/cpp/src/arrow/compute/kernels/vector_selection_test.cc b/cpp/src/arrow/compute/kernels/vector_selection_test.cc index eb5268dfb4b8..71872b84c3a4 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_test.cc @@ -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" @@ -2597,6 +2598,25 @@ TEST_F(TestDropNullKernelWithStruct, DropNullStruct) { class TestDropNullKernelWithUnion : public TestDropNullKernelTyped {}; +namespace { + +Result> MakeBool8UnionArray( + bool is_dense, const std::shared_ptr& type_ids, + const std::shared_ptr& value_offsets, + const std::shared_ptr& bool8_storage, + const std::shared_ptr& 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{"bool8", "other"}); + } + return SparseUnionArray::Make(*type_ids, std::move(children), + std::vector{"bool8", "other"}); +} + +} // namespace + TEST_F(TestDropNullKernelWithUnion, DropNullUnion) { for (const auto& union_type : {dense_union({field("a", int32()), field("b", utf8())}, {2, 5}), @@ -2650,6 +2670,78 @@ TEST_F(TestDropNullKernelWithUnion, DropNullChunkedUnion) { } } +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( + 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(*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( + 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}), diff --git a/cpp/src/arrow/compute/registry_test.cc b/cpp/src/arrow/compute/registry_test.cc index 3dc14bcff83e..8fae8d04099c 100644 --- a/cpp/src/arrow/compute/registry_test.cc +++ b/cpp/src/arrow/compute/registry_test.cc @@ -21,11 +21,15 @@ #include +#include "arrow/array/array_nested.h" +#include "arrow/compute/api_vector.h" #include "arrow/compute/function.h" #include "arrow/compute/function_options.h" #include "arrow/compute/registry.h" +#include "arrow/extension/bool8.h" #include "arrow/result.h" #include "arrow/status.h" +#include "arrow/testing/builder.h" #include "arrow/testing/gtest_util.h" #include "arrow/util/macros.h" @@ -47,6 +51,51 @@ TEST(TestRegistry, CreateBuiltInRegistry) { ARROW_UNUSED(registry); } +TEST(TestRegistry, DropNullUnionWithoutIsValid) { + // The core registry intentionally doesn't register optional compute kernels. + ASSERT_RAISES(KeyError, GetFunctionRegistry()->GetFunction("is_valid")); + + std::shared_ptr type_ids; + std::shared_ptr dense_offsets; + std::shared_ptr dense_bool8_storage; + std::shared_ptr dense_other; + std::shared_ptr sparse_bool8_storage; + std::shared_ptr sparse_other; + ArrayFromVector(std::vector{0, 1, 0}, &type_ids); + ArrayFromVector(std::vector{0, 0, 1}, &dense_offsets); + ArrayFromVector(int8(), {false, true}, {0, 1}, &dense_bool8_storage); + ArrayFromVector(std::vector{42}, &dense_other); + ArrayFromVector(int8(), {false, true, true}, {0, 1, 1}, + &sparse_bool8_storage); + ArrayFromVector(std::vector{0, 42, 0}, &sparse_other); + + for (bool is_dense : {true, false}) { + std::shared_ptr values; + if (is_dense) { + auto bool8_values = + ExtensionType::WrapArray(extension::bool8(), dense_bool8_storage); + ASSERT_OK_AND_ASSIGN( + values, + DenseUnionArray::Make(*type_ids, *dense_offsets, {bool8_values, dense_other}, + std::vector{"bool8", "int32"})); + } else { + auto bool8_values = + ExtensionType::WrapArray(extension::bool8(), sparse_bool8_storage); + ASSERT_OK_AND_ASSIGN( + values, SparseUnionArray::Make(*type_ids, {bool8_values, sparse_other}, + std::vector{"bool8", "int32"})); + } + + ASSERT_OK_AND_ASSIGN(auto actual, DropNull(*values)); + ASSERT_EQ(actual->length(), 2); + ASSERT_EQ(actual->ComputeLogicalNullCount(), 0); + ASSERT_TRUE(actual->type()->Equals(values->type())); + const auto& actual_union = static_cast(*actual); + ASSERT_EQ(actual_union.type_code(0), 1); + ASSERT_EQ(actual_union.type_code(1), 0); + } +} + TEST_P(TestRegistry, Basics) { auto registry_factory = std::get<0>(GetParam()); auto registry_ = registry_factory(); From 295917f9dcd75955753395859963732844a9bc27 Mon Sep 17 00:00:00 2001 From: xgbtxy <115146404+xgbtxy@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:12:58 +0800 Subject: [PATCH 3/3] GH-43851: Handle extension sibling columns when dropping all rows --- .../arrow/compute/kernels/vector_selection.cc | 22 ++++++- .../compute/kernels/vector_selection_test.cc | 62 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/kernels/vector_selection.cc b/cpp/src/arrow/compute/kernels/vector_selection.cc index 11c4a40fd629..a34809380f02 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection.cc @@ -80,9 +80,29 @@ int64_t DropNullCount(const Array& values) { return values.null_count(); } +bool HasExtensionType(const DataType& type) { + if (type.id() == Type::EXTENSION) { + return true; + } + if (type.id() == Type::RUN_END_ENCODED) { + // Keep the existing run-end encoded empty-array path unchanged. + return false; + } + if (type.id() == Type::DICTIONARY) { + const auto& dictionary_type = checked_cast(type); + return HasExtensionType(*dictionary_type.value_type()); + } + for (const auto& field : type.fields()) { + if (HasExtensionType(*field->type())) { + return true; + } + } + return false; +} + Result> MakeDropNullEmptyArray( const std::shared_ptr& type, ExecContext* ctx) { - if (is_union(type->id())) { + if (is_union(type->id()) || HasExtensionType(*type)) { // Unlike MakeEmptyArray, this doesn't require builders for extension children. return MakeArrayOfNull(type, 0, ctx->memory_pool()); } diff --git a/cpp/src/arrow/compute/kernels/vector_selection_test.cc b/cpp/src/arrow/compute/kernels/vector_selection_test.cc index 71872b84c3a4..ddeb5c890a74 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_test.cc @@ -2615,6 +2615,27 @@ Result> MakeBool8UnionArray( std::vector{"bool8", "other"}); } +Result, std::shared_ptr>> +MakeBool8ExtensionSiblingColumns(bool is_dense) { + 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]"); + ARROW_ASSIGN_OR_RAISE(auto all_null_union, + MakeBool8UnionArray(is_dense, all_null_type_ids, all_null_offsets, + all_null_storage, all_null_other)); + + auto list_offsets = ArrayFromJSON(int32(), "[0, 0, 0]"); + auto empty_bool8_storage = ArrayFromJSON(int8(), "[]"); + auto empty_bool8_values = + ExtensionType::WrapArray(extension::bool8(), empty_bool8_storage); + ARROW_ASSIGN_OR_RAISE(auto empty_lists, + ListArray::FromArrays(list(extension::bool8()), *list_offsets, + *empty_bool8_values)); + + return std::make_pair(std::move(all_null_union), std::move(empty_lists)); +} + } // namespace TEST_F(TestDropNullKernelWithUnion, DropNullUnion) { @@ -2742,6 +2763,47 @@ TEST_F(TestDropNullKernelWithUnion, DropNullBool8ExtensionChild) { } } +TEST_F(TestDropNullKernelWithUnion, DropNullBool8ExtensionSiblingColumnRecordBatch) { + for (bool is_dense : {true, false}) { + SCOPED_TRACE(is_dense ? "dense union" : "sparse union"); + ASSERT_OK_AND_ASSIGN(auto columns, MakeBool8ExtensionSiblingColumns(is_dense)); + auto schm = schema( + {field("u", columns.first->type()), field("lists", columns.second->type())}); + auto batch = RecordBatch::Make(schm, 2, {columns.first, columns.second}); + ASSERT_OK(batch->ValidateFull()); + ASSERT_EQ(columns.second->null_count(), 0); + ASSERT_EQ(columns.second->length(), 2); + + ASSERT_OK_AND_ASSIGN(auto actual, DropNull(batch)); + ValidateOutput(actual); + ASSERT_EQ(actual.length(), 0); + ASSERT_TRUE(actual.record_batch()->schema()->Equals(*schm)); + ASSERT_TRUE( + actual.record_batch()->column(1)->type()->Equals(list(extension::bool8()))); + } +} + +TEST_F(TestDropNullKernelWithUnion, DropNullBool8ExtensionSiblingColumnTable) { + for (bool is_dense : {true, false}) { + SCOPED_TRACE(is_dense ? "dense union" : "sparse union"); + ASSERT_OK_AND_ASSIGN(auto columns, MakeBool8ExtensionSiblingColumns(is_dense)); + auto schm = schema( + {field("u", columns.first->type()), field("lists", columns.second->type())}); + auto table = Table::Make(schm, {std::make_shared(columns.first), + std::make_shared(columns.second)}); + ASSERT_OK(table->ValidateFull()); + ASSERT_EQ(columns.second->null_count(), 0); + ASSERT_EQ(columns.second->length(), 2); + + ASSERT_OK_AND_ASSIGN(auto actual, DropNull(table)); + ValidateOutput(actual); + ASSERT_EQ(actual.length(), 0); + ASSERT_TRUE(actual.table()->schema()->Equals(*schm)); + ASSERT_TRUE(actual.table()->column(1)->type()->Equals(list(extension::bool8()))); + ASSERT_EQ(actual.table()->column(1)->length(), 0); + } +} + TEST_F(TestDropNullKernelWithUnion, DropNullUnionRecordBatchAndTable) { for (const auto& union_type : {dense_union({field("a", int32()), field("b", utf8())}, {2, 5}),