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 9462ece25a9..c03b960d438 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,7 @@ #include +#include "arrow/array/array_dict.h" #include "arrow/array/array_run_end.h" #include "arrow/array/builder_base.h" #include "arrow/array/builder_primitive.h" @@ -573,6 +574,24 @@ struct RunEndDecodeExec { DCHECK(span.values[0].is_array()); const auto& input_array = span.values[0].array; const auto& ree_type = checked_cast(input_array.type); + if (ree_type->value_type()->id() == Type::DICTIONARY) { + // Expand dictionary indices, retaining the original dictionary and index type. + // The existing primitive decoder handles run ends, slices, and null indices. + auto encoded = std::static_pointer_cast(input_array.ToArray()); + auto values = std::static_pointer_cast(encoded->values()); + ARROW_ASSIGN_OR_RAISE(auto encoded_indices, + RunEndEncodedArray::Make(encoded->length(), + encoded->run_ends(), values->indices(), + encoded->offset())); + ARROW_ASSIGN_OR_RAISE(auto decoded_indices, + RunEndDecode(encoded_indices, ctx->exec_context())); + ARROW_ASSIGN_OR_RAISE(auto decoded, + DictionaryArray::FromArrays(values->type(), + decoded_indices.make_array(), + values->dictionary())); + result->value = decoded->data(); + return Status::OK(); + } switch (ree_type->run_end_type()->id()) { case Type::INT16: return DoExecNested(ctx, input_array, result); @@ -778,6 +797,7 @@ void RegisterVectorRunEndDecode(FunctionRegistry* registry) { add_nested_kernel(Type::LARGE_LIST_VIEW); add_nested_kernel(Type::MAP); add_nested_kernel(Type::STRUCT); + add_nested_kernel(Type::DICTIONARY); 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 68ab4543922..d619391099f 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 @@ -550,6 +550,62 @@ TEST(TestRunEndEncodeDecodeNested, PreservesDictionaryIndexType) { AssertNestedRunEndEncodeDecode(input, "[2, 3, 4]", expected_values); } +TEST(TestRunEndEncodeDecodeNested, DictionaryValues) { + auto dict_values = ArrayFromJSON(utf8(), R"(["alpha", "beta", "unused"])"); + + for (const auto& run_end_type : {int16(), int32(), int64()}) { + for (const auto& index_type : {int8(), uint8(), int64()}) { + ARROW_SCOPED_TRACE("run end type = ", *run_end_type, ", index type = ", + *index_type); + auto dict_type = dictionary(index_type, utf8(), /*ordered=*/true); + ASSERT_OK_AND_ASSIGN( + auto values, + DictionaryArray::FromArrays( + dict_type, ArrayFromJSON(index_type, "[1, null, 0, 1]"), dict_values)); + auto run_ends = ArrayFromJSON(run_end_type, "[2, 4, 5, 8]"); + ASSERT_OK_AND_ASSIGN(auto encoded, + RunEndEncodedArray::Make(8, run_ends, values)); + + auto expected_indices = + ArrayFromJSON(index_type, "[1, 1, null, null, 0, 1, 1, 1]"); + ASSERT_OK_AND_ASSIGN( + auto expected, + DictionaryArray::FromArrays(dict_type, expected_indices, dict_values)); + ASSERT_OK_AND_ASSIGN(Datum decoded_datum, RunEndDecode(encoded)); + auto decoded = decoded_datum.make_array(); + ASSERT_OK(decoded->ValidateFull()); + ASSERT_ARRAYS_EQUAL(*decoded, *expected); + auto decoded_dict = std::static_pointer_cast(decoded); + ASSERT_ARRAYS_EQUAL(*decoded_dict->indices(), *expected_indices); + ASSERT_ARRAYS_EQUAL(*decoded_dict->dictionary(), *dict_values); + ASSERT_EQ(decoded_dict->type()->id(), Type::DICTIONARY); + ASSERT_TRUE(decoded_dict->dict_type()->ordered()); + + // Test a slice inside runs, and decoding two chunks independently. + auto expected_slice = expected->Slice(1, 6); + ASSERT_OK_AND_ASSIGN(Datum slice_datum, RunEndDecode(encoded->Slice(1, 6))); + ASSERT_ARRAYS_EQUAL(*slice_datum.make_array(), *expected_slice); + auto chunked = std::make_shared( + std::vector>{encoded, encoded->Slice(1, 6)}); + ASSERT_OK_AND_ASSIGN(Datum chunked_datum, RunEndDecode(chunked)); + ASSERT_EQ(chunked_datum.chunked_array()->num_chunks(), 2); + ASSERT_ARRAYS_EQUAL(*chunked_datum.chunked_array()->chunk(0), *expected); + ASSERT_ARRAYS_EQUAL(*chunked_datum.chunked_array()->chunk(1), *expected_slice); + + // An empty input still preserves the dictionary. + ASSERT_OK_AND_ASSIGN( + auto empty_values, + DictionaryArray::FromArrays(dict_type, ArrayFromJSON(index_type, "[]"), + dict_values)); + ASSERT_OK_AND_ASSIGN(auto empty_encoded, + RunEndEncodedArray::Make( + 0, ArrayFromJSON(run_end_type, "[]"), empty_values)); + ASSERT_OK_AND_ASSIGN(Datum empty_decoded, RunEndDecode(empty_encoded)); + ASSERT_ARRAYS_EQUAL(*empty_decoded.make_array(), *empty_values); + } + } +} + TEST(TestRunEndEncodeDecodeNested, DecodeWithOffsetInValuesArray) { auto value_type = list(int32()); auto values = ArrayFromJSON(value_type, "[[9], [1], [2]]")->Slice(1); diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 63c0eb0bea4..634a714907d 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -4339,6 +4339,25 @@ def test_run_end_encoded_to_pandas(): ree_array.to_pandas(zero_copy_only=True) +@pytest.mark.pandas +@pytest.mark.parametrize('run_end_type', [pa.int16(), pa.int32(), pa.int64()]) +def test_run_end_encoded_dictionary_to_pandas(run_end_type): + source = pa.array(['a', 'a', None, None, 'b', 'a', 'a', 'a']).dictionary_encode() + runs = pc.run_end_encode(source.indices, options=pc.RunEndEncodeOptions(run_end_type)) + values = pa.DictionaryArray.from_arrays(runs.values, source.dictionary) + ree = pa.RunEndEncodedArray.from_arrays(runs.run_ends, values) + + decoded = pc.run_end_decode(ree) + assert decoded.type == source.type + assert decoded.indices.equals(source.indices) + assert decoded.dictionary.equals(source.dictionary) + + import pandas as pd + result = pa.table({'c': ree}).to_pandas() + expected = pa.table({'c': source}).to_pandas() + pd.testing.assert_frame_equal(result, expected) + + @pytest.mark.parametrize(('list_array_type', 'list_type_factory'), [(pa.ListViewArray, pa.list_view), (pa.LargeListViewArray, pa.large_list_view)])