diff --git a/cpp/src/arrow/compute/kernels/vector_run_end_encode.cc b/cpp/src/arrow/compute/kernels/vector_run_end_encode.cc index bc8b25de4ec6..9462ece25a9a 100644 --- a/cpp/src/arrow/compute/kernels/vector_run_end_encode.cc +++ b/cpp/src/arrow/compute/kernels/vector_run_end_encode.cc @@ -17,6 +17,10 @@ #include +#include "arrow/array/array_run_end.h" +#include "arrow/array/builder_base.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/compare.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernel.h" #include "arrow/compute/kernels/common_internal.h" @@ -289,6 +293,68 @@ struct RunEndEncodeExec { return Status::Invalid("Invalid run end type: ", *state->run_end_type); } + template + static Status DoExecNested(KernelContext* ctx, const ArraySpan& input_array, + ExecResult* result) { + using RunEndCType = typename RunEndType::c_type; + const int64_t input_length = input_array.length; + auto run_end_type = TypeTraits::type_singleton(); + RETURN_NOT_OK(ValidateRunEndType(run_end_type, input_length)); + + NumericBuilder run_ends_builder(ctx->memory_pool()); + if (input_length > 0) { + auto input = input_array.ToArray(); + // Avoid merging floating-point values whose representations may differ. Signed + // zeros compare unequal, and NaNs remain in separate runs to preserve payloads. + const auto equal_options = + EqualOptions::Defaults().nans_equal(false).signed_zeros_equal(false); + for (int64_t i = 1; i < input_length; ++i) { + if (!ArrayRangeEquals(*input, *input, i - 1, i, i, equal_options)) { + RETURN_NOT_OK(run_ends_builder.Append(static_cast(i))); + } + } + RETURN_NOT_OK(run_ends_builder.Append(static_cast(input_length))); + } + + ARROW_ASSIGN_OR_RAISE( + auto values_builder, + MakeBuilderExactIndex(input_array.type->GetSharedPtr(), ctx->memory_pool())); + RETURN_NOT_OK(values_builder->Reserve(run_ends_builder.length())); + int64_t run_start = 0; + for (int64_t i = 0; i < run_ends_builder.length(); ++i) { + if (input_array.IsNull(run_start)) { + RETURN_NOT_OK(values_builder->AppendNull()); + } else { + RETURN_NOT_OK(values_builder->AppendArraySlice(input_array, run_start, 1)); + } + run_start = run_ends_builder.GetValue(i); + } + + ARROW_ASSIGN_OR_RAISE(auto run_ends, run_ends_builder.Finish()); + ARROW_ASSIGN_OR_RAISE(auto values, values_builder->Finish()); + ARROW_ASSIGN_OR_RAISE(auto output, + RunEndEncodedArray::Make(input_length, run_ends, values)); + result->value = output->data(); + return Status::OK(); + } + + static Status ExecNested(KernelContext* ctx, const ExecSpan& span, ExecResult* result) { + DCHECK(span.values[0].is_array()); + const auto& input_array = span.values[0].array; + const auto* state = checked_cast(ctx->state()); + switch (state->run_end_type->id()) { + case Type::INT16: + return DoExecNested(ctx, input_array, result); + case Type::INT32: + return DoExecNested(ctx, input_array, result); + case Type::INT64: + return DoExecNested(ctx, input_array, result); + default: + break; + } + return Status::Invalid("Invalid run end type: ", *state->run_end_type); + } + /// \brief The OutputType::Resolver of the "run_end_decode" function. static Result ResolveOutputType( KernelContext* ctx, const std::vector& input_types) { @@ -470,6 +536,56 @@ struct RunEndDecodeExec { return Status::Invalid("Invalid run end type: ", *ree_type->run_end_type()); } + template + static Status DoExecNested(KernelContext* ctx, const ArraySpan& input_array, + ExecResult* result) { + using RunEndCType = typename RunEndType::c_type; + const auto& input_values = arrow::ree_util::ValuesArray(input_array); + + ARROW_ASSIGN_OR_RAISE( + auto output_builder, + MakeBuilderExactIndex(input_values.type->GetSharedPtr(), ctx->memory_pool())); + RETURN_NOT_OK(output_builder->Reserve(input_array.length)); + + const arrow::ree_util::RunEndEncodedArraySpan ree_array_span( + input_array); + if (input_array.length > 0) { + for (auto it = ree_array_span.begin(); !it.is_end(ree_array_span); ++it) { + const int64_t physical_index = it.index_into_array(); + const int64_t run_length = it.run_length(); + if (input_values.IsNull(physical_index)) { + RETURN_NOT_OK(output_builder->AppendNulls(run_length)); + continue; + } + for (int64_t i = 0; i < run_length; ++i) { + RETURN_NOT_OK( + output_builder->AppendArraySlice(input_values, physical_index, 1)); + } + } + } + + ARROW_ASSIGN_OR_RAISE(auto output, output_builder->Finish()); + result->value = output->data(); + return Status::OK(); + } + + static Status ExecNested(KernelContext* ctx, const ExecSpan& span, ExecResult* result) { + DCHECK(span.values[0].is_array()); + const auto& input_array = span.values[0].array; + const auto& ree_type = checked_cast(input_array.type); + switch (ree_type->run_end_type()->id()) { + case Type::INT16: + return DoExecNested(ctx, input_array, result); + case Type::INT32: + return DoExecNested(ctx, input_array, result); + case Type::INT64: + return DoExecNested(ctx, input_array, result); + default: + break; + } + return Status::Invalid("Invalid run end type: ", *ree_type->run_end_type()); + } + /// \brief The OutputType::Resolver of the "run_end_decode" function. static Result ResolveOutputType(KernelContext*, const std::vector& in_types) { @@ -562,6 +678,15 @@ void RegisterVectorRunEndEncode(FunctionRegistry* registry) { DCHECK_OK(function->AddKernel(std::move(kernel))); }; + auto add_nested_kernel = [&function](Type::type type_id) { + auto sig = KernelSignature::Make({InputType(match::SameTypeId(type_id))}, + OutputType(RunEndEncodeExec::ResolveOutputType)); + VectorKernel kernel(sig, RunEndEncodeExec::ExecNested, RunEndEncodeInit); + // A REE has null_count=0, so no need to allocate a validity bitmap for them. + kernel.null_handling = NullHandling::OUTPUT_NOT_NULL; + DCHECK_OK(function->AddKernel(std::move(kernel))); + }; + add_kernel(Type::NA); add_kernel(Type::BOOL); for (const auto& ty : NumericTypes()) { @@ -585,6 +710,13 @@ void RegisterVectorRunEndEncode(FunctionRegistry* registry) { add_kernel(Type::BINARY); add_kernel(Type::LARGE_STRING); add_kernel(Type::LARGE_BINARY); + add_nested_kernel(Type::FIXED_SIZE_LIST); + add_nested_kernel(Type::LIST); + add_nested_kernel(Type::LARGE_LIST); + add_nested_kernel(Type::LIST_VIEW); + add_nested_kernel(Type::LARGE_LIST_VIEW); + add_nested_kernel(Type::MAP); + add_nested_kernel(Type::STRUCT); DCHECK_OK(registry->AddFunction(std::move(function))); } @@ -605,6 +737,17 @@ void RegisterVectorRunEndDecode(FunctionRegistry* registry) { } }; + auto add_nested_kernel = [&function](Type::type type_id) { + for (const auto& run_end_type_id : {Type::INT16, Type::INT32, Type::INT64}) { + auto input_type_matcher = match::RunEndEncoded(match::SameTypeId(run_end_type_id), + match::SameTypeId(type_id)); + auto sig = KernelSignature::Make({InputType(std::move(input_type_matcher))}, + OutputType(RunEndDecodeExec::ResolveOutputType)); + VectorKernel kernel(sig, RunEndDecodeExec::ExecNested); + DCHECK_OK(function->AddKernel(std::move(kernel))); + } + }; + add_kernel(Type::NA); add_kernel(Type::BOOL); for (const auto& ty : NumericTypes()) { @@ -628,6 +771,13 @@ void RegisterVectorRunEndDecode(FunctionRegistry* registry) { add_kernel(Type::BINARY); add_kernel(Type::LARGE_STRING); add_kernel(Type::LARGE_BINARY); + add_nested_kernel(Type::FIXED_SIZE_LIST); + add_nested_kernel(Type::LIST); + add_nested_kernel(Type::LARGE_LIST); + add_nested_kernel(Type::LIST_VIEW); + add_nested_kernel(Type::LARGE_LIST_VIEW); + add_nested_kernel(Type::MAP); + add_nested_kernel(Type::STRUCT); DCHECK_OK(registry->AddFunction(std::move(function))); } diff --git a/cpp/src/arrow/compute/kernels/vector_run_end_encode_test.cc b/cpp/src/arrow/compute/kernels/vector_run_end_encode_test.cc index a78e9fe957b1..68ab4543922f 100644 --- a/cpp/src/arrow/compute/kernels/vector_run_end_encode_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_run_end_encode_test.cc @@ -413,5 +413,220 @@ INSTANTIATE_TEST_SUITE_P(EncodeArrayTests, TestRunEndEncodeDecode, ::testing::Values(int16(), int32(), int64()))); +void AssertNestedRunEndEncodeDecode(const std::shared_ptr& input, + const std::string& expected_run_ends_json, + const std::shared_ptr& expected_values) { + for (const auto& run_end_type : {int16(), int32(), int64()}) { + ARROW_SCOPED_TRACE("run end type = ", *run_end_type); + ASSERT_OK_AND_ASSIGN(Datum encoded_datum, + RunEndEncode(input, RunEndEncodeOptions{run_end_type})); + auto encoded = + std::dynamic_pointer_cast(encoded_datum.make_array()); + + ASSERT_NE(encoded, NULLPTR); + ASSERT_OK(encoded->ValidateFull()); + ASSERT_EQ(encoded->length(), input->length()); + ASSERT_EQ(*encoded->type(), *run_end_encoded(run_end_type, input->type())); + ASSERT_ARRAYS_EQUAL(*encoded->run_ends(), + *ArrayFromJSON(run_end_type, expected_run_ends_json)); + ASSERT_ARRAYS_EQUAL(*encoded->values(), *expected_values); + + ASSERT_OK_AND_ASSIGN(Datum decoded_datum, RunEndDecode(encoded)); + auto decoded = decoded_datum.make_array(); + ASSERT_OK(decoded->ValidateFull()); + ASSERT_ARRAYS_EQUAL(*decoded, *input); + + if (input->length() > 0) { + ASSERT_OK_AND_ASSIGN(Datum decoded_slice_datum, RunEndDecode(encoded->Slice(1))); + auto decoded_slice = decoded_slice_datum.make_array(); + ASSERT_OK(decoded_slice->ValidateFull()); + ASSERT_ARRAYS_EQUAL(*decoded_slice, *input->Slice(1)); + + ASSERT_OK_AND_ASSIGN(Datum decoded_prefix_datum, + RunEndDecode(encoded->Slice(0, encoded->length() - 1))); + auto decoded_prefix = decoded_prefix_datum.make_array(); + ASSERT_OK(decoded_prefix->ValidateFull()); + ASSERT_ARRAYS_EQUAL(*decoded_prefix, *input->Slice(0, input->length() - 1)); + } + } +} + +TEST(TestRunEndEncodeDecodeNested, VariableSizeList) { + auto value_type = list(int32()); + auto input = ArrayFromJSON(value_type, R"([ + [9], [1, 2], [1, 2], [], [], null, null, [null], [null], [3], [3], [4], [9] + ])"); + input = input->Slice(1, 11); + auto expected_values = + ArrayFromJSON(value_type, R"([[1, 2], [], null, [null], [3], [4]])"); + AssertNestedRunEndEncodeDecode(input, "[2, 4, 6, 8, 10, 11]", expected_values); + + AssertNestedRunEndEncodeDecode(ArrayFromJSON(value_type, "[]"), "[]", + ArrayFromJSON(value_type, "[]")); + + auto signed_zeros = ArrayFromJSON(list(float64()), "[[0.0], [-0.0]]"); + AssertNestedRunEndEncodeDecode(signed_zeros, "[1, 2]", signed_zeros); + + auto null_lists = ArrayFromJSON(list(null()), "[[null], [null], [], [null]]"); + auto expected_null_lists = ArrayFromJSON(list(null()), "[[null], [], [null]]"); + AssertNestedRunEndEncodeDecode(null_lists, "[2, 3, 4]", expected_null_lists); + + auto nested_value_type = list(list(int32())); + auto nested_lists = + ArrayFromJSON(nested_value_type, "[[[1], [2]], [[1], [2]], [], null, [[3, null]]]"); + auto expected_nested_lists = + ArrayFromJSON(nested_value_type, "[[[1], [2]], [], null, [[3, null]]]"); + AssertNestedRunEndEncodeDecode(nested_lists, "[2, 3, 4, 5]", expected_nested_lists); +} + +TEST(TestRunEndEncodeDecodeNested, LargeList) { + auto value_type = large_list(int32()); + auto input = ArrayFromJSON(value_type, R"([ + [9], [1, 2], [1, 2], [], [], null, null, [null], [null], [3], [3], [4], [9] + ])"); + input = input->Slice(1, 11); + auto expected_values = + ArrayFromJSON(value_type, R"([[1, 2], [], null, [null], [3], [4]])"); + AssertNestedRunEndEncodeDecode(input, "[2, 4, 6, 8, 10, 11]", expected_values); + + AssertNestedRunEndEncodeDecode(ArrayFromJSON(value_type, "[]"), "[]", + ArrayFromJSON(value_type, "[]")); +} + +template +void AssertListViewRunEndEncodeDecode(const std::shared_ptr& offset_type, + const std::shared_ptr& value_type) { + auto offsets = ArrayFromJSON(offset_type, "[0, 1, 1, 5, 2, 0, 0, 3, 3, 5, 0]"); + auto sizes = ArrayFromJSON(offset_type, "[1, 2, 2, 0, 0, null, null, 2, 2, 1, 1]"); + auto child_values = ArrayFromJSON(int32(), "[9, 1, 2, 3, null, 4]"); + ASSERT_OK_AND_ASSIGN(auto list_view, + ListViewArrayType::FromArrays(*offsets, *sizes, *child_values)); + auto input = list_view->Slice(1, 9); + auto expected_values = ArrayFromJSON(value_type, "[[1, 2], [], null, [3, null], [4]]"); + AssertNestedRunEndEncodeDecode(input, "[2, 4, 6, 8, 9]", expected_values); + + AssertNestedRunEndEncodeDecode(ArrayFromJSON(value_type, "[]"), "[]", + ArrayFromJSON(value_type, "[]")); +} + +TEST(TestRunEndEncodeDecodeNested, ListView) { + AssertListViewRunEndEncodeDecode(int32(), list_view(int32())); + AssertListViewRunEndEncodeDecode(int64(), large_list_view(int32())); +} + +TEST(TestRunEndEncodeDecodeNested, FixedSizeList) { + auto value_type = fixed_size_list(int32(), 2); + auto input = ArrayFromJSON(value_type, R"([ + [9, 9], [1, 2], [1, 2], [null, 2], [null, 2], null, null, + [3, 4], [3, 4], [5, 6], [9, 9] + ])"); + input = input->Slice(1, 9); + auto expected_values = + ArrayFromJSON(value_type, "[[1, 2], [null, 2], null, [3, 4], [5, 6]]"); + AssertNestedRunEndEncodeDecode(input, "[2, 4, 6, 8, 9]", expected_values); + + AssertNestedRunEndEncodeDecode(ArrayFromJSON(value_type, "[]"), "[]", + ArrayFromJSON(value_type, "[]")); +} + +TEST(TestRunEndEncodeDecodeNested, PreservesDictionaryIndexType) { + auto dictionary_type = dictionary(uint8(), utf8()); + auto dictionary_values = ArrayFromJSON(utf8(), R"(["a", "b"])"); + ASSERT_OK_AND_ASSIGN( + auto input_values, + DictionaryArray::FromArrays(dictionary_type, ArrayFromJSON(uint8(), "[0, 0, 1, 0]"), + dictionary_values)); + ASSERT_OK_AND_ASSIGN( + auto input, + ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 2, 3, 4]"), *input_values)); + + ASSERT_OK_AND_ASSIGN( + auto expected_dictionary_values, + DictionaryArray::FromArrays(dictionary_type, ArrayFromJSON(uint8(), "[0, 1, 0]"), + dictionary_values)); + ASSERT_OK_AND_ASSIGN(auto expected_values, + ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 2, 3]"), + *expected_dictionary_values)); + AssertNestedRunEndEncodeDecode(input, "[2, 3, 4]", expected_values); +} + +TEST(TestRunEndEncodeDecodeNested, DecodeWithOffsetInValuesArray) { + auto value_type = list(int32()); + auto values = ArrayFromJSON(value_type, "[[9], [1], [2]]")->Slice(1); + auto expected = ArrayFromJSON(value_type, "[[1], [1], [2], [2], [2]]"); + + for (const auto& run_end_type : {int16(), int32(), int64()}) { + ARROW_SCOPED_TRACE("run end type = ", *run_end_type); + auto run_ends = ArrayFromJSON(run_end_type, "[1, 2, 5]")->Slice(1); + ASSERT_OK_AND_ASSIGN(auto encoded, RunEndEncodedArray::Make(5, run_ends, values)); + ASSERT_OK(encoded->ValidateFull()); + + ASSERT_OK_AND_ASSIGN(Datum decoded_datum, RunEndDecode(encoded)); + auto decoded = decoded_datum.make_array(); + ASSERT_OK(decoded->ValidateFull()); + ASSERT_ARRAYS_EQUAL(*decoded, *expected); + + ASSERT_OK_AND_ASSIGN(Datum decoded_slice_datum, RunEndDecode(encoded->Slice(1, 3))); + auto decoded_slice = decoded_slice_datum.make_array(); + ASSERT_OK(decoded_slice->ValidateFull()); + ASSERT_ARRAYS_EQUAL(*decoded_slice, *expected->Slice(1, 3)); + } +} + +TEST(TestRunEndEncodeDecodeNested, Map) { + auto value_type = map(utf8(), int32(), /*keys_sorted=*/true); + auto input = ArrayFromJSON(value_type, R"([ + [["skip", 9]], + [["a", 1], ["b", 2]], + [["a", 1], ["b", 2]], + [], + [], + null, + null, + [["a", null]], + [["a", null]], + [["c", 3]], + [["skip", 9]] + ])"); + input = input->Slice(1, 9); + auto expected_values = ArrayFromJSON(value_type, R"([ + [["a", 1], ["b", 2]], + [], + null, + [["a", null]], + [["c", 3]] + ])"); + AssertNestedRunEndEncodeDecode(input, "[2, 4, 6, 8, 9]", expected_values); + + AssertNestedRunEndEncodeDecode(ArrayFromJSON(value_type, "[]"), "[]", + ArrayFromJSON(value_type, "[]")); +} + +TEST(TestRunEndEncodeDecodeNested, Struct) { + auto value_type = struct_({field("age", int32()), field("name", utf8())}); + auto input = ArrayFromJSON(value_type, R"([ + {"age": 99, "name": "skip"}, + {"age": 20, "name": "a"}, + {"age": 20, "name": "a"}, + {"age": 20, "name": "b"}, + null, + null, + {"age": null, "name": null}, + {"age": null, "name": null}, + {"age": 99, "name": "skip"} + ])"); + input = input->Slice(1, 7); + auto expected_values = ArrayFromJSON(value_type, R"([ + {"age": 20, "name": "a"}, + {"age": 20, "name": "b"}, + null, + {"age": null, "name": null} + ])"); + AssertNestedRunEndEncodeDecode(input, "[2, 3, 5, 7]", expected_values); + + AssertNestedRunEndEncodeDecode(ArrayFromJSON(value_type, "[]"), "[]", + ArrayFromJSON(value_type, "[]")); +} + } // namespace compute } // namespace arrow