diff --git a/src/iceberg/parquet/parquet_data_util.cc b/src/iceberg/parquet/parquet_data_util.cc index 77524aac6..449657df8 100644 --- a/src/iceberg/parquet/parquet_data_util.cc +++ b/src/iceberg/parquet/parquet_data_util.cc @@ -47,6 +47,16 @@ Result> ProjectNestedArray( const arrow::MetadataColumnContext& metadata_context, const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool); +// The reader currently supplies row positions only for top-level projection. +Result> ProjectNestedArray( + const std::shared_ptr<::arrow::Array>& array, + const std::shared_ptr<::arrow::DataType>& output_arrow_type, + const NestedType& nested_type, std::span projections, + const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) { + return ProjectNestedArray(array, output_arrow_type, nested_type, projections, + metadata_context, /*row_positions=*/nullptr, pool); +} + /// \brief Create a null array of the given type and length. Result> MakeNullArray( const std::shared_ptr<::arrow::DataType>& type, int64_t length, @@ -142,6 +152,9 @@ Result> ProjectStructArray( "Inconsistent number of fields ({}) and number of output fields ({})", struct_type.fields().size(), output_struct_type->num_fields()); } + ICEBERG_PRECHECK( + row_positions != nullptr || !NeedsRowPositions(struct_type, metadata_context), + "Row positions are required to project _pos or inherit _row_id"); std::vector> projected_arrays; projected_arrays.reserve(projections.size()); @@ -164,12 +177,10 @@ Result> ProjectStructArray( if (projected_field.type()->is_nested()) { const auto& nested_type = internal::checked_cast(*projected_field.type()); - // Metadata columns belong to top-level rows; nested fields need no positions. ICEBERG_ASSIGN_OR_RAISE( projected_array, ProjectNestedArray(parquet_array, output_arrow_type, nested_type, - field_projection.children, metadata_context, - /*row_positions=*/nullptr, pool)); + field_projection.children, metadata_context, pool)); } else if (MetadataColumns::IsRowLineageColumn(projected_field.field_id())) { ICEBERG_ASSIGN_OR_RAISE( projected_array, @@ -255,12 +266,10 @@ Result> ProjectListArrayImpl( } else if (element_field.type()->is_nested()) { const auto& nested_type = internal::checked_cast(*element_field.type()); - // Metadata columns belong to top-level rows; list elements need no positions. ICEBERG_ASSIGN_OR_RAISE( projected_values, ProjectNestedArray(list_array->values(), output_element_type, nested_type, - element_projection.children, metadata_context, - /*row_positions=*/nullptr, pool)); + element_projection.children, metadata_context, pool)); } else { ICEBERG_ASSIGN_OR_RAISE( projected_values, @@ -316,12 +325,10 @@ Result> ProjectMapArray( ToString(key_projection.kind)); } else if (key_type->is_nested()) { const auto& nested_type = internal::checked_cast(*key_type); - // Metadata columns belong to top-level rows; map keys need no positions. ICEBERG_ASSIGN_OR_RAISE( projected_keys, ProjectNestedArray(map_array->keys(), output_map_type->key_type(), nested_type, - key_projection.children, metadata_context, - /*row_positions=*/nullptr, pool)); + key_projection.children, metadata_context, pool)); } else { ICEBERG_ASSIGN_OR_RAISE( projected_keys, @@ -339,12 +346,10 @@ Result> ProjectMapArray( ToString(value_projection.kind)); } else if (value_type->is_nested()) { const auto& nested_type = internal::checked_cast(*value_type); - // Metadata columns belong to top-level rows; map values need no positions. ICEBERG_ASSIGN_OR_RAISE( projected_items, ProjectNestedArray(map_array->items(), output_map_type->item_type(), nested_type, - value_projection.children, metadata_context, - /*row_positions=*/nullptr, pool)); + value_projection.children, metadata_context, pool)); } else { ICEBERG_ASSIGN_OR_RAISE( projected_items, @@ -419,9 +424,9 @@ Result> ProjectNestedArray( } // namespace -bool NeedsRowPositions(const Schema& read_schema, +bool NeedsRowPositions(const StructType& struct_type, const arrow::MetadataColumnContext& metadata_context) { - return std::ranges::any_of(read_schema.fields(), [&](const SchemaField& field) { + return std::ranges::any_of(struct_type.fields(), [&](const SchemaField& field) { return field.field_id() == MetadataColumns::kFilePositionColumnId || (field.field_id() == MetadataColumns::kRowIdColumnId && metadata_context.first_row_id.has_value()); @@ -435,15 +440,13 @@ Result> ProjectRecordBatch( const arrow::MetadataColumnContext& metadata_context, const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool) { - if (row_positions == nullptr && NeedsRowPositions(projected_schema, metadata_context)) { - return InvalidArgument( - "Row positions are required to project _pos or inherit _row_id"); - } - if (row_positions != nullptr && (row_positions->length() != record_batch->num_rows() || - row_positions->null_count() != 0)) { - return InvalidArgument("Expected {} non-null row positions", - record_batch->num_rows()); - } + ICEBERG_PRECHECK( + row_positions != nullptr || !NeedsRowPositions(projected_schema, metadata_context), + "Row positions are required to project _pos or inherit _row_id"); + ICEBERG_PRECHECK( + row_positions == nullptr || (row_positions->length() == record_batch->num_rows() && + row_positions->null_count() == 0), + "Expected {} non-null row positions", record_batch->num_rows()); auto array = std::make_shared<::arrow::StructArray>( ::arrow::struct_(record_batch->schema()->fields()), record_batch->num_rows(), record_batch->columns()); diff --git a/src/iceberg/parquet/parquet_data_util_internal.h b/src/iceberg/parquet/parquet_data_util_internal.h index 77bb5645c..ad70c223e 100644 --- a/src/iceberg/parquet/parquet_data_util_internal.h +++ b/src/iceberg/parquet/parquet_data_util_internal.h @@ -26,8 +26,8 @@ namespace iceberg::parquet { -/// \brief Whether the schema projects _pos, or _row_id with a known first_row_id. -bool NeedsRowPositions(const Schema& read_schema, +/// \brief Whether direct fields need positions for _pos or inherited _row_id. +bool NeedsRowPositions(const StructType& struct_type, const arrow::MetadataColumnContext& metadata_context); /// \brief Convert record batch read from a Parquet file to projected Iceberg Schema. @@ -38,6 +38,7 @@ bool NeedsRowPositions(const Schema& read_schema, /// \param projection The projection from projected Iceberg schema to the record batch. /// \param metadata_context The file-level constants for populating metadata columns. /// \param row_positions Optional file row positions for _pos and inherited _row_id. +/// Currently available only to top-level fields. /// Required when NeedsRowPositions(projected_schema, metadata_context) is true. /// When non-null, must contain one non-null value per row in record_batch. /// \param pool The arrow memory pool. diff --git a/src/iceberg/test/parquet_data_test.cc b/src/iceberg/test/parquet_data_test.cc index f70723244..1ec4d8065 100644 --- a/src/iceberg/test/parquet_data_test.cc +++ b/src/iceberg/test/parquet_data_test.cc @@ -160,6 +160,34 @@ TEST(ProjectRecordBatchTest, RowPositionsRequiredForPhysicalRowIdInheritance) { IsError(ErrorKind::kInvalidArgument)); } +TEST(ProjectRecordBatchTest, NestedMetadataRequiresRowPositions) { + auto id_field = SchemaField::MakeRequired(2, "id", int32()); + Schema source_schema({SchemaField::MakeRequired( + 1, "nested", std::make_shared(std::vector{id_field}))}); + auto nested_type = ::arrow::struct_({::arrow::field("id", ::arrow::int32(), false)}); + auto input = RecordBatchFromJson( + ::arrow::schema({::arrow::field("nested", nested_type, false)}), "[[[10]]]"); + + for (const auto& field : {MetadataColumns::kRowPosition, MetadataColumns::kRowId}) { + SCOPED_TRACE(field.name()); + Schema projected_schema({SchemaField::MakeRequired( + 1, "nested", + std::make_shared(std::vector{id_field, field}))}); + ICEBERG_UNWRAP_OR_FAIL(auto projection, Project(projected_schema, source_schema, + /*prune_source=*/false)); + auto output_type = ::arrow::struct_( + {::arrow::field("id", ::arrow::int32(), false), + ::arrow::field(std::string(field.name()), ::arrow::int64(), field.optional())}); + auto output_schema = ::arrow::schema({::arrow::field("nested", output_type, false)}); + // Query engines are expected to project these metadata columns at the top level. + // Construct a nested projection deliberately to exercise the check. + EXPECT_THAT(ProjectRecordBatch(input, output_schema, projected_schema, projection, + {.first_row_id = 100}, nullptr, + ::arrow::default_memory_pool()), + IsError(ErrorKind::kInvalidArgument)); + } +} + TEST(ProjectRecordBatchTest, LineageWithoutRowPositions) { Schema data_schema({SchemaField::MakeRequired(1, "id", int32())}); Schema lineage_schema({SchemaField::MakeRequired(1, "id", int32()),