Skip to content

Commit 5a93503

Browse files
advancedxyadvancedxy
authored
feat(parquet): add row number reader for selected row groups (#1006)
1 parent f8eca4f commit 5a93503

13 files changed

Lines changed: 755 additions & 73 deletions

‎src/iceberg/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,8 @@ if(ICEBERG_BUILD_BUNDLE)
322322
parquet/parquet_reader.cc
323323
parquet/parquet_register.cc
324324
parquet/parquet_schema_util.cc
325-
parquet/parquet_writer.cc)
325+
parquet/parquet_writer.cc
326+
parquet/row_number_reader.cc)
326327

327328
# Libraries to link with exported libiceberg_bundle.{so,a}.
328329
set(ICEBERG_BUNDLE_STATIC_BUILD_INTERFACE_LIBS)

‎src/iceberg/arrow/metadata_column_util.cc‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -69,21 +69,9 @@ Result<std::shared_ptr<::arrow::Array>> MakeFilePathArray(const std::string& fil
6969
return array;
7070
}
7171

72-
Result<std::shared_ptr<::arrow::Array>> MakeRowPositionArray(int64_t start_position,
73-
int64_t num_rows,
74-
::arrow::MemoryPool* pool) {
75-
::arrow::Int64Builder builder(pool);
76-
ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
77-
for (int64_t i = 0; i < num_rows; ++i) {
78-
ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(start_position + i));
79-
}
80-
std::shared_ptr<::arrow::Array> array;
81-
ICEBERG_ARROW_RETURN_NOT_OK(builder.Finish(&array));
82-
return array;
83-
}
84-
8572
Result<std::shared_ptr<::arrow::Array>> MakeRowIdArray(
86-
std::optional<int64_t> first_row_id, int64_t start_position, int64_t num_rows,
73+
std::optional<int64_t> first_row_id,
74+
const std::shared_ptr<::arrow::Int64Array>& row_positions, int64_t num_rows,
8775
::arrow::MemoryPool* pool) {
8876
::arrow::Int64Builder builder(pool);
8977
ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
@@ -92,7 +80,7 @@ Result<std::shared_ptr<::arrow::Array>> MakeRowIdArray(
9280
} else {
9381
for (int64_t row_index = 0; row_index < num_rows; ++row_index) {
9482
ICEBERG_ARROW_RETURN_NOT_OK(
95-
builder.Append(first_row_id.value() + start_position + row_index));
83+
builder.Append(first_row_id.value() + row_positions->Value(row_index)));
9684
}
9785
}
9886

‎src/iceberg/arrow/metadata_column_util_internal.h‎

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,9 @@ namespace iceberg::arrow {
3939

4040
/// \brief Context for populating metadata columns during reading.
4141
struct MetadataColumnContext {
42-
std::string file_path; // The file path to populate _file column.
43-
int64_t next_file_pos = 0; // The next row position for populating _pos column.
42+
std::string file_path; // The file path to populate _file column.
43+
// The file row position used when appending metadata values.
44+
int64_t next_file_pos = 0;
4445
std::optional<int64_t> first_row_id; // First row ID used to populate _row_id.
4546
std::optional<int64_t> data_sequence_number; // Data sequence number for lineage.
4647
};
@@ -66,27 +67,18 @@ Result<std::shared_ptr<::arrow::Array>> MakeFilePathArray(const std::string& fil
6667
int64_t num_rows,
6768
::arrow::MemoryPool* pool);
6869

69-
/// \brief Create a sequential int64 array for _pos column.
70-
///
71-
/// Creates an Arrow Int64Array with sequential values starting from start_position.
72-
///
73-
/// \param start_position Starting row position (inclusive).
74-
/// \param num_rows Number of rows in the batch.
75-
/// \param pool Arrow memory pool.
76-
/// \return Arrow Int64Array with values [start_position, start_position + num_rows).
77-
Result<std::shared_ptr<::arrow::Array>> MakeRowPositionArray(int64_t start_position,
78-
int64_t num_rows,
79-
::arrow::MemoryPool* pool);
80-
8170
/// \brief Create a row ID array for _row_id column.
8271
///
8372
/// \param first_row_id First row ID of the data file.
84-
/// \param start_position Starting row position (inclusive).
73+
/// \param row_positions File row positions. May be null when first_row_id is unknown;
74+
/// otherwise must contain num_rows non-null positions.
8575
/// \param num_rows Number of rows in the batch.
8676
/// \param pool Arrow memory pool.
87-
/// \return Arrow Int64Array for _row_id.
77+
/// \return Arrow Int64Array with `first_row_id + pos`, or all nulls when
78+
/// `first_row_id` is unknown.
8879
Result<std::shared_ptr<::arrow::Array>> MakeRowIdArray(
89-
std::optional<int64_t> first_row_id, int64_t start_position, int64_t num_rows,
80+
std::optional<int64_t> first_row_id,
81+
const std::shared_ptr<::arrow::Int64Array>& row_positions, int64_t num_rows,
9082
::arrow::MemoryPool* pool);
9183

9284
/// \brief Create a sequence number array for _last_updated_sequence_number column.

‎src/iceberg/parquet/parquet_data_util.cc‎

Lines changed: 69 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
* under the License.
1818
*/
1919

20+
#include <algorithm>
21+
2022
#include <arrow/array.h>
2123
#include <arrow/array/builder_primitive.h>
2224
#include <arrow/compute/api.h>
@@ -42,7 +44,8 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
4244
const std::shared_ptr<::arrow::Array>& array,
4345
const std::shared_ptr<::arrow::DataType>& output_arrow_type,
4446
const NestedType& nested_type, std::span<const FieldProjection> projections,
45-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool);
47+
const arrow::MetadataColumnContext& metadata_context,
48+
const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool);
4649

4750
/// \brief Create a null array of the given type and length.
4851
Result<std::shared_ptr<::arrow::Array>> MakeNullArray(
@@ -72,10 +75,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectPrimitiveArray(
7275

7376
Result<std::shared_ptr<::arrow::Array>> MakeInheritedRowLineageArray(
7477
int32_t field_id, int64_t num_rows,
75-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
78+
const arrow::MetadataColumnContext& metadata_context,
79+
const std::shared_ptr<::arrow::Int64Array>& row_positions,
80+
::arrow::MemoryPool* pool) {
7681
if (field_id == MetadataColumns::kRowIdColumnId) {
77-
return arrow::MakeRowIdArray(metadata_context.first_row_id,
78-
metadata_context.next_file_pos, num_rows, pool);
82+
return arrow::MakeRowIdArray(metadata_context.first_row_id, row_positions, num_rows,
83+
pool);
7984
}
8085
if (field_id == MetadataColumns::kLastUpdatedSequenceNumberColumnId) {
8186
return arrow::MakeLastUpdatedSequenceNumberArray(
@@ -84,9 +89,14 @@ Result<std::shared_ptr<::arrow::Array>> MakeInheritedRowLineageArray(
8489
return NotSupported("Unsupported row lineage field id: {}", field_id);
8590
}
8691

92+
// Replace nulls in a physical row lineage column with inherited values:
93+
// `first_row_id + _pos` for _row_id and the data sequence number for
94+
// _last_updated_sequence_number.
8795
Result<std::shared_ptr<::arrow::Array>> ProjectRowLineageArray(
8896
const std::shared_ptr<::arrow::Array>& physical_array, int32_t field_id,
89-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
97+
const arrow::MetadataColumnContext& metadata_context,
98+
const std::shared_ptr<::arrow::Int64Array>& row_positions,
99+
::arrow::MemoryPool* pool) {
90100
if (physical_array->null_count() == 0) {
91101
return physical_array;
92102
}
@@ -99,7 +109,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectRowLineageArray(
99109
if (!values->IsNull(row_index)) {
100110
ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(values->Value(row_index)));
101111
} else {
102-
row_metadata_context.next_file_pos = metadata_context.next_file_pos + row_index;
112+
// Only _row_id with a known first_row_id needs a file position. Use the
113+
// original position because selected row groups may leave gaps.
114+
if (field_id == MetadataColumns::kRowIdColumnId &&
115+
metadata_context.first_row_id.has_value()) {
116+
row_metadata_context.next_file_pos = row_positions->Value(row_index);
117+
}
103118
ICEBERG_RETURN_UNEXPECTED(arrow::AppendInheritedRowLineageValue(
104119
field_id, row_metadata_context, &builder));
105120
}
@@ -114,7 +129,9 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
114129
const std::shared_ptr<::arrow::StructArray>& struct_array,
115130
const std::shared_ptr<::arrow::StructType>& output_struct_type,
116131
const StructType& struct_type, std::span<const FieldProjection> projections,
117-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
132+
const arrow::MetadataColumnContext& metadata_context,
133+
const std::shared_ptr<::arrow::Int64Array>& row_positions,
134+
::arrow::MemoryPool* pool) {
118135
if (struct_type.fields().size() != projections.size()) {
119136
return InvalidSchema(
120137
"Inconsistent number of fields ({}) and number of projections ({})",
@@ -147,15 +164,17 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
147164
if (projected_field.type()->is_nested()) {
148165
const auto& nested_type =
149166
internal::checked_cast<const NestedType&>(*projected_field.type());
167+
// Metadata columns belong to top-level rows; nested fields need no positions.
150168
ICEBERG_ASSIGN_OR_RAISE(
151169
projected_array,
152170
ProjectNestedArray(parquet_array, output_arrow_type, nested_type,
153-
field_projection.children, metadata_context, pool));
171+
field_projection.children, metadata_context,
172+
/*row_positions=*/nullptr, pool));
154173
} else if (MetadataColumns::IsRowLineageColumn(projected_field.field_id())) {
155174
ICEBERG_ASSIGN_OR_RAISE(
156175
projected_array,
157176
ProjectRowLineageArray(parquet_array, projected_field.field_id(),
158-
metadata_context, pool));
177+
metadata_context, row_positions, pool));
159178
} else {
160179
ICEBERG_ASSIGN_OR_RAISE(
161180
projected_array,
@@ -177,13 +196,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
177196
arrow::MakeFilePathArray(metadata_context.file_path,
178197
struct_array->length(), pool));
179198
} else if (field_id == MetadataColumns::kFilePositionColumnId) {
180-
ICEBERG_ASSIGN_OR_RAISE(
181-
projected_array, arrow::MakeRowPositionArray(metadata_context.next_file_pos,
182-
struct_array->length(), pool));
199+
projected_array = row_positions;
183200
} else if (MetadataColumns::IsRowLineageColumn(field_id)) {
184-
ICEBERG_ASSIGN_OR_RAISE(projected_array, MakeInheritedRowLineageArray(
185-
field_id, struct_array->length(),
186-
metadata_context, pool));
201+
ICEBERG_ASSIGN_OR_RAISE(
202+
projected_array,
203+
MakeInheritedRowLineageArray(field_id, struct_array->length(),
204+
metadata_context, row_positions, pool));
187205
} else {
188206
return NotSupported("Unsupported metadata field id: {}", field_id);
189207
}
@@ -237,10 +255,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectListArrayImpl(
237255
} else if (element_field.type()->is_nested()) {
238256
const auto& nested_type =
239257
internal::checked_cast<const NestedType&>(*element_field.type());
258+
// Metadata columns belong to top-level rows; list elements need no positions.
240259
ICEBERG_ASSIGN_OR_RAISE(
241260
projected_values,
242261
ProjectNestedArray(list_array->values(), output_element_type, nested_type,
243-
element_projection.children, metadata_context, pool));
262+
element_projection.children, metadata_context,
263+
/*row_positions=*/nullptr, pool));
244264
} else {
245265
ICEBERG_ASSIGN_OR_RAISE(
246266
projected_values,
@@ -296,10 +316,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
296316
ToString(key_projection.kind));
297317
} else if (key_type->is_nested()) {
298318
const auto& nested_type = internal::checked_cast<const NestedType&>(*key_type);
319+
// Metadata columns belong to top-level rows; map keys need no positions.
299320
ICEBERG_ASSIGN_OR_RAISE(
300321
projected_keys,
301322
ProjectNestedArray(map_array->keys(), output_map_type->key_type(), nested_type,
302-
key_projection.children, metadata_context, pool));
323+
key_projection.children, metadata_context,
324+
/*row_positions=*/nullptr, pool));
303325
} else {
304326
ICEBERG_ASSIGN_OR_RAISE(
305327
projected_keys,
@@ -317,10 +339,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectMapArray(
317339
ToString(value_projection.kind));
318340
} else if (value_type->is_nested()) {
319341
const auto& nested_type = internal::checked_cast<const NestedType&>(*value_type);
342+
// Metadata columns belong to top-level rows; map values need no positions.
320343
ICEBERG_ASSIGN_OR_RAISE(
321344
projected_items,
322345
ProjectNestedArray(map_array->items(), output_map_type->item_type(), nested_type,
323-
value_projection.children, metadata_context, pool));
346+
value_projection.children, metadata_context,
347+
/*row_positions=*/nullptr, pool));
324348
} else {
325349
ICEBERG_ASSIGN_OR_RAISE(
326350
projected_items,
@@ -337,7 +361,9 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
337361
const std::shared_ptr<::arrow::Array>& array,
338362
const std::shared_ptr<::arrow::DataType>& output_arrow_type,
339363
const NestedType& nested_type, std::span<const FieldProjection> projections,
340-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
364+
const arrow::MetadataColumnContext& metadata_context,
365+
const std::shared_ptr<::arrow::Int64Array>& row_positions,
366+
::arrow::MemoryPool* pool) {
341367
switch (nested_type.type_id()) {
342368
case TypeId::kStruct: {
343369
if (output_arrow_type->id() != ::arrow::Type::STRUCT) {
@@ -349,7 +375,7 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
349375
internal::checked_pointer_cast<::arrow::StructType>(output_arrow_type);
350376
const auto& struct_type = internal::checked_cast<const StructType&>(nested_type);
351377
return ProjectStructArray(struct_array, output_struct_type, struct_type,
352-
projections, metadata_context, pool);
378+
projections, metadata_context, row_positions, pool);
353379
}
354380
case TypeId::kList: {
355381
const auto& list_type = internal::checked_cast<const ListType&>(nested_type);
@@ -393,18 +419,39 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
393419

394420
} // namespace
395421

422+
bool NeedsRowPositions(const Schema& read_schema,
423+
const arrow::MetadataColumnContext& metadata_context) {
424+
return std::ranges::any_of(read_schema.fields(), [&](const SchemaField& field) {
425+
return field.field_id() == MetadataColumns::kFilePositionColumnId ||
426+
(field.field_id() == MetadataColumns::kRowIdColumnId &&
427+
metadata_context.first_row_id.has_value());
428+
});
429+
}
430+
396431
Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
397432
std::shared_ptr<::arrow::RecordBatch> record_batch,
398433
const std::shared_ptr<::arrow::Schema>& output_arrow_schema,
399434
const Schema& projected_schema, const SchemaProjection& projection,
400-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
435+
const arrow::MetadataColumnContext& metadata_context,
436+
const std::shared_ptr<::arrow::Int64Array>& row_positions,
437+
::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+
}
401447
auto array = std::make_shared<::arrow::StructArray>(
402448
::arrow::struct_(record_batch->schema()->fields()), record_batch->num_rows(),
403449
record_batch->columns());
404450
ICEBERG_ASSIGN_OR_RAISE(
405451
auto output_array,
406452
ProjectNestedArray(array, ::arrow::struct_(output_arrow_schema->fields()),
407-
projected_schema, projection.fields, metadata_context, pool));
453+
projected_schema, projection.fields, metadata_context,
454+
row_positions, pool));
408455
auto* struct_array = internal::checked_cast<::arrow::StructArray*>(output_array.get());
409456
return ::arrow::RecordBatch::Make(output_arrow_schema, record_batch->num_rows(),
410457
struct_array->fields());

‎src/iceberg/parquet/parquet_data_util_internal.h‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,19 +26,27 @@
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,
31+
const arrow::MetadataColumnContext& metadata_context);
32+
2933
/// \brief Convert record batch read from a Parquet file to projected Iceberg Schema.
3034
///
3135
/// \param record_batch The record batch to convert.
3236
/// \param output_arrow_schema The Arrow schema to convert to.
3337
/// \param projected_schema The projected Iceberg schema.
3438
/// \param projection The projection from projected Iceberg schema to the record batch.
35-
/// \param metadata_context The metadata column context for populating _file and _pos.
39+
/// \param metadata_context The file-level constants for populating metadata columns.
40+
/// \param row_positions Optional file row positions for _pos and inherited _row_id.
41+
/// Required when NeedsRowPositions(projected_schema, metadata_context) is true.
42+
/// When non-null, must contain one non-null value per row in record_batch.
3643
/// \param pool The arrow memory pool.
3744
/// \return The converted record batch.
3845
Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
3946
std::shared_ptr<::arrow::RecordBatch> record_batch,
4047
const std::shared_ptr<::arrow::Schema>& output_arrow_schema,
4148
const Schema& projected_schema, const SchemaProjection& projection,
42-
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool);
49+
const arrow::MetadataColumnContext& metadata_context,
50+
const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool);
4351

4452
} // namespace iceberg::parquet

0 commit comments

Comments
 (0)