-
Notifications
You must be signed in to change notification settings - Fork 120
feat(parquet): support reading list columns as Arrow large_list #714
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
ebecef5
430a53c
af0c94f
a5a2d5b
9e960a2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,6 +19,7 @@ | |
|
|
||
| #include "iceberg/parquet/parquet_reader.h" | ||
|
|
||
| #include <algorithm> | ||
| #include <numeric> | ||
|
|
||
| #include <arrow/c/bridge.h> | ||
|
|
@@ -41,6 +42,7 @@ | |
| #include "iceberg/result.h" | ||
| #include "iceberg/schema_internal.h" | ||
| #include "iceberg/schema_util.h" | ||
| #include "iceberg/util/checked_cast.h" | ||
| #include "iceberg/util/macros.h" | ||
|
|
||
| namespace iceberg::parquet { | ||
|
|
@@ -84,6 +86,78 @@ class EmptyRecordBatchReader : public ::arrow::RecordBatchReader { | |
| } | ||
| }; | ||
|
|
||
| // forward declaration to unblock cycle dependence. | ||
| std::shared_ptr<::arrow::Field> UseLargeListField( | ||
| const std::shared_ptr<::arrow::Field>& field); | ||
|
|
||
| // Rebuild a data type with all nested list types replaced by large_list. | ||
| std::shared_ptr<::arrow::DataType> UseLargeListType( | ||
| const std::shared_ptr<::arrow::DataType>& type) { | ||
| switch (type->id()) { | ||
| case ::arrow::Type::LIST: { | ||
| const auto& list_type = internal::checked_cast<const ::arrow::ListType&>(*type); | ||
| return ::arrow::large_list(UseLargeListField(list_type.value_field())); | ||
| } | ||
| case ::arrow::Type::STRUCT: { | ||
| ::arrow::FieldVector fields; | ||
| fields.reserve(type->num_fields()); | ||
| for (const auto& field : type->fields()) { | ||
| fields.push_back(UseLargeListField(field)); | ||
| } | ||
| return ::arrow::struct_(std::move(fields)); | ||
| } | ||
| case ::arrow::Type::MAP: { | ||
| const auto& map_type = internal::checked_cast<const ::arrow::MapType&>(*type); | ||
| return std::make_shared<::arrow::MapType>(UseLargeListField(map_type.key_field()), | ||
| UseLargeListField(map_type.item_field()), | ||
| map_type.keys_sorted()); | ||
| } | ||
| default: | ||
| return type; | ||
| } | ||
| } | ||
|
|
||
| std::shared_ptr<::arrow::Field> UseLargeListField( | ||
| const std::shared_ptr<::arrow::Field>& field) { | ||
| return field->WithType(UseLargeListType(field->type())); | ||
| } | ||
|
|
||
| // Rewrite all fields in a field vector to use large_list instead of list. | ||
| ::arrow::FieldVector UseLargeListFields(const ::arrow::FieldVector& fields) { | ||
| ::arrow::FieldVector rewritten; | ||
| rewritten.reserve(fields.size()); | ||
| for (const auto& field : fields) { | ||
| rewritten.push_back(UseLargeListField(field)); | ||
| } | ||
| return rewritten; | ||
| } | ||
|
|
||
| // Returns true if the type contains a large_list, at any level of nesting. | ||
| bool ContainsLargeList(const ::arrow::DataType& type) { | ||
| if (type.id() == ::arrow::Type::LARGE_LIST) { | ||
| return true; | ||
| } | ||
| return std::ranges::any_of( | ||
| type.fields(), [](const auto& field) { return ContainsLargeList(*field->type()); }); | ||
| } | ||
|
|
||
| // Returns true if the reader produces large_list arrays. | ||
| // | ||
| // Arrow honors the requested large_list type only when it derives the Arrow schema from | ||
| // the Parquet schema. A file that carries serialized ARROW:schema metadata keeps its | ||
| // original list type instead, so whether large lists are produced can only be told from | ||
| // the schema of the reader. | ||
| bool ProducesLargeList(const ::arrow::RecordBatchReader& reader) { | ||
| const auto& schema = reader.schema(); | ||
| if (schema == nullptr) { | ||
| // an empty reader produces no arrays to be described | ||
| return false; | ||
| } | ||
| return std::ranges::any_of(schema->fields(), [](const auto& field) { | ||
| return ContainsLargeList(*field->type()); | ||
| }); | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| // A stateful context to keep track of the reading progress. | ||
|
|
@@ -118,6 +192,10 @@ class ParquetReader::Impl { | |
| arrow_reader_properties.set_batch_size( | ||
| options.properties.Get(ReaderProperties::kBatchSize)); | ||
| arrow_reader_properties.set_arrow_extensions_enabled(true); | ||
| use_large_list_ = options.properties.Get(ReaderProperties::kArrowUseLargeList); | ||
| if (use_large_list_) { | ||
| arrow_reader_properties.set_list_type(::arrow::Type::LARGE_LIST); | ||
| } | ||
|
|
||
| // Open the Parquet file reader | ||
| ICEBERG_ASSIGN_OR_RAISE(input_stream_, OpenInputStream(options)); | ||
|
|
@@ -212,12 +290,6 @@ class ParquetReader::Impl { | |
| Status InitReadContext() { | ||
| context_ = std::make_unique<ReadContext>(); | ||
|
|
||
| // Build the output Arrow schema | ||
| ArrowSchema arrow_schema; | ||
| ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*read_schema_, &arrow_schema)); | ||
| ICEBERG_ARROW_ASSIGN_OR_RETURN(context_->output_arrow_schema_, | ||
| ::arrow::ImportSchema(&arrow_schema)); | ||
|
|
||
| // Row group pruning based on the split | ||
| // TODO(gangwu): add row group filtering based on zone map, bloom filter, etc. | ||
| std::vector<int> row_group_indices; | ||
|
|
@@ -250,6 +322,24 @@ class ParquetReader::Impl { | |
| reader_->GetRecordBatchReader(row_group_indices, column_indices)); | ||
| } | ||
|
|
||
| // Build the output Arrow schema from the projected Iceberg schema. This schema is the | ||
| // target of ProjectRecordBatch, so it must describe the projected schema rather than | ||
| // the schema of the file. | ||
| ArrowSchema arrow_schema; | ||
| ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*read_schema_, &arrow_schema)); | ||
| ICEBERG_ARROW_ASSIGN_OR_RETURN(context_->output_arrow_schema_, | ||
| ::arrow::ImportSchema(&arrow_schema)); | ||
|
|
||
| if (use_large_list_ && ProducesLargeList(*context_->record_batch_reader_)) { | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please align the output schema with the reader schema even when this option is false. Arrow restores a stored |
||
| // Align the output schema with the large_list arrays produced by the Parquet | ||
| // reader. Note that Arrow ignores the requested list type when the file carries | ||
| // serialized ARROW:schema metadata, in which case the reader keeps producing plain | ||
| // list arrays and the output schema must keep describing them as such. | ||
| context_->output_arrow_schema_ = | ||
| ::arrow::schema(UseLargeListFields(context_->output_arrow_schema_->fields()), | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This rewrite is global: one |
||
| context_->output_arrow_schema_->metadata()); | ||
| } | ||
|
|
||
| return {}; | ||
| } | ||
|
|
||
|
|
@@ -258,6 +348,8 @@ class ParquetReader::Impl { | |
| ::arrow::MemoryPool* pool_ = ::arrow::default_memory_pool(); | ||
| // The split to read from the Parquet file. | ||
| std::optional<Split> split_; | ||
| // Whether to read list columns as large_list (64-bit offsets). | ||
| bool use_large_list_ = false; | ||
| // Schema to read from the Parquet file. | ||
| std::shared_ptr<::iceberg::Schema> read_schema_; | ||
| // The projection result to apply to the read schema. | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
add comment: