felipecrv commented on code in PR #41295:
URL: https://github.com/apache/arrow/pull/41295#discussion_r1583446385
##########
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 recursively = false);
Review Comment:
```suggestion
explicit ListFlattenOptions(bool recursive = false);
```
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -29,41 +30,114 @@ namespace compute {
using arrow::internal::checked_cast;
-TEST(TestVectorNested, ListFlatten) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- auto input = ArrayFromJSON(ty, "[[0, null, 1], null, [2, 3], []]");
- auto expected = ArrayFromJSON(int16(), "[0, null, 1, 2, 3]");
+using ListAndListViewTypes =
+ ::testing::Types<ListType, LargeListType, ListViewType, LargeListViewType>;
+
+// ----------------------------------------------------------------------
+// [Large]List and [Large]ListView tests
+template <typename T>
+class TestVectorLogicalList : public ::testing::Test {
Review Comment:
The goal of this file is to have tests for vector kernels on nested types.
It's incidental that it only has tests for list-related functions.
Please rename this to `TestVectorNestedSpecialized` so that we can use
`TestVectorNested` below.
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -29,41 +30,114 @@ namespace compute {
using arrow::internal::checked_cast;
-TEST(TestVectorNested, ListFlatten) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- auto input = ArrayFromJSON(ty, "[[0, null, 1], null, [2, 3], []]");
- auto expected = ArrayFromJSON(int16(), "[0, null, 1, 2, 3]");
+using ListAndListViewTypes =
+ ::testing::Types<ListType, LargeListType, ListViewType, LargeListViewType>;
+
+// ----------------------------------------------------------------------
+// [Large]List and [Large]ListView tests
+template <typename T>
+class TestVectorLogicalList : public ::testing::Test {
+ public:
+ using TypeClass = T;
+
+ void SetUp() override {
+ value_type_ = int16();
+ type_ = std::make_shared<T>(value_type_);
+ }
+
+ public:
+ void TestListFlatten() {
+ auto input = ArrayFromJSON(type_, "[[0, null, 1], null, [2, 3], []]");
+ auto expected = ArrayFromJSON(value_type_, "[0, null, 1, 2, 3]");
CheckVectorUnary("list_flatten", input, expected);
// Construct a list with a non-empty null slot
auto tweaked = TweakValidityBit(input, 0, false);
- expected = ArrayFromJSON(int16(), "[2, 3]");
+ expected = ArrayFromJSON(value_type_, "[2, 3]");
CheckVectorUnary("list_flatten", tweaked, expected);
}
-}
-TEST(TestVectorNested, ListFlattenNulls) {
- const auto ty = list(int32());
- auto input = ArrayFromJSON(ty, "[null, null]");
- auto expected = ArrayFromJSON(int32(), "[]");
- CheckVectorUnary("list_flatten", input, expected);
-}
+ void TestListFlattenNulls() {
+ value_type_ = int32();
+ type_ = std::make_shared<T>(value_type_);
+ auto input = ArrayFromJSON(type_, "[null, null]");
+ auto expected = ArrayFromJSON(value_type_, "[]");
+ CheckVectorUnary("list_flatten", input, expected);
+ }
-TEST(TestVectorNested, ListFlattenChunkedArray) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- ARROW_SCOPED_TRACE(ty->ToString());
- auto input = ChunkedArrayFromJSON(ty, {"[[0, null, 1], null]", "[[2, 3],
[]]"});
- auto expected = ChunkedArrayFromJSON(int16(), {"[0, null, 1]", "[2, 3]"});
+ void TestListFlattenChunkedArray() {
+ ARROW_SCOPED_TRACE(type_->ToString());
+ auto input = ChunkedArrayFromJSON(type_, {"[[0, null, 1], null]", "[[2,
3], []]"});
+ auto expected = ChunkedArrayFromJSON(value_type_, {"[0, null, 1]", "[2,
3]"});
CheckVectorUnary("list_flatten", input, expected);
ARROW_SCOPED_TRACE("empty");
- input = ChunkedArrayFromJSON(ty, {});
- expected = ChunkedArrayFromJSON(int16(), {});
+ input = ChunkedArrayFromJSON(type_, {});
+ expected = ChunkedArrayFromJSON(value_type_, {});
CheckVectorUnary("list_flatten", input, expected);
}
+
+ void TestListFlattenRecursively() {
+ auto inner_type = std::make_shared<T>(value_type_);
+ type_ = std::make_shared<T>(inner_type);
+
+ ListFlattenOptions opts;
+ opts.recursive = true;
+
+ // List types with two nested level: list<list<int16>>
+ auto input = ArrayFromJSON(type_, R"([
+ [[0, 1, 2], null, [3, null]],
+ [null],
+ [[2, 9], [4], [], [6, 5]]
+ ])");
+ auto expected = ArrayFromJSON(value_type_, "[0, 1, 2, 3, null, 2, 9, 4, 6,
5]");
+ CheckVectorUnary("list_flatten", input, expected, &opts);
+
+ // Empty nested list should flatten until non-list type is reached
+ input = ArrayFromJSON(type_, R"([null])");
+ expected = ArrayFromJSON(value_type_, "[]");
+ CheckVectorUnary("list_flatten", input, expected, &opts);
+
+ // List types with three nested level: list<list<fixed_size_list<int32,
2>>>
Review Comment:
```suggestion
// List types with three nesting levels:
list<list<fixed_size_list<int32, 2>>>
```
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -29,41 +30,114 @@ namespace compute {
using arrow::internal::checked_cast;
-TEST(TestVectorNested, ListFlatten) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- auto input = ArrayFromJSON(ty, "[[0, null, 1], null, [2, 3], []]");
- auto expected = ArrayFromJSON(int16(), "[0, null, 1, 2, 3]");
+using ListAndListViewTypes =
+ ::testing::Types<ListType, LargeListType, ListViewType, LargeListViewType>;
+
+// ----------------------------------------------------------------------
+// [Large]List and [Large]ListView tests
+template <typename T>
+class TestVectorLogicalList : public ::testing::Test {
+ public:
+ using TypeClass = T;
+
+ void SetUp() override {
+ value_type_ = int16();
+ type_ = std::make_shared<T>(value_type_);
+ }
+
+ public:
+ void TestListFlatten() {
+ auto input = ArrayFromJSON(type_, "[[0, null, 1], null, [2, 3], []]");
+ auto expected = ArrayFromJSON(value_type_, "[0, null, 1, 2, 3]");
CheckVectorUnary("list_flatten", input, expected);
// Construct a list with a non-empty null slot
auto tweaked = TweakValidityBit(input, 0, false);
- expected = ArrayFromJSON(int16(), "[2, 3]");
+ expected = ArrayFromJSON(value_type_, "[2, 3]");
CheckVectorUnary("list_flatten", tweaked, expected);
}
-}
-TEST(TestVectorNested, ListFlattenNulls) {
- const auto ty = list(int32());
- auto input = ArrayFromJSON(ty, "[null, null]");
- auto expected = ArrayFromJSON(int32(), "[]");
- CheckVectorUnary("list_flatten", input, expected);
-}
+ void TestListFlattenNulls() {
+ value_type_ = int32();
+ type_ = std::make_shared<T>(value_type_);
+ auto input = ArrayFromJSON(type_, "[null, null]");
+ auto expected = ArrayFromJSON(value_type_, "[]");
+ CheckVectorUnary("list_flatten", input, expected);
+ }
-TEST(TestVectorNested, ListFlattenChunkedArray) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- ARROW_SCOPED_TRACE(ty->ToString());
- auto input = ChunkedArrayFromJSON(ty, {"[[0, null, 1], null]", "[[2, 3],
[]]"});
- auto expected = ChunkedArrayFromJSON(int16(), {"[0, null, 1]", "[2, 3]"});
+ void TestListFlattenChunkedArray() {
+ ARROW_SCOPED_TRACE(type_->ToString());
+ auto input = ChunkedArrayFromJSON(type_, {"[[0, null, 1], null]", "[[2,
3], []]"});
+ auto expected = ChunkedArrayFromJSON(value_type_, {"[0, null, 1]", "[2,
3]"});
CheckVectorUnary("list_flatten", input, expected);
ARROW_SCOPED_TRACE("empty");
- input = ChunkedArrayFromJSON(ty, {});
- expected = ChunkedArrayFromJSON(int16(), {});
+ input = ChunkedArrayFromJSON(type_, {});
+ expected = ChunkedArrayFromJSON(value_type_, {});
CheckVectorUnary("list_flatten", input, expected);
}
+
+ void TestListFlattenRecursively() {
+ auto inner_type = std::make_shared<T>(value_type_);
+ type_ = std::make_shared<T>(inner_type);
+
+ ListFlattenOptions opts;
+ opts.recursive = true;
+
+ // List types with two nested level: list<list<int16>>
Review Comment:
```suggestion
// List types with two nesting levels: list<list<int16>>
```
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -85,14 +159,29 @@ TEST(TestVectorNested, ListFlattenFixedSizeList) {
}
}
-TEST(TestVectorNested, ListFlattenFixedSizeListNulls) {
+TEST(TestVectorFixedSizeList, ListFlattenFixedSizeListNulls) {
const auto ty = fixed_size_list(int32(), 1);
auto input = ArrayFromJSON(ty, "[null, null]");
auto expected = ArrayFromJSON(int32(), "[]");
CheckVectorUnary("list_flatten", input, expected);
}
-TEST(TestVectorNested, ListParentIndices) {
+TEST(TestVectorFixedSizeList, ListFlattenFixedSizeListRecursively) {
Review Comment:
Same here -- `TestVectorNested`
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -121,7 +210,7 @@ TEST(TestVectorNested, ListParentIndicesChunkedArray) {
}
}
-TEST(TestVectorNested, ListParentIndicesFixedSizeList) {
+TEST(TestVectorListParentIndices, FixedSizeListArray) {
Review Comment:
Same here -- `TestVectorNested`
##########
python/pyarrow/_compute.pyx:
##########
@@ -2035,6 +2035,26 @@ class PairwiseOptions(_PairwiseOptions):
self._set_options(period)
+cdef class _ListFlattenOptions(FunctionOptions):
+ def _set_options(self, recursive):
+ self.wrapped.reset(new CListFlattenOptions(recursive))
+
+
+class ListFlattenOptions(_ListFlattenOptions):
+ """
+ Options for `list_flatten` function
+
+ Parameters
+ ----------
+ recursive : bool, defalut false
Review Comment:
```suggestion
recursive : bool, default False
```
##########
python/pyarrow/array.pxi:
##########
@@ -2141,22 +2141,99 @@ cdef class Decimal256Array(FixedSizeBinaryArray):
cdef class BaseListArray(Array):
- def flatten(self):
+ def flatten(self, recursive=False):
"""
- Unnest this ListArray/LargeListArray by one level.
-
- The returned Array is logically a concatenation of all the sub-lists
- in this Array.
+ Unnest this [Large]ListArray/[Large]ListViewArray/FixedSizeListArray
+ according to 'recursive'.
Note that this method is different from ``self.values`` in that
it takes care of the slicing offset as well as null elements backed
by non-empty sub-lists.
+ Parameters
+ ----------
+ recursive : bool, defalut false, optional
+ When true, flatten this logical list-array recursivey until an
+ array of non-list values is reached.
+ When false, flatten this logical list-array by one level
Review Comment:
```suggestion
recursive : bool, default False, optional
When True, flatten this logical list-array recursively until an
array of non-list values is formed.
When False, flatten only the top level.
```
##########
cpp/src/arrow/compute/kernels/vector_nested.cc:
##########
@@ -107,10 +113,13 @@ 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"
+ ("`lists` must have a logical list type like `[Large]ListType`, \n"
+ "`[Large]ListViewType` and `FixedSizeListType`. \n"
+ "Whether to flatten the top list level or the bottom list level \n"
+ "will be decided based on the `recursive` option specified in \n"
+ ":struct:`ListFlattenOptions`. \n"
"Top-level null values in `lists` do not emit anything in the input."),
Review Comment:
```suggestion
("`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."),
```
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -85,14 +159,29 @@ TEST(TestVectorNested, ListFlattenFixedSizeList) {
}
}
-TEST(TestVectorNested, ListFlattenFixedSizeListNulls) {
+TEST(TestVectorFixedSizeList, ListFlattenFixedSizeListNulls) {
Review Comment:
Same here -- `TestVectorNested`
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -29,41 +30,114 @@ namespace compute {
using arrow::internal::checked_cast;
-TEST(TestVectorNested, ListFlatten) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- auto input = ArrayFromJSON(ty, "[[0, null, 1], null, [2, 3], []]");
- auto expected = ArrayFromJSON(int16(), "[0, null, 1, 2, 3]");
+using ListAndListViewTypes =
+ ::testing::Types<ListType, LargeListType, ListViewType, LargeListViewType>;
+
+// ----------------------------------------------------------------------
+// [Large]List and [Large]ListView tests
+template <typename T>
+class TestVectorLogicalList : public ::testing::Test {
+ public:
+ using TypeClass = T;
+
+ void SetUp() override {
+ value_type_ = int16();
+ type_ = std::make_shared<T>(value_type_);
+ }
+
+ public:
+ void TestListFlatten() {
+ auto input = ArrayFromJSON(type_, "[[0, null, 1], null, [2, 3], []]");
+ auto expected = ArrayFromJSON(value_type_, "[0, null, 1, 2, 3]");
CheckVectorUnary("list_flatten", input, expected);
// Construct a list with a non-empty null slot
auto tweaked = TweakValidityBit(input, 0, false);
- expected = ArrayFromJSON(int16(), "[2, 3]");
+ expected = ArrayFromJSON(value_type_, "[2, 3]");
CheckVectorUnary("list_flatten", tweaked, expected);
}
-}
-TEST(TestVectorNested, ListFlattenNulls) {
- const auto ty = list(int32());
- auto input = ArrayFromJSON(ty, "[null, null]");
- auto expected = ArrayFromJSON(int32(), "[]");
- CheckVectorUnary("list_flatten", input, expected);
-}
+ void TestListFlattenNulls() {
+ value_type_ = int32();
+ type_ = std::make_shared<T>(value_type_);
+ auto input = ArrayFromJSON(type_, "[null, null]");
+ auto expected = ArrayFromJSON(value_type_, "[]");
+ CheckVectorUnary("list_flatten", input, expected);
+ }
-TEST(TestVectorNested, ListFlattenChunkedArray) {
- for (auto ty : {list(int16()), large_list(int16())}) {
- ARROW_SCOPED_TRACE(ty->ToString());
- auto input = ChunkedArrayFromJSON(ty, {"[[0, null, 1], null]", "[[2, 3],
[]]"});
- auto expected = ChunkedArrayFromJSON(int16(), {"[0, null, 1]", "[2, 3]"});
+ void TestListFlattenChunkedArray() {
+ ARROW_SCOPED_TRACE(type_->ToString());
+ auto input = ChunkedArrayFromJSON(type_, {"[[0, null, 1], null]", "[[2,
3], []]"});
+ auto expected = ChunkedArrayFromJSON(value_type_, {"[0, null, 1]", "[2,
3]"});
CheckVectorUnary("list_flatten", input, expected);
ARROW_SCOPED_TRACE("empty");
- input = ChunkedArrayFromJSON(ty, {});
- expected = ChunkedArrayFromJSON(int16(), {});
+ input = ChunkedArrayFromJSON(type_, {});
+ expected = ChunkedArrayFromJSON(value_type_, {});
CheckVectorUnary("list_flatten", input, expected);
}
+
+ void TestListFlattenRecursively() {
+ auto inner_type = std::make_shared<T>(value_type_);
+ type_ = std::make_shared<T>(inner_type);
+
+ ListFlattenOptions opts;
+ opts.recursive = true;
+
+ // List types with two nested level: list<list<int16>>
+ auto input = ArrayFromJSON(type_, R"([
+ [[0, 1, 2], null, [3, null]],
+ [null],
+ [[2, 9], [4], [], [6, 5]]
+ ])");
+ auto expected = ArrayFromJSON(value_type_, "[0, 1, 2, 3, null, 2, 9, 4, 6,
5]");
+ CheckVectorUnary("list_flatten", input, expected, &opts);
+
+ // Empty nested list should flatten until non-list type is reached
+ input = ArrayFromJSON(type_, R"([null])");
+ expected = ArrayFromJSON(value_type_, "[]");
+ CheckVectorUnary("list_flatten", input, expected, &opts);
+
+ // List types with three nested level: list<list<fixed_size_list<int32,
2>>>
+ type_ =
std::make_shared<T>(std::make_shared<T>(fixed_size_list(value_type_, 2)));
+ input = ArrayFromJSON(type_, R"([
+ [
+ [[null, 0]],
+ [[3, 7], null]
+ ],
+ [
+ [[4, null], [5, 8]],
+ [[8, null]],
+ null
+ ],
+ [
+ null
+ ]
+ ])");
+ expected = ArrayFromJSON(value_type_, "[null, 0, 3, 7, 4, null, 5, 8, 8,
null]");
+ CheckVectorUnary("list_flatten", input, expected, &opts);
+ }
+
+ protected:
+ std::shared_ptr<DataType> type_;
+ std::shared_ptr<DataType> value_type_;
+};
+
+TYPED_TEST_SUITE(TestVectorLogicalList, ListAndListViewTypes);
+
+TYPED_TEST(TestVectorLogicalList, ListFlatten) { this->TestListFlatten(); }
+
+TYPED_TEST(TestVectorLogicalList, ListFlattenNulls) {
this->TestListFlattenNulls(); }
+
+TYPED_TEST(TestVectorLogicalList, ListFlattenChunkedArray) {
+ this->TestListFlattenChunkedArray();
+}
+
+TYPED_TEST(TestVectorLogicalList, ListFlattenRecursively) {
+ this->TestListFlattenRecursively();
}
-TEST(TestVectorNested, ListFlattenFixedSizeList) {
+TEST(TestVectorFixedSizeList, ListFlattenFixedSizeList) {
Review Comment:
Keep using `TestVectorNested` instead of repeating `FixedSizeList` twice.
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -107,7 +196,7 @@ TEST(TestVectorNested, ListParentIndices) {
CheckVectorUnary("list_parent_indices", tweaked, expected);
}
-TEST(TestVectorNested, ListParentIndicesChunkedArray) {
Review Comment:
Same here -- `TestVectorNested`
##########
python/pyarrow/_compute.pyx:
##########
@@ -2035,6 +2035,26 @@ class PairwiseOptions(_PairwiseOptions):
self._set_options(period)
+cdef class _ListFlattenOptions(FunctionOptions):
+ def _set_options(self, recursive):
+ self.wrapped.reset(new CListFlattenOptions(recursive))
+
+
+class ListFlattenOptions(_ListFlattenOptions):
+ """
+ Options for `list_flatten` function
+
+ Parameters
+ ----------
+ recursive : bool, defalut false
+ When true, do list flatten recursively until an array of
+ non-list values is reached.
Review Comment:
```suggestion
When True, the list array is flattened recursively until an array
of non-list values is formed.
```
##########
python/pyarrow/array.pxi:
##########
@@ -2141,22 +2141,99 @@ cdef class Decimal256Array(FixedSizeBinaryArray):
cdef class BaseListArray(Array):
- def flatten(self):
+ def flatten(self, recursive=False):
"""
- Unnest this ListArray/LargeListArray by one level.
-
- The returned Array is logically a concatenation of all the sub-lists
- in this Array.
+ Unnest this [Large]ListArray/[Large]ListViewArray/FixedSizeListArray
+ according to 'recursive'.
Note that this method is different from ``self.values`` in that
it takes care of the slicing offset as well as null elements backed
by non-empty sub-lists.
+ Parameters
+ ----------
+ recursive : bool, defalut false, optional
+ When true, flatten this logical list-array recursivey until an
+ array of non-list values is reached.
+ When false, flatten this logical list-array by one level
+
Returns
-------
result : Array
+
+ Examples
+ --------
+
+ Basic logical list-array's flatten
+ >>> import pyarrow as pa
+ >>> values = [1, 2, 3, 4]
+ >>> offsets = [2, 1, 0]
+ >>> sizes = [2, 2, 2]
+ >>> array = pa.ListViewArray.from_arrays(offsets, sizes, values)
+ >>> array
+ <pyarrow.lib.ListViewArray object at ...>
+ [
+ [
+ 3,
+ 4
+ ],
+ [
+ 2,
+ 3
+ ],
+ [
+ 1,
+ 2
+ ]
+ ]
+ >>> array.flatten()
+ <pyarrow.lib.Int64Array object at ...>
+ [
+ 3,
+ 4,
+ 2,
+ 3,
+ 1,
+ 2
+ ]
+
+ If an logical list-array is nested with multi-level, the array will
+ be flattened recursively until an array of non-list values is reached
+ if we enable recursive=True.
Review Comment:
```suggestion
When recursive=True, nested list arrays are flattened recursively
until an array of non-list values is formed.
```
##########
cpp/src/arrow/compute/kernels/vector_nested_test.cc:
##########
@@ -85,14 +159,29 @@ TEST(TestVectorNested, ListFlattenFixedSizeList) {
}
}
-TEST(TestVectorNested, ListFlattenFixedSizeListNulls) {
+TEST(TestVectorFixedSizeList, ListFlattenFixedSizeListNulls) {
const auto ty = fixed_size_list(int32(), 1);
auto input = ArrayFromJSON(ty, "[null, null]");
auto expected = ArrayFromJSON(int32(), "[]");
CheckVectorUnary("list_flatten", input, expected);
}
-TEST(TestVectorNested, ListParentIndices) {
+TEST(TestVectorFixedSizeList, ListFlattenFixedSizeListRecursively) {
+ ListFlattenOptions opts;
+ opts.recursive = true;
+
+ auto inner_type = fixed_size_list(int32(), 2);
+ auto type = fixed_size_list(inner_type, 2);
+ auto input = ArrayFromJSON(type, R"([
+ [[0, 1], [null, 3]],
+ [[7, null], [2, 5]],
+ [null, null]
+ ])");
+ auto expected = ArrayFromJSON(int32(), "[0, 1, null, 3, 7, null, 2, 5]");
+ CheckVectorUnary("list_flatten", input, expected, &opts);
+}
+
+TEST(TestVectorListParentIndices, BasicListArray) {
Review Comment:
Same here -- `TestVectorNested`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
