Skip to content

feat(parquet): add row number reader for selected row groups #1006

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 3 commits into from
Oct 15, 2026
Merged

feat(parquet): add row number reader for selected row groups #1006

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3582ec1
feat(parquet): add row number reader for selected row groups
advancedxy Oct 13, 2026
e3b3302
fix(parquet): resolve row number reader clang-tidy warnings
advancedxy Oct 14, 2026
a2ee01b
chore: merge main into parquet row number reader branch
wgtmac Oct 15, 2026
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
3 changes: 2 additions & 1 deletion src/iceberg/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,8 @@ if(ICEBERG_BUILD_BUNDLE)
parquet/parquet_reader.cc
parquet/parquet_register.cc
parquet/parquet_schema_util.cc
parquet/parquet_writer.cc)
parquet/parquet_writer.cc
parquet/row_number_reader.cc)

# Libraries to link with exported libiceberg_bundle.{so,a}.
set(ICEBERG_BUNDLE_STATIC_BUILD_INTERFACE_LIBS)
Expand Down
18 changes: 3 additions & 15 deletions src/iceberg/arrow/metadata_column_util.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,21 +69,9 @@ Result<std::shared_ptr<::arrow::Array>> MakeFilePathArray(const std::string& fil
return array;
}

Result<std::shared_ptr<::arrow::Array>> MakeRowPositionArray(int64_t start_position,
int64_t num_rows,
::arrow::MemoryPool* pool) {
::arrow::Int64Builder builder(pool);
ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
for (int64_t i = 0; i < num_rows; ++i) {
ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(start_position + i));
}
std::shared_ptr<::arrow::Array> array;
ICEBERG_ARROW_RETURN_NOT_OK(builder.Finish(&array));
return array;
}

Result<std::shared_ptr<::arrow::Array>> MakeRowIdArray(
std::optional<int64_t> first_row_id, int64_t start_position, int64_t num_rows,
std::optional<int64_t> first_row_id,
const std::shared_ptr<::arrow::Int64Array>& row_positions, int64_t num_rows,
::arrow::MemoryPool* pool) {
::arrow::Int64Builder builder(pool);
ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
Expand All @@ -92,7 +80,7 @@ Result<std::shared_ptr<::arrow::Array>> MakeRowIdArray(
} else {
for (int64_t row_index = 0; row_index < num_rows; ++row_index) {
ICEBERG_ARROW_RETURN_NOT_OK(
builder.Append(first_row_id.value() + start_position + row_index));
builder.Append(first_row_id.value() + row_positions->Value(row_index)));
}
}

Expand Down
26 changes: 9 additions & 17 deletions src/iceberg/arrow/metadata_column_util_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,9 @@ namespace iceberg::arrow {

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

/// \brief Create a sequential int64 array for _pos column.
///
/// Creates an Arrow Int64Array with sequential values starting from start_position.
///
/// \param start_position Starting row position (inclusive).
/// \param num_rows Number of rows in the batch.
/// \param pool Arrow memory pool.
/// \return Arrow Int64Array with values [start_position, start_position + num_rows).
Result<std::shared_ptr<::arrow::Array>> MakeRowPositionArray(int64_t start_position,
int64_t num_rows,
::arrow::MemoryPool* pool);

/// \brief Create a row ID array for _row_id column.
///
/// \param first_row_id First row ID of the data file.
/// \param start_position Starting row position (inclusive).
/// \param row_positions File row positions. May be null when first_row_id is unknown;
/// otherwise must contain num_rows non-null positions.
/// \param num_rows Number of rows in the batch.
/// \param pool Arrow memory pool.
/// \return Arrow Int64Array for _row_id.
/// \return Arrow Int64Array with `first_row_id + pos`, or all nulls when
/// `first_row_id` is unknown.
Result<std::shared_ptr<::arrow::Array>> MakeRowIdArray(
std::optional<int64_t> first_row_id, int64_t start_position, int64_t num_rows,
std::optional<int64_t> first_row_id,
const std::shared_ptr<::arrow::Int64Array>& row_positions, int64_t num_rows,
::arrow::MemoryPool* pool);

/// \brief Create a sequence number array for _last_updated_sequence_number column.
Expand Down
91 changes: 69 additions & 22 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 @@ -17,6 +17,8 @@
* under the License.
*/

#include <algorithm>

#include <arrow/array.h>
#include <arrow/array/builder_primitive.h>
#include <arrow/compute/api.h>
Expand All @@ -42,7 +44,8 @@ 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);
const arrow::MetadataColumnContext& metadata_context,
const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool);

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

Result<std::shared_ptr<::arrow::Array>> MakeInheritedRowLineageArray(
int32_t field_id, int64_t num_rows,
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
const arrow::MetadataColumnContext& metadata_context,
const std::shared_ptr<::arrow::Int64Array>& row_positions,
::arrow::MemoryPool* pool) {
if (field_id == MetadataColumns::kRowIdColumnId) {
return arrow::MakeRowIdArray(metadata_context.first_row_id,
metadata_context.next_file_pos, num_rows, pool);
return arrow::MakeRowIdArray(metadata_context.first_row_id, row_positions, num_rows,
pool);
}
if (field_id == MetadataColumns::kLastUpdatedSequenceNumberColumnId) {
return arrow::MakeLastUpdatedSequenceNumberArray(
Expand All @@ -84,9 +89,14 @@ Result<std::shared_ptr<::arrow::Array>> MakeInheritedRowLineageArray(
return NotSupported("Unsupported row lineage field id: {}", field_id);
}

// Replace nulls in a physical row lineage column with inherited values:
// `first_row_id + _pos` for _row_id and the data sequence number for
// _last_updated_sequence_number.
Result<std::shared_ptr<::arrow::Array>> ProjectRowLineageArray(
const std::shared_ptr<::arrow::Array>& physical_array, int32_t field_id,
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
const arrow::MetadataColumnContext& metadata_context,
const std::shared_ptr<::arrow::Int64Array>& row_positions,
::arrow::MemoryPool* pool) {
if (physical_array->null_count() == 0) {
return physical_array;
}
Expand All @@ -99,7 +109,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectRowLineageArray(
if (!values->IsNull(row_index)) {
ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(values->Value(row_index)));
} else {
row_metadata_context.next_file_pos = metadata_context.next_file_pos + row_index;
// Only _row_id with a known first_row_id needs a file position. Use the
// original position because selected row groups may leave gaps.
if (field_id == MetadataColumns::kRowIdColumnId &&
metadata_context.first_row_id.has_value()) {
row_metadata_context.next_file_pos = row_positions->Value(row_index);
}
ICEBERG_RETURN_UNEXPECTED(arrow::AppendInheritedRowLineageValue(
field_id, row_metadata_context, &builder));
}
Expand All @@ -114,7 +129,9 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
const std::shared_ptr<::arrow::StructArray>& struct_array,
const std::shared_ptr<::arrow::StructType>& output_struct_type,
const StructType& struct_type, std::span<const FieldProjection> projections,
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
const arrow::MetadataColumnContext& metadata_context,
const std::shared_ptr<::arrow::Int64Array>& row_positions,
::arrow::MemoryPool* pool) {
if (struct_type.fields().size() != projections.size()) {
return InvalidSchema(
"Inconsistent number of fields ({}) and number of projections ({})",
Expand Down Expand Up @@ -147,15 +164,17 @@ 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, pool));
field_projection.children, metadata_context,
/*row_positions=*/nullptr, pool));
Comment on lines +167 to +172

@advancedxy advancedxy Oct 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not satisfied about these nullptr changes yet, however I don't find a more elegant way to passing row positions as it will touch all the ProjectedXXXMethods.

happy to take suggestions.

@wgtmac wgtmac Oct 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A cleaner option would be to keep row positions only in a top-level struct projection helper, and use a separate ProjectNestedArray overload without the row_positions parameter for nested fields.

@advancedxy advancedxy Oct 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good idea， let me address this

@wgtmac wgtmac Oct 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me just merge this PR. You can create followup PR if you wish :)

@advancedxy advancedxy Oct 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can create followup PR if you wish :)

thanks. Let me create a follow-up pr then.

} else if (MetadataColumns::IsRowLineageColumn(projected_field.field_id())) {
ICEBERG_ASSIGN_OR_RAISE(
projected_array,
ProjectRowLineageArray(parquet_array, projected_field.field_id(),
metadata_context, pool));
metadata_context, row_positions, pool));
} else {
ICEBERG_ASSIGN_OR_RAISE(
projected_array,
Expand All @@ -177,13 +196,12 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
arrow::MakeFilePathArray(metadata_context.file_path,
struct_array->length(), pool));
} else if (field_id == MetadataColumns::kFilePositionColumnId) {
ICEBERG_ASSIGN_OR_RAISE(
projected_array, arrow::MakeRowPositionArray(metadata_context.next_file_pos,
struct_array->length(), pool));
projected_array = row_positions;
} else if (MetadataColumns::IsRowLineageColumn(field_id)) {
ICEBERG_ASSIGN_OR_RAISE(projected_array, MakeInheritedRowLineageArray(
field_id, struct_array->length(),
metadata_context, pool));
ICEBERG_ASSIGN_OR_RAISE(
projected_array,
MakeInheritedRowLineageArray(field_id, struct_array->length(),
metadata_context, row_positions, pool));
} else {
return NotSupported("Unsupported metadata field id: {}", field_id);
}
Expand Down Expand Up @@ -237,10 +255,12 @@ 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, pool));
element_projection.children, metadata_context,
/*row_positions=*/nullptr, pool));
} else {
ICEBERG_ASSIGN_OR_RAISE(
projected_values,
Expand Down Expand Up @@ -296,10 +316,12 @@ 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, pool));
key_projection.children, metadata_context,
/*row_positions=*/nullptr, pool));
} else {
ICEBERG_ASSIGN_OR_RAISE(
projected_keys,
Expand All @@ -317,10 +339,12 @@ 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, pool));
value_projection.children, metadata_context,
/*row_positions=*/nullptr, pool));
} else {
ICEBERG_ASSIGN_OR_RAISE(
projected_items,
Expand All @@ -337,7 +361,9 @@ 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) {
const arrow::MetadataColumnContext& metadata_context,
const std::shared_ptr<::arrow::Int64Array>& row_positions,
::arrow::MemoryPool* pool) {
switch (nested_type.type_id()) {
case TypeId::kStruct: {
if (output_arrow_type->id() != ::arrow::Type::STRUCT) {
Expand All @@ -349,7 +375,7 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(
internal::checked_pointer_cast<::arrow::StructType>(output_arrow_type);
const auto& struct_type = internal::checked_cast<const StructType&>(nested_type);
return ProjectStructArray(struct_array, output_struct_type, struct_type,
projections, metadata_context, pool);
projections, metadata_context, row_positions, pool);
}
case TypeId::kList: {
const auto& list_type = internal::checked_cast<const ListType&>(nested_type);
Expand Down Expand Up @@ -393,18 +419,39 @@ Result<std::shared_ptr<::arrow::Array>> ProjectNestedArray(

} // namespace

bool NeedsRowPositions(const Schema& read_schema,
const arrow::MetadataColumnContext& metadata_context) {
return std::ranges::any_of(read_schema.fields(), [&](const SchemaField& field) {
return field.field_id() == MetadataColumns::kFilePositionColumnId ||
(field.field_id() == MetadataColumns::kRowIdColumnId &&
metadata_context.first_row_id.has_value());
});
}

Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
std::shared_ptr<::arrow::RecordBatch> record_batch,
const std::shared_ptr<::arrow::Schema>& output_arrow_schema,
const Schema& projected_schema, const SchemaProjection& projection,
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) {
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());
}
auto array = std::make_shared<::arrow::StructArray>(
::arrow::struct_(record_batch->schema()->fields()), record_batch->num_rows(),
record_batch->columns());
ICEBERG_ASSIGN_OR_RAISE(
auto output_array,
ProjectNestedArray(array, ::arrow::struct_(output_arrow_schema->fields()),
projected_schema, projection.fields, metadata_context, pool));
projected_schema, projection.fields, metadata_context,
row_positions, pool));
auto* struct_array = internal::checked_cast<::arrow::StructArray*>(output_array.get());
return ::arrow::RecordBatch::Make(output_arrow_schema, record_batch->num_rows(),
struct_array->fields());
Expand Down
12 changes: 10 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,19 +26,27 @@

namespace iceberg::parquet {

/// \brief Whether the schema projects _pos, or _row_id with a known first_row_id.
bool NeedsRowPositions(const Schema& read_schema,
const arrow::MetadataColumnContext& metadata_context);

/// \brief Convert record batch read from a Parquet file to projected Iceberg Schema.
///
/// \param record_batch The record batch to convert.
/// \param output_arrow_schema The Arrow schema to convert to.
/// \param projected_schema The projected Iceberg schema.
/// \param projection The projection from projected Iceberg schema to the record batch.
/// \param metadata_context The metadata column context for populating _file and _pos.
/// \param metadata_context The file-level constants for populating metadata columns.
/// \param row_positions Optional file row positions for _pos and inherited _row_id.
/// 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.
/// \return The converted record batch.
Result<std::shared_ptr<::arrow::RecordBatch>> ProjectRecordBatch(
std::shared_ptr<::arrow::RecordBatch> record_batch,
const std::shared_ptr<::arrow::Schema>& output_arrow_schema,
const Schema& projected_schema, const SchemaProjection& projection,
const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool);
const arrow::MetadataColumnContext& metadata_context,
const std::shared_ptr<::arrow::Int64Array>& row_positions, ::arrow::MemoryPool* pool);

} // namespace iceberg::parquet
Loading
Loading