diff --git a/cpp/src/arrow/compute/kernels/scalar_cast_internal.cc b/cpp/src/arrow/compute/kernels/scalar_cast_internal.cc index 0484e0b259c5..88ee26cc9cb1 100644 --- a/cpp/src/arrow/compute/kernels/scalar_cast_internal.cc +++ b/cpp/src/arrow/compute/kernels/scalar_cast_internal.cc @@ -280,10 +280,8 @@ void AddZeroCopyCast(Type::type in_type_id, InputType in_type, OutputType out_ty } static bool CanCastFromDictionary(Type::type type_id) { - /// TODO(GH-43010): add is_binary_view_like() here once array_take - /// can handle string-views return (is_primitive(type_id) || is_base_binary_like(type_id) || - is_fixed_size_binary(type_id)); + is_binary_view_like(type_id) || is_fixed_size_binary(type_id)); } void AddCommonCasts(Type::type out_type_id, OutputType out_ty, CastFunction* func) { diff --git a/cpp/src/arrow/compute/kernels/scalar_cast_test.cc b/cpp/src/arrow/compute/kernels/scalar_cast_test.cc index 4ff58040e05e..ee6906a8a710 100644 --- a/cpp/src/arrow/compute/kernels/scalar_cast_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_cast_test.cc @@ -4368,6 +4368,53 @@ TEST(Cast, FromDictionary) { } } +TEST(Cast, DictionaryDecodeFromViewDictionary) { + for (const auto& value_type : {binary_view(), utf8_view()}) { + ARROW_SCOPED_TRACE(value_type->ToString()); + auto dict_values = ArrayFromJSON( + value_type, R"(["alpha", "long-value-over-inline-limit", "omega"])"); + auto indices = ArrayFromJSON(int8(), "[0, 1, null, 2, 1]"); + ASSERT_OK_AND_ASSIGN(auto dict_arr, + DictionaryArray::FromArrays(dictionary(int8(), value_type), + indices, dict_values)); + auto expected = ArrayFromJSON( + value_type, + R"(["alpha", "long-value-over-inline-limit", null, "omega", "long-value-over-inline-limit"])"); + + ASSERT_OK_AND_ASSIGN(Datum decoded, CallFunction("dictionary_decode", {dict_arr})); + ValidateOutput(decoded); + AssertArraysEqual(*expected, *decoded.make_array(), /*verbose=*/true); + CheckCast(dict_arr, expected); + + auto chunked_dict = std::make_shared( + ArrayVector{dict_arr->Slice(0, 2), dict_arr->Slice(2, 3)}); + ASSERT_OK_AND_ASSIGN(Datum decoded_chunked, + CallFunction("dictionary_decode", {chunked_dict})); + ValidateOutput(decoded_chunked); + AssertChunkedEqual( + *ChunkedArrayFromJSON(value_type, + {R"(["alpha", "long-value-over-inline-limit"])", + R"([null, "omega", "long-value-over-inline-limit"])"}), + *decoded_chunked.chunked_array()); + + auto dict_values_with_null = ArrayFromJSON(value_type, R"(["alpha", null, "omega"])"); + auto indices_with_null_source = ArrayFromJSON(int8(), "[0, 1, null, 2]"); + ASSERT_OK_AND_ASSIGN( + auto dict_arr_with_null, + DictionaryArray::FromArrays(dictionary(int8(), value_type), + indices_with_null_source, dict_values_with_null)); + auto expected_with_null_source = + ArrayFromJSON(value_type, R"(["alpha", null, null, "omega"])"); + + ASSERT_OK_AND_ASSIGN(Datum decoded_with_null_source, + CallFunction("dictionary_decode", {dict_arr_with_null})); + ValidateOutput(decoded_with_null_source); + AssertArraysEqual(*expected_with_null_source, *decoded_with_null_source.make_array(), + /*verbose=*/true); + CheckCast(dict_arr_with_null, expected_with_null_source); + } +} + std::shared_ptr SmallintArrayFromJSON(const std::string& json_data) { auto arr = ArrayFromJSON(int16(), json_data); auto ext_data = arr->data()->Copy(); 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 cca8c7ae745d..8645359f21a4 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_filter_internal.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_filter_internal.cc @@ -919,6 +919,11 @@ Status SparseUnionFilterExec(KernelContext* ctx, const ExecSpan& batch, ExecResu return FilterWithTakeExec(SparseUnionTakeExec, ctx, batch, out); } +Status VarBinaryViewFilterExec(KernelContext* ctx, const ExecSpan& batch, + ExecResult* out) { + return FilterWithTakeExec(VarBinaryViewTakeExec, ctx, batch, out); +} + // ---------------------------------------------------------------------- // Implement Filter metafunction @@ -1094,6 +1099,8 @@ void PopulateFilterKernels(std::vector* out) { {InputType(match::Primitive()), plain_filter, PrimitiveFilterExec}, {InputType(match::BinaryLike()), plain_filter, BinaryFilterExec}, {InputType(match::LargeBinaryLike()), plain_filter, BinaryFilterExec}, + {InputType(Type::BINARY_VIEW), plain_filter, VarBinaryViewFilterExec}, + {InputType(Type::STRING_VIEW), plain_filter, VarBinaryViewFilterExec}, {InputType(null()), plain_filter, NullFilterExec}, {InputType(Type::FIXED_SIZE_BINARY), plain_filter, PrimitiveFilterExec}, {InputType(Type::DECIMAL32), plain_filter, PrimitiveFilterExec}, @@ -1116,6 +1123,8 @@ void PopulateFilterKernels(std::vector* out) { {InputType(match::Primitive()), ree_filter, PrimitiveFilterExec}, {InputType(match::BinaryLike()), ree_filter, BinaryFilterExec}, {InputType(match::LargeBinaryLike()), ree_filter, BinaryFilterExec}, + {InputType(Type::BINARY_VIEW), ree_filter, VarBinaryViewFilterExec}, + {InputType(Type::STRING_VIEW), ree_filter, VarBinaryViewFilterExec}, {InputType(null()), ree_filter, NullFilterExec}, {InputType(Type::FIXED_SIZE_BINARY), ree_filter, PrimitiveFilterExec}, {InputType(Type::DECIMAL32), ree_filter, PrimitiveFilterExec}, diff --git a/cpp/src/arrow/compute/kernels/vector_selection_internal.h b/cpp/src/arrow/compute/kernels/vector_selection_internal.h index 887bf0835412..e5af5420ffa4 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_internal.h +++ b/cpp/src/arrow/compute/kernels/vector_selection_internal.h @@ -75,6 +75,7 @@ Status MapFilterExec(KernelContext*, const ExecSpan&, ExecResult*); Status VarBinaryTakeExec(KernelContext*, const ExecSpan&, ExecResult*); Status LargeVarBinaryTakeExec(KernelContext*, const ExecSpan&, ExecResult*); +Status VarBinaryViewTakeExec(KernelContext*, const ExecSpan&, ExecResult*); Status FixedWidthTakeExec(KernelContext*, const ExecSpan&, ExecResult*); Status ListTakeExec(KernelContext*, const ExecSpan&, ExecResult*); Status LargeListTakeExec(KernelContext*, const ExecSpan&, ExecResult*); diff --git a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc index fedafeb5beaa..d0000e2dd2f6 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc @@ -24,6 +24,7 @@ #include "arrow/array/builder_primitive.h" #include "arrow/array/concatenate.h" +#include "arrow/buffer.h" #include "arrow/buffer_builder.h" #include "arrow/chunked_array.h" #include "arrow/compute/api_vector.h" @@ -40,6 +41,8 @@ #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_writer.h" #include "arrow/util/fixed_width_internal.h" #include "arrow/util/int_util.h" #include "arrow/util/logging_internal.h" @@ -488,6 +491,145 @@ Status FixedWidthTakeExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* namespace { +template +int64_t VarBinaryViewTakeTyped(const ArraySpan& values, const ArraySpan& indices, + BinaryViewType::c_type* out_views, uint8_t* out_validity, + std::vector* used_buffers) { + const auto* source_views = values.GetValues(1); + const auto* index_values = indices.GetValues(1); + const auto* indices_validity = + indices.MayHaveNulls() ? indices.buffers[0].data : nullptr; + int64_t valid_count = 0; + + auto copy_view = [&](int64_t out_i) { + auto view = source_views[static_cast(index_values[out_i])]; + out_views[out_i] = view; + if (!view.is_inline()) { + (*used_buffers)[view.ref.buffer_index] = true; + } + ++valid_count; + }; + + if (!values.MayHaveNulls()) { + int64_t out_i = 0; + VisitBitBlocksVoid( + indices_validity, indices.offset, indices.length, + [&](int64_t i) { + copy_view(i); + ++out_i; + }, + [&]() { out_views[out_i++] = {}; }); + } else { + // Take gathers validity through arbitrary indices. VisitTwoBitBlocksVoid and + // OptionalBitmapAnd cannot combine the source and index bitmaps elementwise. + ::arrow::internal::FirstTimeBitmapWriter writer(out_validity, 0, indices.length); + VisitBitBlocksVoid( + indices_validity, indices.offset, indices.length, + [&](int64_t i) { + if (values.IsValid(static_cast(index_values[i]))) { + copy_view(i); + writer.Set(); + } else { + out_views[i] = {}; + } + writer.Next(); + }, + [&]() { + out_views[writer.position()] = {}; + writer.Next(); + }); + writer.Finish(); + } + return valid_count; +} + +int64_t VarBinaryViewTakeDispatch(const ArraySpan& values, const ArraySpan& indices, + BinaryViewType::c_type* out_views, + uint8_t* out_validity, + std::vector* used_buffers) { + switch (indices.type->byte_width()) { + case 1: + return VarBinaryViewTakeTyped(values, indices, out_views, out_validity, + used_buffers); + case 2: + return VarBinaryViewTakeTyped(values, indices, out_views, out_validity, + used_buffers); + case 4: + return VarBinaryViewTakeTyped(values, indices, out_views, out_validity, + used_buffers); + default: + DCHECK_EQ(indices.type->byte_width(), 8); + return VarBinaryViewTakeTyped(values, indices, out_views, out_validity, + used_buffers); + } +} + +} // namespace + +Status VarBinaryViewTakeExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + const ArraySpan& values = batch[0].array; + const ArraySpan& indices = batch[1].array; + + if (TakeState::Get(ctx).boundscheck) { + RETURN_NOT_OK(CheckIndexBounds(indices, values.length)); + } + + const int64_t out_length = indices.length; + const auto data_buffers = values.GetVariadicBuffers(); + std::vector used_buffers(data_buffers.size(), 0); + + ARROW_ASSIGN_OR_RAISE( + auto views_buf, + AllocateBuffer(out_length * static_cast(sizeof(BinaryViewType::c_type)), + ctx->memory_pool())); + auto* out_views = reinterpret_cast(views_buf->mutable_data()); + + std::shared_ptr validity_buf; + uint8_t* out_validity = nullptr; + if (values.MayHaveNulls()) { + ARROW_ASSIGN_OR_RAISE(validity_buf, AllocateBitmap(out_length, ctx->memory_pool())); + out_validity = validity_buf->mutable_data(); + } else if (indices.MayHaveNulls()) { + // All source values are valid, so reuse or realign the index bitmap. + ARROW_ASSIGN_OR_RAISE(validity_buf, ::arrow::internal::OptionalBitmapAnd( + ctx->memory_pool(), indices.GetBuffer(0), + indices.offset, nullptr, 0, out_length)); + } + + const int64_t valid_count = + VarBinaryViewTakeDispatch(values, indices, out_views, out_validity, &used_buffers); + const int64_t null_count = out_length - valid_count; + // Preserve buffer indices without retaining unused payload allocations. C++ arrays + // require non-null variadic buffer slots, even though the C Data Interface permits + // null pointers for zero-sized buffers. + size_t num_data_buffers = used_buffers.size(); + while (num_data_buffers > 0 && !used_buffers[num_data_buffers - 1]) { + --num_data_buffers; + } + BufferVector buffers; + buffers.reserve(2 + num_data_buffers); + buffers.push_back(null_count == 0 ? nullptr : std::move(validity_buf)); + buffers.push_back(std::move(views_buf)); + + std::shared_ptr empty_buffer; + for (size_t i = 0; i < num_data_buffers; ++i) { + if (used_buffers[i]) { + buffers.push_back(data_buffers[i]); + } else { + if (!empty_buffer) { + empty_buffer = std::make_shared(nullptr, 0); + } + buffers.push_back(empty_buffer); + } + } + + out->value = ArrayData::Make(values.type->GetSharedPtr(), out_length, + std::move(buffers), null_count, /*offset=*/0); + return Status::OK(); +} + +namespace { + // ---------------------------------------------------------------------- // Null take @@ -740,6 +882,8 @@ void PopulateTakeKernels(std::vector* out) { {InputType(match::Primitive()), take_indices, FixedWidthTakeExec}, {InputType(match::BinaryLike()), take_indices, VarBinaryTakeExec}, {InputType(match::LargeBinaryLike()), take_indices, LargeVarBinaryTakeExec}, + {InputType(Type::BINARY_VIEW), take_indices, VarBinaryViewTakeExec}, + {InputType(Type::STRING_VIEW), take_indices, VarBinaryViewTakeExec}, {InputType(match::FixedSizeBinaryLike()), take_indices, FixedWidthTakeExec}, {InputType(null()), take_indices, NullTakeExec}, {InputType(Type::DICTIONARY), take_indices, DictionaryTake}, diff --git a/cpp/src/arrow/compute/kernels/vector_selection_test.cc b/cpp/src/arrow/compute/kernels/vector_selection_test.cc index 5fa2d6824dc1..c17214a7a5a6 100644 --- a/cpp/src/arrow/compute/kernels/vector_selection_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_selection_test.cc @@ -16,6 +16,7 @@ // under the License. #include +#include #include #include #include @@ -23,8 +24,12 @@ #include #include +#include "arrow/array/array_binary.h" #include "arrow/array/builder_nested.h" #include "arrow/array/concatenate.h" +#include "arrow/buffer.h" +#include "arrow/c/abi.h" +#include "arrow/c/bridge.h" #include "arrow/chunked_array.h" #include "arrow/compute/api.h" #include "arrow/compute/kernels/test_util_internal.h" @@ -35,7 +40,9 @@ #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" #include "arrow/testing/util.h" +#include "arrow/util/binary_view_util.h" #include "arrow/util/logging.h" +#include "arrow/util/string.h" namespace arrow { @@ -71,6 +78,42 @@ Result> FilterFromJSON( } } +std::shared_ptr MakeViewArrayWithMultiplePayloadBuffers( + const std::shared_ptr& type) { + auto payload0 = Buffer::FromString("prefix-first-long-value-suffix"); + auto payload1 = Buffer::FromString("header-second-long-value-tail"); + constexpr int32_t first_offset = 7; + constexpr int32_t second_offset = 7; + constexpr int32_t first_size = sizeof("first-long-value") - 1; + constexpr int32_t second_size = sizeof("second-long-value") - 1; + + // Null slots may contain an out-of-line view, but must not retain its payload. + auto null_view = util::ToNonInlineBinaryView(payload0->data() + first_offset, + first_size, 0, first_offset); + std::vector views = { + util::ToInlineBinaryView(""), + util::ToInlineBinaryView("tiny"), + util::ToNonInlineBinaryView(payload0->data() + first_offset, first_size, 0, + first_offset), + null_view, + util::ToNonInlineBinaryView(payload1->data() + second_offset, second_size, 1, + second_offset), + util::ToInlineBinaryView("z")}; + + std::shared_ptr null_bitmap; + BitmapFromVector({true, true, true, false, true, true}, &null_bitmap); + BufferVector data_buffers{payload0, payload1}; + auto views_buffer = Buffer::FromVector(views); + if (type->id() == Type::STRING_VIEW) { + return std::make_shared(type, static_cast(views.size()), + views_buffer, data_buffers, null_bitmap, + /*null_count=*/1); + } + return std::make_shared(type, static_cast(views.size()), + views_buffer, data_buffers, null_bitmap, + /*null_count=*/1); +} + Result> REEncode(const std::shared_ptr& array) { ARROW_ASSIGN_OR_RAISE(auto datum, RunEndEncode(array)); return datum.make_array(); @@ -665,12 +708,65 @@ class TestFilterKernelWithString : public TestFilterKernel { } }; -TYPED_TEST_SUITE(TestFilterKernelWithString, BaseBinaryArrowTypes); +TYPED_TEST_SUITE(TestFilterKernelWithString, BaseBinaryOrBinaryViewLikeArrowTypes); TYPED_TEST(TestFilterKernelWithString, FilterString) { this->AssertFilter(R"(["a", "b", "c"])", "[0, 1, 0]", R"(["b"])"); this->AssertFilter(R"([null, "b", "c"])", "[0, 1, 0]", R"(["b"])"); this->AssertFilter(R"(["a", "b", "c"])", "[null, 1, 0]", R"([null, "b"])"); + this->AssertFilter(R"(["a", "long-value-over-inline-limit", null, "z"])", + "[1, 0, null, 1]", R"(["a", null, "z"])"); + this->AssertFilter(R"(["a", null, "z"])", "[0, 1, 1]", R"([null, "z"])"); +} + +TEST_F(TestFilterKernel, FilterBinaryViewMultiplePayloadBuffers) { + for (const auto& value_type : BinaryViewTypes()) { + ARROW_SCOPED_TRACE(value_type->ToString()); + auto values = MakeViewArrayWithMultiplePayloadBuffers(value_type); + auto filter = ArrayFromJSON(boolean(), "[false, false, true, true, true, false]"); + auto expected = + ArrayFromJSON(value_type, R"(["first-long-value", null, "second-long-value"])"); + this->AssertFilter(values, filter, expected); + + ASSERT_OK_AND_ASSIGN(Datum actual_datum, + CallFunction("array_filter", {values, filter}, &this->drop_)); + auto actual = actual_datum.make_array(); + ValidateOutput(*actual); + AssertArraysEqual(*expected, *actual, /*verbose=*/true); + + ASSERT_OK_AND_ASSIGN(auto ree_filter, REEncode(filter)); + ASSERT_OK_AND_ASSIGN( + Datum ree_actual_datum, + CallFunction("array_filter", {values, ree_filter}, &this->drop_)); + auto ree_actual = ree_actual_datum.make_array(); + ValidateOutput(*ree_actual); + AssertArraysEqual(*expected, *ree_actual, /*verbose=*/true); + } +} + +TEST_F(TestFilterKernel, FilterViewDropsUnusedPayloadBuffers) { + for (const auto& value_type : BinaryViewTypes()) { + auto values = MakeViewArrayWithMultiplePayloadBuffers(value_type)->Slice(1, 4); + auto filter = ArrayFromJSON(boolean(), "[true, false, null, true, true]")->Slice(1); + ASSERT_OK_AND_ASSIGN(auto ree_filter, REEncode(filter)); + for (const auto& options : {drop_, emit_null_}) { + auto expected = + ArrayFromJSON(value_type, options.null_selection_behavior == FilterOptions::DROP + ? R"([null, "second-long-value"])" + : R"([null, null, "second-long-value"])"); + for (const auto& selection : {filter, ree_filter}) { + ASSERT_OK_AND_ASSIGN(auto actual_datum, Filter(values, selection, options)); + auto actual = actual_datum.make_array(); + ValidateOutput(*actual); + AssertArraysEqual(*expected, *actual); + ASSERT_EQ(actual->data()->buffers.size(), 4); + ASSERT_NE(actual->data()->buffers[2], nullptr); + ASSERT_EQ(actual->data()->buffers[2]->size(), 0); + ASSERT_EQ(actual->data()->buffers[2]->parent(), nullptr); + ASSERT_EQ(actual->data()->buffers[3], values->data()->buffers[3]); + } + } + } } TYPED_TEST(TestFilterKernelWithString, FilterDictionary) { @@ -1656,12 +1752,19 @@ class TestTakeKernelWithString : public TestTakeKernelTyped { } }; -TYPED_TEST_SUITE(TestTakeKernelWithString, BaseBinaryArrowTypes); +TYPED_TEST_SUITE(TestTakeKernelWithString, BaseBinaryOrBinaryViewLikeArrowTypes); TYPED_TEST(TestTakeKernelWithString, TakeString) { this->CheckTakeXA(R"(["a", "b", "c"])", "[0, 1, 0]", R"(["a", "b", "a"])"); this->CheckTakeXA(R"([null, "b", "c"])", "[0, 1, 0]", "[null, \"b\", null]"); this->CheckTakeXA(R"(["a", "b", "c"])", "[null, 1, 0]", R"([null, "b", "a"])"); + this->CheckTakeXA( + R"(["", "a", "long-value-over-inline-limit", null, "z"])", "[2, 0, 2, 4, 1]", + R"(["long-value-over-inline-limit", "", "long-value-over-inline-limit", "z", "a"])"); + this->CheckTakeXA(R"(["a", null, "long-value-over-inline-limit"])", "[0, 1, null, 2]", + R"(["a", null, null, "long-value-over-inline-limit"])"); + this->CheckTakeXA("[]", "[]", "[]"); + this->CheckTakeXA(R"(["a", "long-value-over-inline-limit"])", "[]", "[]"); this->TestNoValidityBitmapButUnknownNullCount(R"(["a", "b", "c"])", "[0, 1, 0]"); @@ -1675,6 +1778,139 @@ TYPED_TEST(TestTakeKernelWithString, TakeString) { ASSERT_RAISES(IndexError, TakeCAC(type, {kABC, kABC}, "[4, 10]").Value(&chunked_arr)); } +TEST_F(TestTakeKernel, TakeViewMultiplePayloadBuffers) { + for (const auto& value_type : BinaryViewTypes()) { + ARROW_SCOPED_TRACE(value_type->ToString()); + auto values = MakeViewArrayWithMultiplePayloadBuffers(value_type); + auto indices = ArrayFromJSON(int32(), "[2, 4, 0, 5]"); + auto expected = ArrayFromJSON( + value_type, R"(["first-long-value", "second-long-value", "", "z"])"); + + DoCheckTakeAAA(values, indices, expected); + + ASSERT_OK_AND_ASSIGN(Datum actual_datum, + CallFunction("array_take", {values, indices})); + auto actual = actual_datum.make_array(); + ValidateOutput(*actual); + AssertArraysEqual(*expected, *actual, /*verbose=*/true); + ASSERT_EQ(actual->data()->buffers.size(), values->data()->buffers.size()); + ASSERT_EQ(values->data()->buffers[2].get(), actual->data()->buffers[2].get()); + ASSERT_EQ(values->data()->buffers[3].get(), actual->data()->buffers[3].get()); + + auto sliced_values = values->Slice(1, 4); + auto sliced_indices = ArrayFromJSON(int32(), "[1, 3, 0]"); + auto sliced_expected = + ArrayFromJSON(value_type, R"(["first-long-value", "second-long-value", "tiny"])"); + DoAssertTakeAAA(sliced_values, sliced_indices, sliced_expected); + + auto indices_with_offset = ArrayFromJSON(int32(), "[0, 4, null, 2, 5]")->Slice(1, 3); + auto expected_with_offset = + ArrayFromJSON(value_type, R"(["second-long-value", null, "first-long-value"])"); + DoAssertTakeAAA(values, indices_with_offset, expected_with_offset); + } +} + +TEST_F(TestTakeKernel, TakeViewDropsUnusedPayloadBuffers) { + struct TestCase { + const char* indices; + const char* expected; + bool uses_first; + bool uses_second; + }; + const TestCase cases[] = { + {"[4, 3, 4, null, 1]", + R"(["second-long-value", null, "second-long-value", null, "tiny"])", false, true}, + {"[2, 1, 2]", R"(["first-long-value", "tiny", "first-long-value"])", true, false}, + {"[4, 2]", R"(["second-long-value", "first-long-value"])", true, true}, + {"[0, 1, 5, null, 3]", R"(["", "tiny", "z", null, null])", false, false}, + {"[3, null, 3]", "[null, null, null]", false, false}, + {"[]", "[]", false, false}, + }; + for (const auto& value_type : BinaryViewTypes()) { + for (const auto& test : cases) { + ARROW_SCOPED_TRACE(value_type->ToString(), " indices = ", test.indices); + auto values = MakeViewArrayWithMultiplePayloadBuffers(value_type); + auto indices = ArrayFromJSON(int32(), test.indices); + auto expected = ArrayFromJSON(value_type, test.expected); + std::weak_ptr first = values->data()->buffers[2]; + std::weak_ptr second = values->data()->buffers[3]; + + ASSERT_OK_AND_ASSIGN(auto actual, TakeAAA(*values, *indices)); + ValidateOutput(*actual); + AssertArraysEqual(*expected, *actual); + const auto& buffers = actual->data()->buffers; + ASSERT_EQ(buffers.size(), test.uses_second ? 4 : (test.uses_first ? 3 : 2)); + if (test.uses_first) { + ASSERT_EQ(buffers[2], first.lock()); + } else if (test.uses_second) { + ASSERT_NE(buffers[2], nullptr); + ASSERT_EQ(buffers[2]->size(), 0); + ASSERT_EQ(buffers[2]->parent(), nullptr); + } + if (test.uses_second) { + ASSERT_EQ(buffers[3], second.lock()); + } + values.reset(); + ASSERT_EQ(first.expired(), !test.uses_first); + ASSERT_EQ(second.expired(), !test.uses_second); + + // Empty placeholders must also be safe to export and import. + ArrowArray c_array{}; + ASSERT_OK(ExportArray(*actual, &c_array)); + ASSERT_OK_AND_ASSIGN(auto imported, ImportArray(&c_array, value_type)); + ValidateOutput(*imported); + AssertArraysEqual(*expected, *imported); + } + } +} + +TEST_F(TestTakeKernel, TakeViewValidityFollowsIndices) { + for (const auto& value_type : BinaryViewTypes()) { + for (const auto& index_type : IntTypes()) { + for (bool source_has_nulls : {false, true}) { + for (int64_t offset : {0, 1, 7, 8}) { + ARROW_SCOPED_TRACE(value_type->ToString(), + " indices = ", index_type->ToString(), + " source nulls = ", source_has_nulls, " offset = ", offset); + auto values = ArrayFromJSON( + value_type, + source_has_nulls + ? R"(["unused", "a", null, "long-value-over-inline-limit", "z"])" + : R"(["unused", "a", "b", "long-value-over-inline-limit", "z"])"); + values = values->Slice(1); + // Cross bitmap word boundaries with all-valid, all-null and mixed blocks. + std::vector index_json(offset, "0"); + std::vector expected_json; + const char* source_json[] = {R"("a")", source_has_nulls ? "null" : R"("b")", + R"("long-value-over-inline-limit")", R"("z")"}; + for (int64_t i = 0; i < 193; ++i) { + const int64_t index = (i * 3 + 2) % 4; + const bool is_null = (i >= 64 && i < 128) || (i >= 128 && i % 3 == 0); + index_json.push_back(is_null ? "null" : std::to_string(index)); + expected_json.push_back(is_null ? "null" : source_json[index]); + } + auto indices = ArrayFromJSON( + index_type, "[" + ::arrow::internal::JoinStrings(index_json, ",") + "]"); + // A null index's payload is unspecified and must never be dereferenced. + auto* index_bytes = indices->data()->buffers[1]->mutable_data(); + for (int64_t i = 0; i < indices->length(); ++i) { + if (indices->IsNull(i)) { + std::memset(index_bytes + i * index_type->byte_width(), 0xff, + index_type->byte_width()); + } + } + indices = indices->Slice(offset); + values->data()->null_count = kUnknownNullCount; + indices->data()->null_count = kUnknownNullCount; + auto expected = ArrayFromJSON( + value_type, "[" + ::arrow::internal::JoinStrings(expected_json, ",") + "]"); + DoAssertTakeAAA(values, indices, expected); + } + } + } + } +} + TYPED_TEST(TestTakeKernelWithString, TakeDictionary) { auto dict = R"(["a", "b", "c", "d", "e"])"; this->AssertTakeXADictionary(dict, "[3, 4, 2]", "[0, 1, 0]", "[3, 4, 3]"); @@ -2447,7 +2683,7 @@ class TestDropNullKernelWithString : public TestDropNullKernelTyped { } }; -TYPED_TEST_SUITE(TestDropNullKernelWithString, BaseBinaryArrowTypes); +TYPED_TEST_SUITE(TestDropNullKernelWithString, BaseBinaryOrBinaryViewLikeArrowTypes); TYPED_TEST(TestDropNullKernelWithString, DropNullString) { this->AssertDropNull(R"(["a", "b", "c"])", R"(["a", "b", "c"])"); @@ -2637,6 +2873,23 @@ TEST_F(TestDropNullKernelWithRecordBatch, DropNullRecordBatch) { this->AssertDropNull(schm, R"([])", R"([])"); } +TEST_F(TestDropNullKernelWithRecordBatch, DropNullRecordBatchWithView) { + for (const auto& value_type : BinaryViewTypes()) { + ARROW_SCOPED_TRACE(value_type->ToString()); + auto schm = schema({field("a", value_type), field("b", int32())}); + auto batch_json = R"([ + {"a": "a", "b": 1}, + {"a": null, "b": 2}, + {"a": "c", "b": null}, + {"a": "long-value-over-inline-limit", "b": 4} + ])"; + this->AssertDropNull(schm, batch_json, R"([ + {"a": "a", "b": 1}, + {"a": "long-value-over-inline-limit", "b": 4} + ])"); + } +} + class TestDropNullKernelWithChunkedArray : public TestDropNullKernelTyped { public: TestDropNullKernelWithChunkedArray() @@ -2692,6 +2945,14 @@ TEST_F(TestDropNullKernelWithChunkedArray, DropNullChunkedArray) { this->AssertDropNull(int8(), {"[null]", "[null, null]"}, {"[]"}); this->AssertDropNull(int8(), {"[7]", "[8, 9]"}, {"[7]", "[8, 9]"}); this->AssertDropNull(int8(), {"[]", "[]"}, {"[]", "[]"}); + + for (const auto& value_type : BinaryViewTypes()) { + ARROW_SCOPED_TRACE(value_type->ToString()); + this->AssertDropNull( + value_type, + {R"(["a", null])", R"([null, "long-value-over-inline-limit"])", R"([""])"}, + {R"(["a"])", R"(["long-value-over-inline-limit"])", R"([""])"}); + } } TEST_F(TestDropNullKernelWithChunkedArray, DropNullChunkedArrayWithSlices) { @@ -2812,6 +3073,26 @@ TEST_F(TestDropNullKernelWithTable, DropNullTable) { AssertSchemaEqual(schm, actual->schema()); ASSERT_EQ(actual->num_rows(), 0); } + + for (const auto& value_type : BinaryViewTypes()) { + ARROW_SCOPED_TRACE(value_type->ToString()); + auto view_schm = schema({field("a", value_type), field("b", int32())}); + std::vector table_json = {R"([ + {"a": "a", "b": 1}, + {"a": null, "b": 2} + ])", + R"([ + {"a": "c", "b": null}, + {"a": "long-value-over-inline-limit", "b": 4} + ])"}; + std::vector expected_table_json = {R"([ + {"a": "a", "b": 1} + ])", + R"([ + {"a": "long-value-over-inline-limit", "b": 4} + ])"}; + this->AssertDropNull(view_schm, table_json, expected_table_json); + } } TEST_F(TestDropNullKernelWithTable, DropNullTableWithSlices) { diff --git a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc index 22b78a016d93..21f98cec7d4a 100644 --- a/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_swizzle_test.cc @@ -17,6 +17,7 @@ #include +#include "arrow/array/array_dict.h" #include "arrow/array/concatenate.h" #include "arrow/chunked_array.h" #include "arrow/compute/api_vector.h" @@ -754,4 +755,44 @@ TEST(Scatter, Binary) { } } +TEST(Scatter, BinaryView) { + for (const auto& value_type : BinaryViewTypes()) { + ARROW_SCOPED_TRACE(value_type->ToString()); + { + auto values = R"(["a", "b", "long-value-over-inline-limit", "d"])"; + auto indices = "[2, 0, 3, 1]"; + int64_t max_index = 3; + auto expected = R"(["b", "d", "a", "long-value-over-inline-limit"])"; + TestScatter(value_type, values, indices, max_index, expected); + } + { + auto values = R"(["a", "b", "c"])"; + auto indices = "[2, null, 0]"; + int64_t max_index = 4; + auto expected = R"(["c", null, "a", null, null])"; + TestScatter(value_type, values, indices, max_index, expected); + } + { + auto values = R"(["first", "second"])"; + auto indices = "[0, 0]"; + int64_t max_index = 0; + auto expected = R"(["second"])"; + TestScatter(value_type, values, indices, max_index, expected); + } + { + auto dict_values = ArrayFromJSON( + value_type, R"(["alpha", "long-value-over-inline-limit", "omega"])"); + auto dict_type = dictionary(int8(), value_type); + ASSERT_OK_AND_ASSIGN( + auto values, DictionaryArray::FromArrays( + dict_type, ArrayFromJSON(int8(), "[0, 1, 2]"), dict_values)); + ASSERT_OK_AND_ASSIGN( + auto expected, DictionaryArray::FromArrays( + dict_type, ArrayFromJSON(int8(), "[1, 2, 0]"), dict_values)); + AssertScatterAAA(values, ArrayFromJSON(int32(), "[2, 0, 1]"), /*max_index=*/2, + expected); + } + } +} + } // namespace arrow::compute