Skip to content

refactor(parquet): simplify row position handling in projections #1020

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.

Sign up for GitHub

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

Merged
wgtmac merged 1 commit into from
Oct 16, 2026
Merged

refactor(parquet): simplify row position handling in projections #1020

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 26 additions & 23 deletions src/iceberg/parquet/parquet_data_util.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,16 @@ Result<std::shared_ptr<::arrow::Array>> 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<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
const std::shared_ptr<::arrow::Array>& array,
const std::shared_ptr<::arrow::DataType>& output_arrow_type,
const NestedType& nested_type, std::span<const FieldProjection> 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<std::shared_ptr<::arrow::Array>> MakeNullArray(
const std::shared_ptr<::arrow::DataType>& type, int64_t length,
Expand Down Expand Up @@ -142,6 +152,9 @@ Result<std::shared_ptr<::arrow::Array>> 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<std::shared_ptr<::arrow::Array>> projected_arrays;
projected_arrays.reserve(projections.size());
Expand All @@ -164,12 +177,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
if (projected_field.type()->is_nested()) {
const auto& nested_type =
internal::checked_cast<const NestedType&>(*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,
Expand Down Expand Up @@ -255,12 +266,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectListArrayImpl(
} else if (element_field.type()->is_nested()) {
const auto& nested_type =
internal::checked_cast<const NestedType&>(*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,
Expand Down Expand Up @@ -316,12 +325,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
ToString(key_projection.kind));
} else if (key_type->is_nested()) {
const auto& nested_type = internal::checked_cast<const NestedType&>(*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,
Expand All @@ -339,12 +346,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
ToString(value_projection.kind));
} else if (value_type->is_nested()) {
const auto& nested_type = internal::checked_cast<const NestedType&>(*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,
Expand Down Expand Up @@ -419,9 +424,9 @@ Result<std::shared_ptr<::arrow::Array>> 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());
Expand All @@ -435,15 +440,13 @@ Result<std::shared_ptr<::arrow::RecordBatch>> 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());
Expand Down
5 changes: 3 additions & 2 deletions src/iceberg/parquet/parquet_data_util_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand Down
28 changes: 28 additions & 0 deletions src/iceberg/test/parquet_data_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<StructType>(std::vector<SchemaField>{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<StructType>(std::vector<SchemaField>{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()),
Expand Down
Loading