GH-44183: [C++][Python] Support run-end encoded struct and list values

cpp/src/arrow/compute/kernels/vector_run_end_encode.cc

@@ -17,6 +17,10 @@
 
 #include <utility>
 
+#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 <typename RunEndType>
+  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<RunEndType>::type_singleton();
+    RETURN_NOT_OK(ValidateRunEndType(run_end_type, input_length));
+
+    NumericBuilder<RunEndType> 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<RunEndCType>(i)));
+        }
+      }
+      RETURN_NOT_OK(run_ends_builder.Append(static_cast<RunEndCType>(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<const RunEndEncodingState*>(ctx->state());
+    switch (state->run_end_type->id()) {
+      case Type::INT16:
+        return DoExecNested<Int16Type>(ctx, input_array, result);
+      case Type::INT32:
+        return DoExecNested<Int32Type>(ctx, input_array, result);
+      case Type::INT64:
+        return DoExecNested<Int64Type>(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<TypeHolder> ResolveOutputType(
       KernelContext* ctx, const std::vector<TypeHolder>& input_types) {
@@ -470,6 +536,56 @@ struct RunEndDecodeExec {
     return Status::Invalid("Invalid run end type: ", *ree_type->run_end_type());
   }
 
+  template <typename RunEndType>
+  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<RunEndCType> 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<const RunEndEncodedType*>(input_array.type);
+    switch (ree_type->run_end_type()->id()) {
+      case Type::INT16:
+        return DoExecNested<Int16Type>(ctx, input_array, result);
+      case Type::INT32:
+        return DoExecNested<Int32Type>(ctx, input_array, result);
+      case Type::INT64:
+        return DoExecNested<Int64Type>(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<TypeHolder> ResolveOutputType(KernelContext*,
                                               const std::vector<TypeHolder>& 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)));
 }

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<Array>& input,
+                                    const std::string& expected_run_ends_json,
+                                    const std::shared_ptr<Array>& 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<RunEndEncodedArray>(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