Added a C++ method Flatten() to ListArray.

This commit is contained in:
Zhuo Peng authored and Antoine Pitrou committed 2019-12-18 11:53:06 +01:00
1 parent 4c63bef8de
commit a3d4d2fbba
2 files changed
+42

No files matched your search

+11
View File
@@ -551,6 +551,17 @@ class BaseListArray : public Array {
/// \brief Return array object containing the list's values
std::shared_ptr<Array> values() const { return values_; }
/// \brief Return an Array that is a concatenation of the lists in this array.
///
/// Note that it's different from `values()` in that it takes into
/// consideration of this array's offsets and slices its child Array
/// accordingly.
std::shared_ptr<Array> Flatten() const {
const offset_type begin_offset = value_offset(0);
const offset_type end_offset = value_offset(length());
return values()->Slice(begin_offset, end_offset - begin_offset);
}
/// Note that this buffer does not account for any slice offset
std::shared_ptr<Buffer> value_offsets() const { return data_->buffers[1]; }
+31
View File
@@ -347,6 +347,34 @@ class TestListArray : public TestBuilder {
ASSERT_EQ("counts", type.value_field()->name());
}
void TestFlattenZeroLength() {
Done();
auto flattened = result_->Flatten();
ASSERT_OK(flattened->ValidateFull());
ASSERT_EQ(0, flattened->length());
}
void TestFlatten() {
auto type = std::make_shared<T>(int32());
auto list_array = std::dynamic_pointer_cast<ArrayType>(
ArrayFromJSON(type, "[[1, 2], [3], [4], null, [5], [], [6]]"));
auto flattened = list_array->Flatten();
ASSERT_OK(flattened->ValidateFull());
EXPECT_TRUE(
list_array->Flatten()->Equals(ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
EXPECT_TRUE(
list_array->values()->Equals(ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
auto sliced_list_array =
std::dynamic_pointer_cast<ArrayType>(list_array->Slice(3, 4));
auto sliced_flattened = sliced_list_array->Flatten();
ASSERT_OK(sliced_flattened->ValidateFull());
// Note the difference between values() and Flatten().
EXPECT_TRUE(sliced_flattened->Equals(ArrayFromJSON(int32(), "[5, 6]")));
EXPECT_TRUE(sliced_list_array->values()->Equals(
ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
}
protected:
std::shared_ptr<DataType> value_type_;
@@ -378,6 +406,9 @@ TYPED_TEST(TestListArray, BuilderPreserveFieldName) {
this->TestBuilderPreserveFieldName();
}
TYPED_TEST(TestListArray, Flatten) { this->TestFlatten(); }
TYPED_TEST(TestListArray, FlattenZeroLength) { this->TestFlattenZeroLength(); }
// ----------------------------------------------------------------------
// Map tests