diff --git a/src/iceberg/file_reader.h b/src/iceberg/file_reader.h index cefc688c0..d9424026e 100644 --- a/src/iceberg/file_reader.h +++ b/src/iceberg/file_reader.h @@ -76,6 +76,9 @@ class ICEBERG_EXPORT ReaderProperties : public ConfigBase { /// \brief The batch size to read. inline static Entry kBatchSize{"read.batch-size", 4096}; + /// \brief Read list columns as Arrow large_list (64-bit offsets) instead of list. + /// Default: false (use 32-bit offset list). + inline static Entry kArrowUseLargeList{"read.arrow.use-large-list", false}; /// \brief Skip GenericDatum in Avro reader for better performance. /// When true, decode directly from Avro to Arrow without GenericDatum intermediate. /// Default: true (skip GenericDatum for better performance). diff --git a/src/iceberg/parquet/parquet_reader.cc b/src/iceberg/parquet/parquet_reader.cc index f0df07ac2..60fc11122 100644 --- a/src/iceberg/parquet/parquet_reader.cc +++ b/src/iceberg/parquet/parquet_reader.cc @@ -19,6 +19,7 @@ #include "iceberg/parquet/parquet_reader.h" +#include #include #include @@ -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(*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(*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(); - // 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 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_)) { + // 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()), + 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_; + // 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. diff --git a/src/iceberg/parquet/parquet_schema_util.cc b/src/iceberg/parquet/parquet_schema_util.cc index 658880814..fde6b24c7 100644 --- a/src/iceberg/parquet/parquet_schema_util.cc +++ b/src/iceberg/parquet/parquet_schema_util.cc @@ -251,7 +251,8 @@ Status ValidateParquetSchemaEvolution( } break; case TypeId::kList: - if (arrow_type->id() == ::arrow::Type::LIST) { + if (arrow_type->id() == ::arrow::Type::LIST || + arrow_type->id() == ::arrow::Type::LARGE_LIST) { return {}; } break; diff --git a/src/iceberg/test/parquet_schema_test.cc b/src/iceberg/test/parquet_schema_test.cc index 75e99ff12..46f50e1e1 100644 --- a/src/iceberg/test/parquet_schema_test.cc +++ b/src/iceberg/test/parquet_schema_test.cc @@ -111,12 +111,14 @@ ::parquet::schema::NodePtr MakeMapNode(const std::string& name, // Helper to create SchemaManifest from Parquet schema ::parquet::arrow::SchemaManifest MakeSchemaManifest( - const ::parquet::schema::NodePtr& parquet_schema) { + const ::parquet::schema::NodePtr& parquet_schema, + ::arrow::Type::type list_type = ::arrow::Type::LIST) { auto parquet_schema_descriptor = std::make_shared<::parquet::SchemaDescriptor>(); parquet_schema_descriptor->Init(parquet_schema); auto properties = ::parquet::default_arrow_reader_properties(); properties.set_arrow_extensions_enabled(true); + properties.set_list_type(list_type); ::parquet::arrow::SchemaManifest manifest; auto status = ::parquet::arrow::SchemaManifest::Make(parquet_schema_descriptor.get(), @@ -340,6 +342,16 @@ TEST(ParquetSchemaProjectionTest, ValidateSchemaEvolutionAllowsNullPhysicalType) ASSERT_THAT(status, IsOk()); } +TEST(ParquetSchemaProjectionTest, ValidateSchemaEvolutionAllowsLargeList) { + ::parquet::arrow::SchemaField parquet_field; + parquet_field.field = ::arrow::field("numbers", ::arrow::large_list(::arrow::int32())); + + ListType expected_type( + SchemaField::MakeOptional(/*field_id=*/101, "element", iceberg::int32())); + auto status = ValidateParquetSchemaEvolution(expected_type, parquet_field); + ASSERT_THAT(status, IsOk()); +} + TEST(ParquetSchemaProjectionTest, ProjectNullPhysicalFieldsAsNull) { Schema expected_schema({ SchemaField::MakeOptional(/*field_id=*/1, "age", iceberg::int32()), @@ -624,6 +636,38 @@ TEST(ParquetSchemaProjectionTest, ProjectListType) { ASSERT_EQ(SelectedColumnIndices(projection), std::vector({0, 1})); } +TEST(ParquetSchemaProjectionTest, ProjectLargeListType) { + Schema expected_schema({ + SchemaField::MakeOptional( + /*field_id=*/2, "numbers", + std::make_shared(SchemaField::MakeOptional( + /*field_id=*/101, "element", iceberg::int32()))), + }); + + auto parquet_schema = MakeGroupNode( + "iceberg_schema", + { + MakeListNode("numbers", MakeInt32Node("element", /*field_id=*/101), + /*field_id=*/2), + }); + + auto schema_manifest = MakeSchemaManifest(parquet_schema, ::arrow::Type::LARGE_LIST); + ASSERT_EQ(schema_manifest.schema_fields[0].field->type()->id(), + ::arrow::Type::LARGE_LIST); + + auto projection_result = Project(expected_schema, schema_manifest); + ASSERT_THAT(projection_result, IsOk()); + + const auto& projection = *projection_result; + ASSERT_EQ(projection.fields.size(), 1); + ASSERT_PROJECTED_FIELD(projection.fields[0], 0); + + ASSERT_EQ(projection.fields[0].children.size(), 1); + ASSERT_PROJECTED_FIELD(projection.fields[0].children[0], 0); + + ASSERT_EQ(SelectedColumnIndices(projection), std::vector({0})); +} + TEST(ParquetSchemaProjectionTest, ProjectMapType) { Schema expected_schema({ SchemaField::MakeOptional( diff --git a/src/iceberg/test/parquet_test.cc b/src/iceberg/test/parquet_test.cc index 6ccff6cb4..4e11a9fb0 100644 --- a/src/iceberg/test/parquet_test.cc +++ b/src/iceberg/test/parquet_test.cc @@ -198,6 +198,32 @@ class ParquetReaderTest : public TempFileTestBase { .properties = std::move(writer_properties)})); } + void CreateListParquetFile() { + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "numbers", + std::make_shared(SchemaField::MakeOptional( + /*field_id=*/101, "element", int32())))}); + + ArrowSchema arrow_c_schema; + ASSERT_THAT(ToArrowSchema(*schema, &arrow_c_schema), IsOk()); + auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); + + auto array = + ::arrow::json::ArrayFromJSONString(::arrow::struct_(arrow_schema->fields()), + R"([[1, [1, 2]], [2, [3]], [3, null]])") + .ValueOrDie(); + + WriterProperties writer_properties; + writer_properties.Set(WriterProperties::kParquetCompression, + std::string("uncompressed")); + + ASSERT_TRUE(WriteArray(array, {.path = temp_parquet_file_, + .schema = schema, + .io = file_io_, + .properties = std::move(writer_properties)})); + } + void CreateRowLineageParquetFile() { auto schema = RowLineageSchema(); @@ -267,6 +293,38 @@ class ParquetReaderTest : public TempFileTestBase { .data_sequence_number = data_sequence_number}); } + // Writes a list parquet file through parquet::arrow::WriteTable, which serializes the + // Arrow schema of the table into the ARROW:schema key value metadata of the file. + void CreateListParquetFileWithArrowSchema() { + const std::string kParquetFieldIdKey = "PARQUET:field_id"; + auto arrow_schema = ::arrow::schema( + {::arrow::field("id", ::arrow::int32(), /*nullable=*/false, + ::arrow::KeyValueMetadata::Make({kParquetFieldIdKey}, {"1"})), + ::arrow::field( + "numbers", + ::arrow::list(::arrow::field( + "element", ::arrow::int32(), /*nullable=*/true, + ::arrow::KeyValueMetadata::Make({kParquetFieldIdKey}, {"101"}))), + /*nullable=*/true, + ::arrow::KeyValueMetadata::Make({kParquetFieldIdKey}, {"2"}))}); + auto batch = + ::arrow::RecordBatch::FromStructArray( + ::arrow::json::ArrayFromJSONString(::arrow::struct_(arrow_schema->fields()), + R"([[1, [1, 2]], [2, [3]]])") + .ValueOrDie()) + .ValueOrDie(); + auto table = ::arrow::Table::FromRecordBatches(arrow_schema, {batch}).ValueOrDie(); + + auto io = internal::checked_cast(*file_io_); + auto outfile = io.fs()->OpenOutputStream(temp_parquet_file_).ValueOrDie(); + + // write a single row group so that one batch holds every row + ASSERT_TRUE(::parquet::arrow::WriteTable(*table, ::arrow::default_memory_pool(), + outfile, table->num_rows()) + .ok()); + ASSERT_TRUE(outfile->Close().ok()); + } + void VerifyNextBatch(Reader& reader, std::string_view expected_json) { // Boilerplate to get Arrow schema auto schema_result = reader.Schema(); @@ -446,6 +504,156 @@ TEST_F(ParquetReaderTest, ReadWithBatchSize) { ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); } +TEST_F(ParquetReaderTest, ReadListType) { + CreateListParquetFile(); + + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "numbers", + std::make_shared(SchemaField::MakeOptional( + /*field_id=*/101, "element", int32())))}); + + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kParquet, + {.path = temp_parquet_file_, .io = file_io_, .projection = schema}); + ASSERT_THAT(reader_result, IsOk()); + auto reader = std::move(reader_result.value()); + + // By default list columns are read as 32-bit offset list arrays. + auto schema_result = reader->Schema(); + ASSERT_THAT(schema_result, IsOk()); + auto arrow_c_schema = std::move(schema_result.value()); + auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); + ASSERT_EQ(arrow_type->field(1)->type()->id(), ::arrow::Type::LIST); + + ASSERT_NO_FATAL_FAILURE( + VerifyNextBatch(*reader, R"([[1, [1, 2]], [2, [3]], [3, null]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); +} + +TEST_F(ParquetReaderTest, ReadListAsLargeList) { + CreateListParquetFile(); + + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "numbers", + std::make_shared(SchemaField::MakeOptional( + /*field_id=*/101, "element", int32())))}); + + ReaderProperties reader_properties; + reader_properties.Set(ReaderProperties::kArrowUseLargeList, true); + + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kParquet, {.path = temp_parquet_file_, + .io = file_io_, + .projection = schema, + .properties = std::move(reader_properties)}); + ASSERT_THAT(reader_result, IsOk()); + auto reader = std::move(reader_result.value()); + + // The output schema should expose list columns as large_list. + auto schema_result = reader->Schema(); + ASSERT_THAT(schema_result, IsOk()); + auto arrow_c_schema = std::move(schema_result.value()); + auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); + ASSERT_EQ(arrow_type->field(1)->type()->id(), ::arrow::Type::LARGE_LIST); + + // JSON parsing creates regular ListArray, so verify large_list data manually. + auto data = reader->Next(); + ASSERT_THAT(data, IsOk()); + ASSERT_TRUE(data.value().has_value()); + auto arrow_c_array = data.value().value(); + auto arrow_array = ::arrow::ImportArray(&arrow_c_array, arrow_type).ValueOrDie(); + + const auto& struct_array = + internal::checked_cast(*arrow_array); + ASSERT_EQ(struct_array.length(), 3); + + const auto& id_array = + internal::checked_cast(*struct_array.field(0)); + ASSERT_EQ(id_array.Value(0), 1); + ASSERT_EQ(id_array.Value(1), 2); + ASSERT_EQ(id_array.Value(2), 3); + + const auto& numbers_array = + internal::checked_cast(*struct_array.field(1)); + ASSERT_EQ(numbers_array.value_slice(0)->length(), 2); + ASSERT_EQ(numbers_array.value_slice(1)->length(), 1); + ASSERT_TRUE(numbers_array.IsNull(2)); + + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); +} + +TEST_F(ParquetReaderTest, ReadListAsLargeListWithArrowSchema) { + // Reading a file that carries serialized ARROW:schema metadata must report an output + // schema that describes the arrays the reader actually produces. Arrow decides the list + // type of the arrays, so the output schema follows the schema of the reader instead of + // assuming that the requested large_list type was applied. + CreateListParquetFileWithArrowSchema(); + + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "numbers", + std::make_shared(SchemaField::MakeOptional( + /*field_id=*/101, "element", int32())))}); + + ReaderProperties reader_properties; + reader_properties.Set(ReaderProperties::kArrowUseLargeList, true); + + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kParquet, {.path = temp_parquet_file_, + .io = file_io_, + .projection = schema, + .properties = std::move(reader_properties)}); + ASSERT_THAT(reader_result, IsOk()); + auto reader = std::move(reader_result.value()); + + auto schema_result = reader->Schema(); + ASSERT_THAT(schema_result, IsOk()); + auto arrow_c_schema = std::move(schema_result.value()); + auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); + auto list_type_id = arrow_type->field(1)->type()->id(); + ASSERT_TRUE(list_type_id == ::arrow::Type::LIST || + list_type_id == ::arrow::Type::LARGE_LIST) + << "unexpected list type: " << arrow_type->field(1)->type()->ToString(); + + auto data = reader->Next(); + ASSERT_THAT(data, IsOk()); + ASSERT_TRUE(data.value().has_value()); + auto arrow_c_array = data.value().value(); + + // Importing the array against the reported schema fails if the two disagree, which is + // what this test guards: the output schema must not claim a list type that the reader + // did not produce. + auto import_result = ::arrow::ImportArray(&arrow_c_array, arrow_type); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + auto arrow_array = import_result.ValueOrDie(); + ASSERT_TRUE(arrow_array->ValidateFull().ok()); + + const auto& struct_array = + internal::checked_cast(*arrow_array); + ASSERT_EQ(struct_array.length(), 2); + + const auto& id_array = + internal::checked_cast(*struct_array.field(0)); + ASSERT_EQ(id_array.Value(0), 1); + ASSERT_EQ(id_array.Value(1), 2); + + // The list offsets are 32 or 64 bit wide depending on the type the reader produced. + ASSERT_EQ(struct_array.field(1)->type()->id(), list_type_id); + if (list_type_id == ::arrow::Type::LARGE_LIST) { + const auto& numbers_array = + internal::checked_cast(*struct_array.field(1)); + ASSERT_EQ(numbers_array.value_slice(0)->length(), 2); + ASSERT_EQ(numbers_array.value_slice(1)->length(), 1); + } else { + const auto& numbers_array = + internal::checked_cast(*struct_array.field(1)); + ASSERT_EQ(numbers_array.value_slice(0)->length(), 2); + ASSERT_EQ(numbers_array.value_slice(1)->length(), 1); + } +} + TEST_F(ParquetReaderTest, ReadSplit) { CreateSplitParquetFile();