Skip to content

Commit ad0f270

Browse files
advancedxyadvancedxy
committed
refactor(parquet): simplify row position handling in projections
1 parent 5a93503 commit ad0f270

3 files changed

Lines changed: 57 additions & 25 deletions

File tree

‎src/iceberg/parquet/parquet_data_util.cc‎

Lines changed: 26 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,16 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
4747
const arrow::MetadataColumnContext& metadata_context,
4848
const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool);
4949

50+
// The reader currently supplies row positions only for top-level projection.
51+
Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
52+
const std::shared_ptr<::arrow::Array>& array,
53+
const std::shared_ptr<::arrow::DataType>& output_arrow_type,
54+
const NestedType& nested_type, std::span<const FieldProjection> projections,
55+
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
56+
return ProjectNestedArray(array, output_arrow_type, nested_type, projections,
57+
metadata_context, /*row_positions=*/nullptr, pool);
58+
}
59+
5060
/// \brief Create a null array of the given type and length.
5161
Result<std::shared_ptr<::arrow::Array>> MakeNullArray(
5262
const std::shared_ptr<::arrow::DataType>& type, int64_t length,
@@ -142,6 +152,9 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
142152
"Inconsistent number of fields ({}) and number of output fields ({})",
143153
struct_type.fields().size(), output_struct_type->num_fields());
144154
}
155+
ICEBERG_PRECHECK(
156+
row_positions != nullptr || !NeedsRowPositions(struct_type, metadata_context),
157+
"Row positions are required to project _pos or inherit _row_id");
145158

146159
std::vector<std::shared_ptr<::arrow::Array>> projected_arrays;
147160
projected_arrays.reserve(projections.size());
@@ -164,12 +177,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
164177
if (projected_field.type()->is_nested()) {
165178
const auto& nested_type =
166179
internal::checked_cast<const NestedType&>(*projected_field.type());
167-
// Metadata columns belong to top-level rows; nested fields need no positions.
168180
ICEBERG_ASSIGN_OR_RAISE(
169181
projected_array,
170182
ProjectNestedArray(parquet_array, output_arrow_type, nested_type,
171-
field_projection.children, metadata_context,
172-
/*row_positions=*/nullptr, pool));
183+
field_projection.children, metadata_context, pool));
173184
} else if (MetadataColumns::IsRowLineageColumn(projected_field.field_id())) {
174185
ICEBERG_ASSIGN_OR_RAISE(
175186
projected_array,
@@ -255,12 +266,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectListArrayImpl(
255266
} else if (element_field.type()->is_nested()) {
256267
const auto& nested_type =
257268
internal::checked_cast<const NestedType&>(*element_field.type());
258-
// Metadata columns belong to top-level rows; list elements need no positions.
259269
ICEBERG_ASSIGN_OR_RAISE(
260270
projected_values,
261271
ProjectNestedArray(list_array->values(), output_element_type, nested_type,
262-
element_projection.children, metadata_context,
263-
/*row_positions=*/nullptr, pool));
272+
element_projection.children, metadata_context, pool));
264273
} else {
265274
ICEBERG_ASSIGN_OR_RAISE(
266275
projected_values,
@@ -316,12 +325,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
316325
ToString(key_projection.kind));
317326
} else if (key_type->is_nested()) {
318327
const auto& nested_type = internal::checked_cast<const NestedType&>(*key_type);
319-
// Metadata columns belong to top-level rows; map keys need no positions.
320328
ICEBERG_ASSIGN_OR_RAISE(
321329
projected_keys,
322330
ProjectNestedArray(map_array->keys(), output_map_type->key_type(), nested_type,
323-
key_projection.children, metadata_context,
324-
/*row_positions=*/nullptr, pool));
331+
key_projection.children, metadata_context, pool));
325332
} else {
326333
ICEBERG_ASSIGN_OR_RAISE(
327334
projected_keys,
@@ -339,12 +346,10 @@ Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
339346
ToString(value_projection.kind));
340347
} else if (value_type->is_nested()) {
341348
const auto& nested_type = internal::checked_cast<const NestedType&>(*value_type);
342-
// Metadata columns belong to top-level rows; map values need no positions.
343349
ICEBERG_ASSIGN_OR_RAISE(
344350
projected_items,
345351
ProjectNestedArray(map_array->items(), output_map_type->item_type(), nested_type,
346-
value_projection.children, metadata_context,
347-
/*row_positions=*/nullptr, pool));
352+
value_projection.children, metadata_context, pool));
348353
} else {
349354
ICEBERG_ASSIGN_OR_RAISE(
350355
projected_items,
@@ -419,9 +424,9 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
419424

420425
} // namespace
421426

422-
bool NeedsRowPositions(const Schema& read_schema,
427+
bool NeedsRowPositions(const StructType& struct_type,
423428
const arrow::MetadataColumnContext& metadata_context) {
424-
return std::ranges::any_of(read_schema.fields(), [&](const SchemaField& field) {
429+
return std::ranges::any_of(struct_type.fields(), [&](const SchemaField& field) {
425430
return field.field_id() == MetadataColumns::kFilePositionColumnId ||
426431
(field.field_id() == MetadataColumns::kRowIdColumnId &&
427432
metadata_context.first_row_id.has_value());
@@ -435,15 +440,13 @@ Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
435440
const arrow::MetadataColumnContext& metadata_context,
436441
const std::shared_ptr<::arrow::Int64Array>& row_positions,
437442
::arrow::MemoryPool* pool) {
438-
if (row_positions == nullptr && NeedsRowPositions(projected_schema, metadata_context)) {
439-
return InvalidArgument(
440-
"Row positions are required to project _pos or inherit _row_id");
441-
}
442-
if (row_positions != nullptr && (row_positions->length() != record_batch->num_rows() ||
443-
row_positions->null_count() != 0)) {
444-
return InvalidArgument("Expected {} non-null row positions",
445-
record_batch->num_rows());
446-
}
443+
ICEBERG_PRECHECK(
444+
row_positions != nullptr || !NeedsRowPositions(projected_schema, metadata_context),
445+
"Row positions are required to project _pos or inherit _row_id");
446+
ICEBERG_PRECHECK(
447+
row_positions == nullptr || (row_positions->length() == record_batch->num_rows() &&
448+
row_positions->null_count() == 0),
449+
"Expected {} non-null row positions", record_batch->num_rows());
447450
auto array = std::make_shared<::arrow::StructArray>(
448451
::arrow::struct_(record_batch->schema()->fields()), record_batch->num_rows(),
449452
record_batch->columns());

‎src/iceberg/parquet/parquet_data_util_internal.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@
2626

2727
namespace iceberg::parquet {
2828

29-
/// \brief Whether the schema projects _pos, or _row_id with a known first_row_id.
30-
bool NeedsRowPositions(const Schema& read_schema,
29+
/// \brief Whether direct fields need positions for _pos or inherited _row_id.
30+
bool NeedsRowPositions(const StructType& struct_type,
3131
const arrow::MetadataColumnContext& metadata_context);
3232

3333
/// \brief Convert record batch read from a Parquet file to projected Iceberg Schema.
@@ -38,6 +38,7 @@ bool NeedsRowPositions(const Schema& read_schema,
3838
/// \param projection The projection from projected Iceberg schema to the record batch.
3939
/// \param metadata_context The file-level constants for populating metadata columns.
4040
/// \param row_positions Optional file row positions for _pos and inherited _row_id.
41+
/// Currently available only to top-level fields.
4142
/// Required when NeedsRowPositions(projected_schema, metadata_context) is true.
4243
/// When non-null, must contain one non-null value per row in record_batch.
4344
/// \param pool The arrow memory pool.

‎src/iceberg/test/parquet_data_test.cc‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,34 @@ TEST(ProjectRecordBatchTest, RowPositionsRequiredForPhysicalRowIdInheritance) {
160160
IsError(ErrorKind::kInvalidArgument));
161161
}
162162

163+
TEST(ProjectRecordBatchTest, NestedMetadataRequiresRowPositions) {
164+
auto id_field = SchemaField::MakeRequired(2, "id", int32());
165+
Schema source_schema({SchemaField::MakeRequired(
166+
1, "nested", std::make_shared<StructType>(std::vector<SchemaField>{id_field}))});
167+
auto nested_type = ::arrow::struct_({::arrow::field("id", ::arrow::int32(), false)});
168+
auto input = RecordBatchFromJson(
169+
::arrow::schema({::arrow::field("nested", nested_type, false)}), "[[[10]]]");
170+
171+
for (const auto& field : {MetadataColumns::kRowPosition, MetadataColumns::kRowId}) {
172+
SCOPED_TRACE(field.name());
173+
Schema projected_schema({SchemaField::MakeRequired(
174+
1, "nested",
175+
std::make_shared<StructType>(std::vector<SchemaField>{id_field, field}))});
176+
ICEBERG_UNWRAP_OR_FAIL(auto projection, Project(projected_schema, source_schema,
177+
/*prune_source=*/false));
178+
auto output_type = ::arrow::struct_(
179+
{::arrow::field("id", ::arrow::int32(), false),
180+
::arrow::field(std::string(field.name()), ::arrow::int64(), field.optional())});
181+
auto output_schema = ::arrow::schema({::arrow::field("nested", output_type, false)});
182+
// Query engines are expected to project these metadata columns at the top level.
183+
// Construct a nested projection deliberately to exercise the check.
184+
EXPECT_THAT(ProjectRecordBatch(input, output_schema, projected_schema, projection,
185+
{.first_row_id = 100}, nullptr,
186+
::arrow::default_memory_pool()),
187+
IsError(ErrorKind::kInvalidArgument));
188+
}
189+
}
190+
163191
TEST(ProjectRecordBatchTest, LineageWithoutRowPositions) {
164192
Schema data_schema({SchemaField::MakeRequired(1, "id", int32())});
165193
Schema lineage_schema({SchemaField::MakeRequired(1, "id", int32()),

0 commit comments

Comments
 (0)