Skip to content

Commit

Permalink
add a new flatten api with a recursion argument
Browse files Browse the repository at this point in the history
  • Loading branch information
@ZhangHuiGui
ZhangHuiGui committed Apr 9, 2024
1 parent 26b7a6f commit c8fc9c1
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 15 deletions.
32 changes: 32 additions & 0 deletions cpp/src/arrow/array/array_nested.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -580,6 +580,10 @@ Result<std::shared_ptr<ListArray>> ListArray::FromArrays(
null_bitmap, null_count);
}

Result<std::shared_ptr<Array>> ListArray::Flatten(MemoryPool* memory_pool) const {
return FlattenListArray(*this, /*with_recursion=*/false, memory_pool);
}

Result<std::shared_ptr<Array>> ListArray::Flatten(bool with_recursion,
MemoryPool* memory_pool) const {
return FlattenListArray(*this, with_recursion, memory_pool);
Expand Down Expand Up @@ -640,6 +644,10 @@ Result<std::shared_ptr<LargeListArray>> LargeListArray::FromArrays(
null_bitmap, null_count);
}

Result<std::shared_ptr<Array>> LargeListArray::Flatten(MemoryPool* memory_pool) const {
return FlattenListArray(*this, /*with_recursion=*/false, memory_pool);
}

Result<std::shared_ptr<Array>> LargeListArray::Flatten(bool with_recursion,
MemoryPool* memory_pool) const {
return FlattenListArray(*this, with_recursion, memory_pool);
Expand Down Expand Up @@ -711,6 +719,15 @@ Result<std::shared_ptr<LargeListViewArray>> LargeListViewArray::FromList(
return std::make_shared<LargeListViewArray>(std::move(data));
}

Result<std::shared_ptr<Array>> ListViewArray::Flatten(MemoryPool* memory_pool) const {
if (null_count() > 0) {
return FlattenListViewArray<ListViewArray, true>(*this, /*with_recursion=*/false,
memory_pool);
}
return FlattenListViewArray<ListViewArray, false>(*this, /*with_recursion=*/false,
memory_pool);
}

Result<std::shared_ptr<Array>> ListViewArray::Flatten(bool with_recursion,
MemoryPool* memory_pool) const {
if (null_count() > 0) {
Expand Down Expand Up @@ -772,6 +789,16 @@ Result<std::shared_ptr<LargeListViewArray>> LargeListViewArray::FromArrays(
std::move(type), offsets, sizes, values, pool, null_bitmap, null_count);
}

Result<std::shared_ptr<Array>> LargeListViewArray::Flatten(
MemoryPool* memory_pool) const {
if (null_count() > 0) {
return FlattenListViewArray<LargeListViewArray, true>(*this, /*with_recursion=*/false,
memory_pool);
}
return FlattenListViewArray<LargeListViewArray, false>(*this, /*with_recursion=*/false,
memory_pool);
}

Result<std::shared_ptr<Array>> LargeListViewArray::Flatten(
bool with_recursion, MemoryPool* memory_pool) const {
if (null_count() > 0) {
Expand Down Expand Up @@ -995,6 +1022,11 @@ Result<std::shared_ptr<Array>> FixedSizeListArray::FromArrays(
null_count);
}

Result<std::shared_ptr<Array>> FixedSizeListArray::Flatten(
MemoryPool* memory_pool) const {
return FlattenListArray(*this, /*with_recursion=*/false, memory_pool);
}

Result<std::shared_ptr<Array>> FixedSizeListArray::Flatten(
bool with_recursion, MemoryPool* memory_pool) const {
return FlattenListArray(*this, with_recursion, memory_pool);
Expand Down
50 changes: 35 additions & 15 deletions cpp/src/arrow/array/array_nested.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -183,13 +183,17 @@ class ARROW_EXPORT ListArray : public BaseListArray<ListType> {

/// \brief Return an Array that is a concatenation of the lists in this array.
///
/// \param[in] with_recursion Flatten recursively until reach non-list type
///
/// Note that it's different from `values()` in that it takes into
/// consideration of this array's offsets as well as null elements backed
/// by non-empty lists (they are skipped, thus copying may be needed).
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion = false, MemoryPool* memory_pool = default_memory_pool()) const;
MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Same as the Flatten function above.
///
/// \param[in] with_recursion Flatten all level recursively until reach non-list type
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion, MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Return list offsets as an Int32Array
///
Expand Down Expand Up @@ -253,13 +257,17 @@ class ARROW_EXPORT LargeListArray : public BaseListArray<LargeListType> {

/// \brief Return an Array that is a concatenation of the lists in this array.
///
/// \param[in] with_recursion Flatten recursively until reach non-list type
///
/// Note that it's different from `values()` in that it takes into
/// consideration of this array's offsets as well as null elements backed
/// by non-empty lists (they are skipped, thus copying may be needed).
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion = false, MemoryPool* memory_pool = default_memory_pool()) const;
MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Same as the Flatten function above.
///
/// \param[in] with_recursion Flatten all level recursively until reach non-list type
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion, MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Return list offsets as an Int64Array
std::shared_ptr<Array> offsets() const;
Expand Down Expand Up @@ -357,8 +365,6 @@ class ARROW_EXPORT ListViewArray : public BaseListViewArray<ListViewType> {

/// \brief Return an Array that is a concatenation of the list-views in this array.
///
/// \param[in] with_recursion Flatten recursively until reach non-list type
///
/// Note that it's different from `values()` in that it takes into
/// consideration this array's offsets (which can be in any order)
/// and sizes. Nulls are skipped.
Expand All @@ -368,7 +374,13 @@ class ARROW_EXPORT ListViewArray : public BaseListViewArray<ListViewType> {
/// maximizing the size of each slice (containing as many contiguous
/// list-views as possible).
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion = false, MemoryPool* memory_pool = default_memory_pool()) const;
MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Same as the 'Flatten' function above.
///
/// \param[in] with_recursion Flatten all level recursively until reach non-list type
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion, MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Return list-view offsets as an Int32Array
///
Expand Down Expand Up @@ -448,13 +460,17 @@ class ARROW_EXPORT LargeListViewArray : public BaseListViewArray<LargeListViewTy
/// \brief Return an Array that is a concatenation of the large list-views in this
/// array.
///
/// \param[in] with_recursion Flatten recursively until reach non-list type
///
/// Note that it's different from `values()` in that it takes into
/// consideration this array's offsets (which can be in any order)
/// and sizes. Nulls are skipped.
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion = false, MemoryPool* memory_pool = default_memory_pool()) const;
MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Same as the Flatten function above.
///
/// \param[in] with_recursion Flatten all level recursively until reach non-list type
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion, MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Return list-view offsets as an Int64Array
///
Expand Down Expand Up @@ -598,12 +614,16 @@ class ARROW_EXPORT FixedSizeListArray : public Array {

/// \brief Return an Array that is a concatenation of the lists in this array.
///
/// \param[in] with_recursion Flatten recursively until reach non-list type
///
/// Note that it's different from `values()` in that it takes into
/// consideration null elements (they are skipped, thus copying may be needed).
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion = false, MemoryPool* memory_pool = default_memory_pool()) const;
MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Same as the Flatten function above.
///
/// \param[in] with_recursion Flatten all level recursively until reach non-list type
Result<std::shared_ptr<Array>> Flatten(
bool with_recursion, MemoryPool* memory_pool = default_memory_pool()) const;

/// \brief Construct FixedSizeListArray from child value array and value_length
///
Expand Down

0 comments on commit c8fc9c1

Please sign in to comment.