GH-41183: [C++][Python] Expose recursive flatten for lists on list_flatten kernel function and pyarrow bindings

cpp/src/arrow/compute/api_vector.cc

@@ -153,6 +153,8 @@ static auto kRankOptionsType = GetFunctionOptionsType<RankOptions>(
  DataMember("tiebreaker", &RankOptions::tiebreaker));
 static auto kPairwiseOptionsType = GetFunctionOptionsType<PairwiseOptions>(
  DataMember("periods", &PairwiseOptions::periods));
+static auto kListFlattenOptionsType = GetFunctionOptionsType<ListFlattenOptions>(
+ DataMember("recursive", &ListFlattenOptions::recursive));
 } // namespace
 } // namespace internal
 
@@ -224,6 +226,10 @@ PairwiseOptions::PairwiseOptions(int64_t periods)
  : FunctionOptions(internal::kPairwiseOptionsType), periods(periods) {}
 constexpr char PairwiseOptions::kTypeName[];
 
+ListFlattenOptions::ListFlattenOptions(bool recursive)
+ : FunctionOptions(internal::kListFlattenOptionsType), recursive(recursive) {}
+constexpr char ListFlattenOptions::kTypeName[];
+
 namespace internal {
 void RegisterVectorOptions(FunctionRegistry* registry) {
  DCHECK_OK(registry->AddFunctionOptionsType(kFilterOptionsType));
@@ -237,6 +243,7 @@ void RegisterVectorOptions(FunctionRegistry* registry) {
  DCHECK_OK(registry->AddFunctionOptionsType(kCumulativeOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kRankOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kPairwiseOptionsType));
+ DCHECK_OK(registry->AddFunctionOptionsType(kListFlattenOptionsType));
 }
 } // namespace internal
 

cpp/src/arrow/compute/api_vector.h

@@ -245,6 +245,18 @@ class ARROW_EXPORT PairwiseOptions : public FunctionOptions {
  int64_t periods = 1;
 };
 
+/// \brief Options for list_flatten function
+class ARROW_EXPORT ListFlattenOptions : public FunctionOptions {
+ public:
+ explicit ListFlattenOptions(bool recursive = false);
+ static constexpr char const kTypeName[] = "ListFlattenOptions";
+ static ListFlattenOptions Defaults() { return ListFlattenOptions(); }
+
+ /// \brief If true, the list is flattened recursively until a non-list
+ /// array is formed.
+ bool recursive = false;
+};
+
 /// @}
 
 /// \brief Filter with a boolean selection filter

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

@@ -23,6 +23,7 @@
 #include <mutex>
 #include <vector>
 
+#include "arrow/compute/api_vector.h"
 #include "arrow/type_fwd.h"
 
 namespace arrow {
@@ -56,9 +57,23 @@ Result<TypeHolder> LastType(KernelContext*, const std::vector<TypeHolder>& types
  return types.back();
 }
 
-Result<TypeHolder> ListValuesType(KernelContext*, const std::vector<TypeHolder>& args) {
- const auto& list_type = checked_cast<const BaseListType&>(*args[0].type);
- return list_type.value_type().get();
+Result<TypeHolder> ListValuesType(KernelContext* ctx,
+ const std::vector<TypeHolder>& args) {
+ auto list_type = checked_cast<const BaseListType*>(args[0].type);
+ auto value_type = list_type->value_type().get();
+
+ auto recursive =
+ ctx->state() ? OptionsWrapper<ListFlattenOptions>::Get(ctx).recursive : false;
+ if (!recursive) {
+ return value_type;
+ }
+
+ for (auto value_kind = value_type->id();
+ is_list(value_kind) || is_list_view(value_kind); value_kind = value_type->id()) {
+ list_type = checked_cast<const BaseListType*>(list_type->value_type().get());
+ value_type = list_type->value_type().get();
+ }
+ return value_type;
 }
 
 void EnsureDictionaryDecoded(std::vector<TypeHolder>* types) {

cpp/src/arrow/compute/kernels/codegen_internal.h

@@ -423,7 +423,8 @@ static void VisitTwoArrayValuesInline(const ArraySpan& arr0, const ArraySpan& ar
 
 Result<TypeHolder> FirstType(KernelContext*, const std::vector<TypeHolder>& types);
 Result<TypeHolder> LastType(KernelContext*, const std::vector<TypeHolder>& types);
-Result<TypeHolder> ListValuesType(KernelContext*, const std::vector<TypeHolder>& types);
+Result<TypeHolder> ListValuesType(KernelContext* ctx,
+ const std::vector<TypeHolder>& types);
 
 // ----------------------------------------------------------------------
 // Helpers for iterating over common DataType instances for adding kernels to

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

@@ -23,6 +23,7 @@
 #include "arrow/compute/api_scalar.h"
 #include "arrow/compute/kernels/common_internal.h"
 #include "arrow/result.h"
+#include "arrow/type_fwd.h"
 #include "arrow/util/bit_block_counter.h"
 #include "arrow/util/bit_util.h"
 #include "arrow/util/bitmap_generate.h"
@@ -41,10 +42,17 @@ Status ListValueLength(KernelContext* ctx, const ExecSpan& batch, ExecResult* ou
  const ArraySpan& arr = batch[0].array;
  ArraySpan* out_arr = out->array_span_mutable();
  auto out_values = out_arr->GetValues<offset_type>(1);
- const offset_type* offsets = arr.GetValues<offset_type>(1);
- // Offsets are always well-defined and monotonic, even for null values
- for (int64_t i = 0; i < arr.length; ++i) {
- *out_values++ = offsets[i + 1] - offsets[i];
+ if (is_list_view(*arr.type)) {
+ const auto* sizes = arr.GetValues<offset_type>(2);
+ if (arr.length > 0) {
+ memcpy(out_values, sizes, arr.length * sizeof(offset_type));
+ }
+ } else {
+ const offset_type* offsets = arr.GetValues<offset_type>(1);
+ // Offsets are always well-defined and monotonic, even for null values
+ for (int64_t i = 0; i < arr.length; ++i) {
+ *out_values++ = offsets[i + 1] - offsets[i];
+ }
  }
  return Status::OK();
 }
@@ -59,6 +67,30 @@ Status FixedSizeListValueLength(KernelContext* ctx, const ExecSpan& batch,
  return Status::OK();
 }
 
+template <typename InListType>
+void AddListValueLengthKernel(ScalarFunction* func,
+ const std::shared_ptr<DataType>& out_type) {
+ auto in_type = {InputType(InListType::type_id)};
+ ScalarKernel kernel(in_type, out_type, ListValueLength<InListType>);
+ DCHECK_OK(func->AddKernel(std::move(kernel)));
+}
+
+template <>
+void AddListValueLengthKernel<FixedSizeListType>(
+ ScalarFunction* func, const std::shared_ptr<DataType>& out_type) {
+ auto in_type = {InputType(Type::FIXED_SIZE_LIST)};
+ ScalarKernel kernel(in_type, out_type, FixedSizeListValueLength);
+ DCHECK_OK(func->AddKernel(std::move(kernel)));
+}
+
+void AddListValueLengthKernels(ScalarFunction* func) {
+ AddListValueLengthKernel<ListType>(func, int32());
+ AddListValueLengthKernel<LargeListType>(func, int64());
+ AddListValueLengthKernel<ListViewType>(func, int32());
+ AddListValueLengthKernel<LargeListViewType>(func, int64());
+ AddListValueLengthKernel<FixedSizeListType>(func, int32());
+}
+
 const FunctionDoc list_value_length_doc{
  "Compute list lengths",
  ("`lists` must have a list-like type.\n"
@@ -399,6 +431,8 @@ void AddListElementKernels(ScalarFunction* func) {
 void AddListElementKernels(ScalarFunction* func) {
  AddListElementKernels<ListType, ListElement>(func);
  AddListElementKernels<LargeListType, ListElement>(func);
+ AddListElementKernels<ListViewType, ListElement>(func);
+ AddListElementKernels<LargeListViewType, ListElement>(func);
  AddListElementKernels<FixedSizeListType, FixedSizeListElement>(func);
 }
 
@@ -824,12 +858,7 @@ const FunctionDoc map_lookup_doc{
 void RegisterScalarNested(FunctionRegistry* registry) {
  auto list_value_length = std::make_shared<ScalarFunction>(
  "list_value_length", Arity::Unary(), list_value_length_doc);
- DCHECK_OK(list_value_length->AddKernel({InputType(Type::LIST)}, int32(),
- ListValueLength<ListType>));
- DCHECK_OK(list_value_length->AddKernel({InputType(Type::FIXED_SIZE_LIST)}, int32(),
- FixedSizeListValueLength));
- DCHECK_OK(list_value_length->AddKernel({InputType(Type::LARGE_LIST)}, int64(),
- ListValueLength<LargeListType>));
+ AddListValueLengthKernels(list_value_length.get());
  DCHECK_OK(registry->AddFunction(std::move(list_value_length)));
 
  auto list_element =

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

@@ -30,11 +30,21 @@ namespace arrow {
 namespace compute {
 
 static std::shared_ptr<DataType> GetOffsetType(const DataType& type) {
- return type.id() == Type::LIST ? int32() : int64();
+ switch (type.id()) {
+ case Type::LIST:
+ case Type::LIST_VIEW:
+ return int32();
+ case Type::LARGE_LIST:
+ case Type::LARGE_LIST_VIEW:
+ return int64();
+ default:
+ Unreachable("Unexpected type");
+ }
 }
 
 TEST(TestScalarNested, ListValueLength) {
- for (auto ty : {list(int32()), large_list(int32())}) {
+ for (auto ty : {list(int32()), large_list(int32()), list_view(int32()),
+ large_list_view(int32())}) {
  CheckScalarUnary("list_value_length", ty, "[[0, null, 1], null, [2, 3], []]",
  GetOffsetType(*ty), "[3, null, 2, 0]");
  }
@@ -47,7 +57,8 @@ TEST(TestScalarNested, ListValueLength) {
 TEST(TestScalarNested, ListElementNonFixedListWithNulls) {
  auto sample = "[[7, 5, 81], [6, null, 4, 7, 8], [3, 12, 2, 0], [1, 9], null]";
  for (auto ty : NumericTypes()) {
- for (auto list_type : {list(ty), large_list(ty)}) {
+ for (auto list_type :
+ {list(ty), large_list(ty), list_view(ty), large_list_view(ty)}) {
  auto input = ArrayFromJSON(list_type, sample);
  auto null_input = ArrayFromJSON(list_type, "[null]");
  for (auto index_type : IntTypes()) {

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

@@ -18,6 +18,7 @@
 // Vector kernels involving nested types
 
 #include "arrow/array/array_base.h"
+#include "arrow/compute/api_vector.h"
 #include "arrow/compute/kernels/common_internal.h"
 #include "arrow/result.h"
 #include "arrow/visit_type_inline.h"
@@ -29,8 +30,13 @@ namespace {
 
 template <typename Type>
 Status ListFlatten(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
+ auto recursive = OptionsWrapper<ListFlattenOptions>::Get(ctx).recursive;
  typename TypeTraits<Type>::ArrayType list_array(batch[0].array.ToArrayData());
- ARROW_ASSIGN_OR_RAISE(auto result, list_array.Flatten(ctx->memory_pool()));
+
+ auto pool = ctx->memory_pool();
+ ARROW_ASSIGN_OR_RAISE(auto result, (recursive ? list_array.FlattenRecursively(pool)
+ : list_array.Flatten(pool)));
+
  out->value = std::move(result->data());
  return Status::OK();
 }
@@ -107,10 +113,15 @@ struct ListParentIndicesArray {
 
 const FunctionDoc list_flatten_doc(
  "Flatten list values",
- ("`lists` must have a list-like type.\n"
- "Return an array with the top list level flattened.\n"
- "Top-level null values in `lists` do not emit anything in the input."),
- {"lists"});
+ ("`lists` must have a list-like type (lists, list-views, and\n"
+ "fixed-size lists).\n"
+ "Return an array with the top list level flattened unless\n"
+ "`recursive` is set to true in ListFlattenOptions. When that\n"
+ "is that case, flattening happens recursively until a non-list\n"
+ "array is formed.\n"
+ "\n"
+ "Null list values do not emit anything to the output."),
+ {"lists"}, "ListFlattenOptions");
 
 const FunctionDoc list_parent_indices_doc(
  "Compute parent indices of nested list values",
@@ -153,17 +164,34 @@ class ListParentIndicesFunction : public MetaFunction {
  }
 };
 
+const ListFlattenOptions* GetDefaultListFlattenOptions() {
+ static const auto kDefaultListFlattenOptions = ListFlattenOptions::Defaults();
+ return &kDefaultListFlattenOptions;
+}
+
+template <typename InListType>
+void AddBaseListFlattenKernels(VectorFunction* func) {
+ auto in_type = {InputType(InListType::type_id)};
+ auto out_type = OutputType(ListValuesType);
+ VectorKernel kernel(in_type, out_type, ListFlatten<InListType>,
+ OptionsWrapper<ListFlattenOptions>::Init);
+ DCHECK_OK(func->AddKernel(std::move(kernel)));
+}
+
+void AddBaseListFlattenKernels(VectorFunction* func) {
+ AddBaseListFlattenKernels<ListType>(func);
+ AddBaseListFlattenKernels<LargeListType>(func);
+ AddBaseListFlattenKernels<FixedSizeListType>(func);
+ AddBaseListFlattenKernels<ListViewType>(func);
+ AddBaseListFlattenKernels<LargeListViewType>(func);
+}
+
 } // namespace
 
 void RegisterVectorNested(FunctionRegistry* registry) {
- auto flatten =
- std::make_shared<VectorFunction>("list_flatten", Arity::Unary(), list_flatten_doc);
- DCHECK_OK(flatten->AddKernel({Type::LIST}, OutputType(ListValuesType),
- ListFlatten<ListType>));
- DCHECK_OK(flatten->AddKernel({Type::FIXED_SIZE_LIST}, OutputType(ListValuesType),
- ListFlatten<FixedSizeListType>));
- DCHECK_OK(flatten->AddKernel({Type::LARGE_LIST}, OutputType(ListValuesType),
- ListFlatten<LargeListType>));
+ auto flatten = std::make_shared<VectorFunction>(
+ "list_flatten", Arity::Unary(), list_flatten_doc, GetDefaultListFlattenOptions());
+ AddBaseListFlattenKernels(flatten.get());
  DCHECK_OK(registry->AddFunction(std::move(flatten)));
 
  DCHECK_OK(registry->AddFunction(std::make_shared<ListParentIndicesFunction>()));