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..fe8a5de23360 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,9 @@ 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::STRUCT); DCHECK_OK(registry->AddFunction(std::move(function))); } @@ -605,6 +733,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 +767,9 @@ 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::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..32edf12b1ad4 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,156 @@ 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, 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, 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 diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index cb010f4387b6..2181d4bf5a6b 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -2535,6 +2535,56 @@ def test_table_from_pylist(cls): assert table.to_pylist() == data2 +@pytest.mark.parametrize( + ("value_type", "values", "encoded_values"), + [ + ( + pa.struct([pa.field("age", pa.int32())]), + [{"age": 20}, {"age": 20}, {"age": 21}], + [{"age": 20}, {"age": 21}], + ), + ( + pa.list_(pa.int32()), + [[20], [20], [21, 22]], + [[20], [21, 22]], + ), + ( + pa.list_(pa.int32(), 2), + [[20, 21], [20, 21], [22, 23]], + [[20, 21], [22, 23]], + ), + ( + pa.list_(pa.list_(pa.int32())), + [[[20], [21]], [[20], [21]], [[22]]], + [[[20], [21]], [[22]]], + ), + ], +) +def test_table_from_pylist_run_end_encoded_nested( + value_type, values, encoded_values +): + ree_type = pa.run_end_encoded(pa.int16(), value_type) + schema = pa.schema([pa.field("data", ree_type)]) + + table = pa.Table.from_pylist( + [{"data": value} for value in values], schema=schema + ) + + encoded = table.column(0).chunk(0) + assert encoded.type == ree_type + assert encoded.run_ends.equals(pa.array([2, 3], type=pa.int16())) + assert encoded.values.equals(pa.array(encoded_values, type=value_type)) + + expected = pa.array(values, type=value_type) + decoded = pc.run_end_decode(encoded) + assert decoded.type == value_type + assert decoded.equals(expected) + assert pc.run_end_decode(encoded.slice(1)).equals(expected.slice(1)) + assert pc.run_end_decode(encoded.slice(0, len(encoded) - 1)).equals( + expected.slice(0, len(expected) - 1) + ) + + @pytest.mark.pandas def test_table_from_pandas_schema(): # passed schema is source of truth for the columns