Skip to content

Commit 8eea173

Browse files
emkornfieldpitrou
emkornfield
and
pitrou
committed
ARROW-7960: [C++] Add support fo reading additional types
New types supported: - Fixed Size list (will throw an incorrect error if nulls). Null slot count is best implemented after we recursively apply metadata. - LargeList - Maps Still missing: LargeString and LargeBytes Unimplemented functionality: - Removing duplicate maps on read. - Single column maps are converted to List of struct. Other: Fixed two bugs for FixedSizeLists: 1. Def levels genereated were incorrect (they only made sense for two level lists). 2. Slices where not handled appropriately Closes apache#8376 from emkornfield/read_most_types Lead-authored-by: Micah Kornfield <emkornfield@gmail.com> Co-authored-by: Antoine Pitrou <antoine@python.org> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent 6a8f84d commit 8eea173

7 files changed

Lines changed: 344 additions & 12 deletions

File tree

‎cpp/src/parquet/arrow/arrow_reader_writer_test.cc‎

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2469,7 +2469,60 @@ TEST(ArrowReadWrite, ListOfStructOfList1) {
24692469
CheckSimpleRoundtrip(table, 2);
24702470
}
24712471

2472-
TEST(ArrowReadWrite, DISABLED_ListOfStructOfList2) {
2472+
TEST(ArrowReadWrite, Map) {
2473+
using ::arrow::field;
2474+
using ::arrow::map;
2475+
2476+
auto type = map(::arrow::int16(), ::arrow::utf8());
2477+
2478+
const char* json = R"([
2479+
[[1, "a"], [2, "b"]],
2480+
[[3, "c"]],
2481+
[],
2482+
null,
2483+
[[4, "d"], [5, "e"], [6, "f"]]
2484+
])";
2485+
auto array = ::arrow::ArrayFromJSON(type, json);
2486+
auto table = ::arrow::Table::Make(::arrow::schema({field("root", type)}), {array});
2487+
auto props_store_schema = ArrowWriterProperties::Builder().store_schema()->build();
2488+
CheckSimpleRoundtrip(table, 2, props_store_schema);
2489+
}
2490+
2491+
TEST(ArrowReadWrite, LargeList) {
2492+
using ::arrow::field;
2493+
using ::arrow::large_list;
2494+
using ::arrow::struct_;
2495+
2496+
auto type = large_list(::arrow::int16());
2497+
2498+
const char* json = R"([
2499+
[1, 2, 3],
2500+
[4, 5, 6],
2501+
[7, 8, 9]])";
2502+
auto array = ::arrow::ArrayFromJSON(type, json);
2503+
auto table = ::arrow::Table::Make(::arrow::schema({field("root", type)}), {array});
2504+
auto props_store_schema = ArrowWriterProperties::Builder().store_schema()->build();
2505+
CheckSimpleRoundtrip(table, 2, props_store_schema);
2506+
}
2507+
2508+
TEST(ArrowReadWrite, FixedSizeList) {
2509+
using ::arrow::field;
2510+
using ::arrow::fixed_size_list;
2511+
using ::arrow::struct_;
2512+
2513+
auto type = fixed_size_list(::arrow::int16(), /*size=*/3);
2514+
2515+
const char* json = R"([
2516+
[1, 2, 3],
2517+
[4, 5, 6],
2518+
[7, 8, 9]])";
2519+
auto array = ::arrow::ArrayFromJSON(type, json);
2520+
auto table = ::arrow::Table::Make(::arrow::schema({field("root", type)}), {array});
2521+
auto props_store_schema = ArrowWriterProperties::Builder().store_schema()->build();
2522+
CheckSimpleRoundtrip(table, 2, props_store_schema);
2523+
}
2524+
2525+
TEST(ArrowReadWrite, ListOfStructOfList2) {
24732526
using ::arrow::field;
24742527
using ::arrow::list;
24752528
using ::arrow::struct_;
@@ -2484,7 +2537,7 @@ TEST(ArrowReadWrite, DISABLED_ListOfStructOfList2) {
24842537
[{"a": 123, "b": [1, 2, 3]}],
24852538
null,
24862539
[],
2487-
[{"a": 456}, {"a": 789, "b": [null]}, {"a": 876, "b": [4, 5, 6]}]])";
2540+
[{"a": 456, "b": []}, {"a": 789, "b": [null]}, {"a": 876, "b": [4, 5, 6]}]])";
24882541
auto array = ::arrow::ArrayFromJSON(type, json);
24892542
auto table = ::arrow::Table::Make(::arrow::schema({field("root", type)}), {array});
24902543
CheckSimpleRoundtrip(table, 2);

‎cpp/src/parquet/arrow/arrow_schema_test.cc‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,44 @@ TEST_F(TestConvertParquetSchema, ParquetFlatDecimals) {
346346
ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema));
347347
}
348348

349+
TEST_F(TestConvertParquetSchema, ParquetMaps) {
350+
std::vector<NodePtr> parquet_fields;
351+
std::vector<std::shared_ptr<Field>> arrow_fields;
352+
353+
// MAP encoding example taken from parquet-format/LogicalTypes.md
354+
355+
// Two column map.
356+
{
357+
auto key = PrimitiveNode::Make("key", Repetition::REQUIRED, ParquetType::BYTE_ARRAY,
358+
ConvertedType::UTF8);
359+
auto value = PrimitiveNode::Make("value", Repetition::OPTIONAL,
360+
ParquetType::BYTE_ARRAY, ConvertedType::UTF8);
361+
362+
auto list = GroupNode::Make("key_value", Repetition::REPEATED, {key, value});
363+
parquet_fields.push_back(
364+
GroupNode::Make("my_map", Repetition::REQUIRED, {list}, LogicalType::Map()));
365+
auto arrow_value = ::arrow::field("string", UTF8, /*nullable=*/true);
366+
auto arrow_map = ::arrow::map(/*key=*/UTF8, arrow_value);
367+
arrow_fields.push_back(::arrow::field("my_map", arrow_map, false));
368+
}
369+
// Single column map (i.e. set) gets converted to list of struct.
370+
{
371+
auto key = PrimitiveNode::Make("key", Repetition::REQUIRED, ParquetType::BYTE_ARRAY,
372+
ConvertedType::UTF8);
373+
374+
auto list = GroupNode::Make("key_value", Repetition::REPEATED, {key});
375+
parquet_fields.push_back(
376+
GroupNode::Make("my_set", Repetition::REQUIRED, {list}, LogicalType::Map()));
377+
auto arrow_list = ::arrow::list({::arrow::field("key", UTF8, /*nullable=*/false)});
378+
arrow_fields.push_back(::arrow::field("my_set", arrow_list, false));
379+
}
380+
381+
auto arrow_schema = ::arrow::schema(arrow_fields);
382+
ASSERT_OK(ConvertSchema(parquet_fields));
383+
384+
ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema));
385+
}
386+
349387
TEST_F(TestConvertParquetSchema, ParquetLists) {
350388
std::vector<NodePtr> parquet_fields;
351389
std::vector<std::shared_ptr<Field>> arrow_fields;
@@ -1217,6 +1255,54 @@ TEST_F(TestLevels, TestPrimitive) {
12171255
/*ancestor_list_def_level*/ 1})); // primitive field
12181256
}
12191257

1258+
TEST_F(TestLevels, TestMaps) {
1259+
// Two column map.
1260+
auto key = PrimitiveNode::Make("key", Repetition::REQUIRED, ParquetType::BYTE_ARRAY,
1261+
ConvertedType::UTF8);
1262+
auto value = PrimitiveNode::Make("value", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY,
1263+
ConvertedType::UTF8);
1264+
1265+
auto list = GroupNode::Make("key_value", Repetition::REPEATED, {key, value});
1266+
SetParquetSchema(
1267+
GroupNode::Make("my_map", Repetition::OPTIONAL, {list}, LogicalType::Map()));
1268+
ASSERT_OK_AND_ASSIGN(std::deque<LevelInfo> levels,
1269+
RootToTreeLeafLevels(*manifest_, /*column_number=*/0));
1270+
EXPECT_THAT(
1271+
levels,
1272+
ElementsAre(LevelInfo{/*null_slot_usage=*/1, /*def_level=*/2, /*rep_level=*/1,
1273+
/*ancestor_list_def_level*/ 0},
1274+
LevelInfo{/*null_slot_usage=*/1, /*def_level=*/2, /*rep_level=*/1,
1275+
/*ancestor_list_def_level*/ 2},
1276+
LevelInfo{/*null_slot_usage=*/1, /*def_level=*/2, /*rep_level=*/1,
1277+
/*ancestor_list_def_level*/ 2}));
1278+
1279+
ASSERT_OK_AND_ASSIGN(levels, RootToTreeLeafLevels(*manifest_, /*column_number=*/1));
1280+
EXPECT_THAT(
1281+
levels,
1282+
ElementsAre(LevelInfo{/*null_slot_usage=*/1, /*def_level=*/2, /*rep_level=*/1,
1283+
/*ancestor_list_def_level*/ 0},
1284+
LevelInfo{/*null_slot_usage=*/1, /*def_level=*/2, /*rep_level=*/1,
1285+
/*ancestor_list_def_level*/ 2},
1286+
LevelInfo{/*null_slot_usage=*/1, /*def_level=*/3, /*rep_level=*/1,
1287+
/*ancestor_list_def_level*/ 2}));
1288+
1289+
// single column map.
1290+
key = PrimitiveNode::Make("key", Repetition::REQUIRED, ParquetType::BYTE_ARRAY,
1291+
ConvertedType::UTF8);
1292+
1293+
list = GroupNode::Make("key_value", Repetition::REPEATED, {key});
1294+
SetParquetSchema(
1295+
GroupNode::Make("my_set", Repetition::REQUIRED, {list}, LogicalType::Map()));
1296+
1297+
ASSERT_OK_AND_ASSIGN(levels, RootToTreeLeafLevels(*manifest_, /*column_number=*/0));
1298+
EXPECT_THAT(
1299+
levels,
1300+
ElementsAre(LevelInfo{/*null_slot_usage=*/1, /*def_level=*/1, /*rep_level=*/1,
1301+
/*ancestor_list_def_level*/ 0},
1302+
LevelInfo{/*null_slot_usage=*/1, /*def_level=*/1, /*rep_level=*/1,
1303+
/*ancestor_list_def_level*/ 1}));
1304+
}
1305+
12201306
TEST_F(TestLevels, TestSimpleGroups) {
12211307
// Arrow schema: struct(child: struct(inner: boolean not null))
12221308
SetParquetSchema(GroupNode::Make(

‎cpp/src/parquet/arrow/path_internal.cc‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -815,13 +815,17 @@ class PathBuilder {
815815
Status Visit(const ::arrow::FixedSizeListArray& array) {
816816
MaybeAddNullable(array);
817817
int32_t list_size = array.list_type()->list_size();
818-
if (list_size == 0) {
819-
info_.max_def_level++;
820-
}
818+
// Technically we could encode fixed size lists with two level encodings
819+
// but since we always use 3 level encoding we increment def levels as
820+
// well.
821+
info_.max_def_level++;
821822
info_.max_rep_level++;
822823
info_.path.push_back(FixedSizeListNode(FixedSizedRangeSelector{list_size},
823824
info_.max_rep_level, info_.max_def_level));
824825
nullable_in_parent_ = array.list_type()->value_field()->nullable();
826+
if (array.offset() > 0) {
827+
return VisitInline(*array.values()->Slice(array.value_offset(0)));
828+
}
825829
return VisitInline(*array.values());
826830
}
827831

‎cpp/src/parquet/arrow/path_internal_test.cc‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,26 @@ TEST_F(MultipathLevelBuilderTest, TestStruct) {
461461
/*def_levels=*/std::vector<int16_t>({2, 2, 0}));
462462
}
463463

464+
TEST_F(MultipathLevelBuilderTest, TestFixedSizeListNullableElements) {
465+
auto entries = field("Entries", ::arrow::int64());
466+
auto list_type = fixed_size_list(entries, 2);
467+
auto array = ::arrow::ArrayFromJSON(list_type, R"([null, [2, 3], [4, 5], null])");
468+
469+
ASSERT_OK(
470+
MultipathLevelBuilder::Write(*array, /*nullable=*/true, &context_, callback_));
471+
472+
ASSERT_THAT(results_, SizeIs(1));
473+
results_[0].CheckLevels(/*def_levels=*/std::vector<int16_t>{0, 3, 3, 3, 3, 0},
474+
/*rep_levels=*/std::vector<int16_t>{0, 0, 1, 0, 1, 0});
475+
476+
// Null slots take up space in a fixed size list (they can in variable size
477+
// lists as well) but the actual written values are only the "middle" elements
478+
// in this case.
479+
ASSERT_THAT(results_[0].post_list_elements, SizeIs(1));
480+
EXPECT_THAT(results_[0].post_list_elements[0].start, Eq(2));
481+
EXPECT_THAT(results_[0].post_list_elements[0].end, Eq(6));
482+
}
483+
464484
TEST_F(MultipathLevelBuilderTest, TestFixedSizeList) {
465485
auto entries = field("Entries", ::arrow::int64(), /*nullable=*/false);
466486
auto list_type = fixed_size_list(entries, 2);
@@ -470,7 +490,7 @@ TEST_F(MultipathLevelBuilderTest, TestFixedSizeList) {
470490
MultipathLevelBuilder::Write(*array, /*nullable=*/true, &context_, callback_));
471491

472492
ASSERT_THAT(results_, SizeIs(1));
473-
results_[0].CheckLevels(/*def_levels=*/std::vector<int16_t>{0, 1, 1, 1, 1, 0},
493+
results_[0].CheckLevels(/*def_levels=*/std::vector<int16_t>{0, 2, 2, 2, 2, 0},
474494
/*rep_levels=*/std::vector<int16_t>{0, 0, 1, 0, 1, 0});
475495

476496
// Null slots take up space in a fixed size list (they can in variable size

‎cpp/src/parquet/arrow/reader.cc‎

Lines changed: 59 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -538,13 +538,20 @@ class ListReader : public ColumnReaderImpl {
538538
return item_reader_->LoadBatch(number_of_records);
539539
}
540540

541+
virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
542+
std::shared_ptr<ArrayData> data) {
543+
std::shared_ptr<Array> result = ::arrow::MakeArray(data);
544+
return std::make_shared<ChunkedArray>(result);
545+
}
546+
541547
Status BuildArray(int64_t length_upper_bound,
542548
std::shared_ptr<ChunkedArray>* out) override {
543549
const int16_t* def_levels;
544550
const int16_t* rep_levels;
545551
int64_t num_levels;
546552
RETURN_NOT_OK(item_reader_->GetDefLevels(&def_levels, &num_levels));
547553
RETURN_NOT_OK(item_reader_->GetRepLevels(&rep_levels, &num_levels));
554+
548555
std::shared_ptr<ResizableBuffer> validity_buffer;
549556
::parquet::internal::ValidityBitmapInputOutput validity_io;
550557
validity_io.values_read_upper_bound = length_upper_bound;
@@ -576,6 +583,7 @@ class ListReader : public ColumnReaderImpl {
576583
if (validity_buffer != nullptr) {
577584
RETURN_NOT_OK(
578585
validity_buffer->Resize(BitUtil::BytesForBits(validity_io.values_read)));
586+
validity_buffer->ZeroPadding();
579587
}
580588
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrayData> item_chunk, ChunksToSingle(**out));
581589

@@ -586,8 +594,7 @@ class ListReader : public ColumnReaderImpl {
586594
/*length=*/validity_io.values_read, std::move(buffers),
587595
std::vector<std::shared_ptr<ArrayData>>{item_chunk}, validity_io.null_count);
588596

589-
std::shared_ptr<Array> result = ::arrow::MakeArray(data);
590-
*out = std::make_shared<ChunkedArray>(result);
597+
ARROW_ASSIGN_OR_RAISE(*out, AssembleArray(std::move(data)));
591598
return Status::OK();
592599
}
593600

@@ -600,6 +607,32 @@ class ListReader : public ColumnReaderImpl {
600607
std::unique_ptr<ColumnReaderImpl> item_reader_;
601608
};
602609

610+
class PARQUET_NO_EXPORT FixedSizeListReader : public ListReader<int32_t> {
611+
public:
612+
FixedSizeListReader(std::shared_ptr<ReaderContext> ctx, std::shared_ptr<Field> field,
613+
::parquet::internal::LevelInfo level_info,
614+
std::unique_ptr<ColumnReaderImpl> child_reader)
615+
: ListReader(std::move(ctx), std::move(field), level_info,
616+
std::move(child_reader)) {}
617+
::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
618+
std::shared_ptr<ArrayData> data) final {
619+
DCHECK_EQ(data->buffers.size(), 2);
620+
DCHECK_EQ(field()->type()->id(), ::arrow::Type::FIXED_SIZE_LIST);
621+
const auto& type = checked_cast<::arrow::FixedSizeListType&>(*field()->type());
622+
const int32_t* offsets = reinterpret_cast<const int32_t*>(data->buffers[1]->data());
623+
for (int x = 1; x <= data->length; x++) {
624+
int32_t size = offsets[x] - offsets[x - 1];
625+
if (size != type.list_size()) {
626+
return Status::Invalid("Expected all lists to be of size=", type.list_size(),
627+
" but index ", x, " had size=", size);
628+
}
629+
}
630+
data->buffers.resize(1);
631+
std::shared_ptr<Array> result = ::arrow::MakeArray(data);
632+
return std::make_shared<ChunkedArray>(result);
633+
}
634+
};
635+
603636
class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl {
604637
public:
605638
explicit StructReader(std::shared_ptr<ReaderContext> ctx,
@@ -714,6 +747,7 @@ Status StructReader::BuildArray(int64_t length_upper_bound,
714747
// Ensure all values are initialized.
715748
if (null_bitmap) {
716749
RETURN_NOT_OK(null_bitmap->Resize(BitUtil::BytesForBits(validity_io.values_read)));
750+
null_bitmap->ZeroPadding();
717751
}
718752

719753
END_PARQUET_CATCH_EXCEPTIONS
@@ -770,16 +804,36 @@ Status GetReader(const SchemaField& field, const std::shared_ptr<Field>& arrow_f
770804
std::unique_ptr<FileColumnIterator> input(
771805
ctx->iterator_factory(field.column_index, ctx->reader));
772806
out->reset(new LeafReader(ctx, arrow_field, std::move(input), field.level_info));
773-
} else if (type_id == ::arrow::Type::LIST) {
807+
} else if (type_id == ::arrow::Type::LIST || type_id == ::arrow::Type::MAP ||
808+
type_id == ::arrow::Type::FIXED_SIZE_LIST ||
809+
type_id == ::arrow::Type::LARGE_LIST) {
810+
auto list_field = arrow_field;
774811
auto child = &field.children[0];
775812
std::unique_ptr<ColumnReaderImpl> child_reader;
776813
RETURN_NOT_OK(GetReader(*child, ctx, &child_reader));
777814
if (child_reader == nullptr) {
778815
*out = nullptr;
779816
return Status::OK();
780817
}
781-
out->reset(new ListReader<int32_t>(ctx, arrow_field, field.level_info,
782-
std::move(child_reader)));
818+
if (type_id == ::arrow::Type::LIST ||
819+
type_id == ::arrow::Type::MAP) { // Map can be reconstructed as list of structs.
820+
if (type_id == ::arrow::Type::MAP &&
821+
child_reader->field()->type()->num_fields() != 2) {
822+
// This case applies if either key or value is filtered.
823+
list_field = list_field->WithType(::arrow::list(child_reader->field()));
824+
}
825+
out->reset(new ListReader<int32_t>(ctx, list_field, field.level_info,
826+
std::move(child_reader)));
827+
} else if (type_id == ::arrow::Type::LARGE_LIST) {
828+
out->reset(new ListReader<int64_t>(ctx, list_field, field.level_info,
829+
std::move(child_reader)));
830+
831+
} else if (type_id == ::arrow::Type::FIXED_SIZE_LIST) {
832+
out->reset(new FixedSizeListReader(ctx, list_field, field.level_info,
833+
std::move(child_reader)));
834+
} else {
835+
return Status::UnknownError("Unknown list type: ", field.field->ToString());
836+
}
783837
} else if (type_id == ::arrow::Type::STRUCT) {
784838
std::vector<std::shared_ptr<Field>> child_fields;
785839
std::vector<std::unique_ptr<ColumnReaderImpl>> child_readers;

0 commit comments

Comments
 (0)