This is an automated email from the ASF dual-hosted git repository.
wgtmac pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-cpp.git
The following commit(s) were added to refs/heads/main by this push:
new 5a935034 feat(parquet): add row number reader for selected row groups
(#1006)
5a935034 is described below
commit 5a9350342b04a102f69aa8f626ca624072f69098
Author: advancedxy <[email protected]>
AuthorDate: Thu Oct 15 14:18:41 2026 +0800
feat(parquet): add row number reader for selected row groups (#1006)
---
src/iceberg/CMakeLists.txt | 3 +-
src/iceberg/arrow/metadata_column_util.cc | 18 +-
src/iceberg/arrow/metadata_column_util_internal.h | 26 +--
src/iceberg/parquet/parquet_data_util.cc | 91 +++++++---
src/iceberg/parquet/parquet_data_util_internal.h | 12 +-
src/iceberg/parquet/parquet_reader.cc | 35 +++-
src/iceberg/parquet/row_number_reader.cc | 123 +++++++++++++
src/iceberg/parquet/row_number_reader_internal.h | 80 +++++++++
src/iceberg/test/CMakeLists.txt | 3 +-
src/iceberg/test/arrow_test.cc | 11 ++
src/iceberg/test/parquet_data_test.cc | 129 ++++++++++++-
src/iceberg/test/parquet_test.cc | 88 +++++++++
src/iceberg/test/row_number_reader_test.cc | 209 ++++++++++++++++++++++
13 files changed, 755 insertions(+), 73 deletions(-)
diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt
index 6939b2be..717f1da4 100644
--- a/src/iceberg/CMakeLists.txt
+++ b/src/iceberg/CMakeLists.txt
@@ -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)
diff --git a/src/iceberg/arrow/metadata_column_util.cc
b/src/iceberg/arrow/metadata_column_util.cc
index 4b81d7b1..0009bf37 100644
--- a/src/iceberg/arrow/metadata_column_util.cc
+++ b/src/iceberg/arrow/metadata_column_util.cc
@@ -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));
@@ -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)));
}
}
diff --git a/src/iceberg/arrow/metadata_column_util_internal.h
b/src/iceberg/arrow/metadata_column_util_internal.h
index 580244e7..a08ea9bf 100644
--- a/src/iceberg/arrow/metadata_column_util_internal.h
+++ b/src/iceberg/arrow/metadata_column_util_internal.h
@@ -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.
};
@@ -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.
diff --git a/src/iceberg/parquet/parquet_data_util.cc
b/src/iceberg/parquet/parquet_data_util.cc
index 052d03b5..77524aac 100644
--- a/src/iceberg/parquet/parquet_data_util.cc
+++ b/src/iceberg/parquet/parquet_data_util.cc
@@ -17,6 +17,8 @@
* under the License.
*/
+#include <algorithm>
+
#include <arrow/array.h>
#include <arrow/array/builder_primitive.h>
#include <arrow/compute/api.h>
@@ -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(
@@ -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(
@@ -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;
}
@@ -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));
}
@@ -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 ({})",
@@ -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));
} 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,
@@ -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);
}
@@ -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,
@@ -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,
@@ -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,
@@ -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) {
@@ -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);
@@ -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());
diff --git a/src/iceberg/parquet/parquet_data_util_internal.h
b/src/iceberg/parquet/parquet_data_util_internal.h
index 5aad68f8..77bb5645 100644
--- a/src/iceberg/parquet/parquet_data_util_internal.h
+++ b/src/iceberg/parquet/parquet_data_util_internal.h
@@ -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
diff --git a/src/iceberg/parquet/parquet_reader.cc
b/src/iceberg/parquet/parquet_reader.cc
index 1f814610..3f15ab7a 100644
--- a/src/iceberg/parquet/parquet_reader.cc
+++ b/src/iceberg/parquet/parquet_reader.cc
@@ -24,6 +24,7 @@
#include <variant>
#include <vector>
+#include <arrow/array.h>
#include <arrow/c/bridge.h>
#include <arrow/memory_pool.h>
#include <arrow/record_batch.h>
@@ -41,6 +42,7 @@
#include "iceberg/parquet/parquet_data_util_internal.h"
#include "iceberg/parquet/parquet_register.h"
#include "iceberg/parquet/parquet_schema_util_internal.h"
+#include "iceberg/parquet/row_number_reader_internal.h"
#include "iceberg/result.h"
#include "iceberg/schema_internal.h"
#include "iceberg/schema_util.h"
@@ -268,6 +270,10 @@ struct ReadContext {
std::shared_ptr<::arrow::Schema> output_arrow_schema_;
// The reader to read record batches from the Parquet file.
std::unique_ptr<::arrow::RecordBatchReader> record_batch_reader_;
+ // Produces file row positions in lockstep with `record_batch_reader_`. Only
set when
+ // the projection needs row positions (_pos or inherited _row_id).
+ // TODO(advancedxy): use Arrow C++'s row number virtual column once
available.
+ std::unique_ptr<RowNumberReader> row_number_reader_;
};
// TODO(gangwu): list of work items
@@ -308,7 +314,6 @@ class ParquetReader::Impl {
// Project read schema onto the Parquet file schema
ICEBERG_ASSIGN_OR_RAISE(projection_, BuildProjection(reader_.get(),
*read_schema_));
metadata_context_ = {.file_path = options.path,
- .next_file_pos = 0,
.first_row_id = options.first_row_id,
.data_sequence_number = options.data_sequence_number};
@@ -326,11 +331,20 @@ class ParquetReader::Impl {
return std::nullopt;
}
+ std::shared_ptr<::arrow::Int64Array> row_positions;
+ if (context_->row_number_reader_) {
+ ICEBERG_ASSIGN_OR_RAISE(row_positions,
+
context_->row_number_reader_->NextBatch(batch->num_rows()));
+ if (row_positions->length() != batch->num_rows()) {
+ return Invalid("Parquet batch has {} rows, but row number reader
returned {}",
+ batch->num_rows(), row_positions->length());
+ }
+ }
+
ICEBERG_ASSIGN_OR_RAISE(
batch, ProjectRecordBatch(std::move(batch),
context_->output_arrow_schema_,
- *read_schema_, projection_,
metadata_context_, pool_));
-
- metadata_context_.next_file_pos += batch->num_rows();
+ *read_schema_, projection_,
metadata_context_,
+ row_positions, pool_));
ArrowArray arrow_array;
ICEBERG_ARROW_RETURN_NOT_OK(::arrow::ExportRecordBatch(*batch,
&arrow_array));
@@ -392,9 +406,9 @@ class ParquetReader::Impl {
// Row group pruning based on the split
// TODO(gangwu): add row group filtering based on zone map, bloom filter,
etc.
+ auto metadata = reader_->parquet_reader()->metadata();
std::vector<int> row_group_indices;
if (split_.has_value()) {
- auto metadata = reader_->parquet_reader()->metadata();
for (int i = 0; i < metadata->num_row_groups(); ++i) {
auto row_group_offset = metadata->RowGroup(i)->file_offset();
if (row_group_offset >= split_->offset &&
@@ -402,12 +416,10 @@ class ParquetReader::Impl {
row_group_indices.push_back(i);
} else if (row_group_offset >= split_->offset + split_->length) {
break;
- } else {
- metadata_context_.next_file_pos += metadata->RowGroup(i)->num_rows();
}
}
} else {
-
row_group_indices.resize(reader_->parquet_reader()->metadata()->num_row_groups());
+ row_group_indices.resize(metadata->num_row_groups());
std::iota(row_group_indices.begin(), row_group_indices.end(), 0); //
NOLINT
}
@@ -422,6 +434,11 @@ class ParquetReader::Impl {
reader_->GetRecordBatchReader(row_group_indices, column_indices));
}
+ if (NeedsRowPositions(*read_schema_, metadata_context_)) {
+ ICEBERG_ASSIGN_OR_RAISE(context_->row_number_reader_,
+ RowNumberReader::Make(*metadata,
row_group_indices, pool_));
+ }
+
// Build the output Arrow schema from the projected Iceberg schema. This
schema is the
// target of ProjectRecordBatch, so it must describe the projected schema
rather than
// the schema of the file.
@@ -460,7 +477,7 @@ class ParquetReader::Impl {
std::shared_ptr<::arrow::io::RandomAccessFile> input_stream_;
// Parquet file reader to create RecordBatchReader.
std::unique_ptr<::parquet::arrow::FileReader> reader_;
- // Metadata column context for populating _file and _pos columns.
+ // File-level constants for metadata and row lineage columns.
arrow::MetadataColumnContext metadata_context_;
// The context to keep track of the reading progress.
std::unique_ptr<ReadContext> context_;
diff --git a/src/iceberg/parquet/row_number_reader.cc
b/src/iceberg/parquet/row_number_reader.cc
new file mode 100644
index 00000000..3ba5a424
--- /dev/null
+++ b/src/iceberg/parquet/row_number_reader.cc
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <limits>
+
+#include <arrow/array.h>
+#include <arrow/array/builder_primitive.h>
+#include <arrow/memory_pool.h>
+#include <parquet/metadata.h>
+
+#include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/parquet/row_number_reader_internal.h"
+
+namespace iceberg::parquet {
+
+Result<std::unique_ptr<RowNumberReader>> RowNumberReader::Make(
+ const ::parquet::FileMetaData& metadata, const std::vector<int>&
row_groups,
+ ::arrow::MemoryPool* pool) {
+ // Prefix sums in footer order, including the end of the last row group.
+ const int num_row_groups = metadata.num_row_groups();
+ std::vector<int64_t> offsets(static_cast<size_t>(num_row_groups) + 1, 0);
+ for (int i = 0; i < num_row_groups; ++i) {
+ const int64_t num_rows = metadata.RowGroup(i)->num_rows();
+ if (num_rows < 0 || num_rows > std::numeric_limits<int64_t>::max() -
offsets[i]) {
+ return InvalidArgument("Invalid row count for row group {}: {}", i,
num_rows);
+ }
+ offsets[i + 1] = offsets[i] + num_rows;
+ }
+
+ // Keep the requested order, including repeated row groups.
+ std::deque<RowRange> ranges;
+ for (int row_group : row_groups) {
+ if (row_group < 0 || row_group >= num_row_groups) {
+ return InvalidArgument("Row group index {} out of bounds [0, {})",
row_group,
+ num_row_groups);
+ }
+ if (offsets[row_group] != offsets[row_group + 1]) {
+ ranges.push_back({.start = offsets[row_group], .end = offsets[row_group
+ 1]});
+ }
+ }
+
+ return std::unique_ptr<RowNumberReader>(new
RowNumberReader(std::move(ranges), pool));
+}
+
+Result<std::shared_ptr<::arrow::Int64Array>> RowNumberReader::NextBatch(
+ int64_t batch_size) {
+ if (batch_size < 0) {
+ return InvalidArgument("Negative batch size: {}", batch_size);
+ }
+
+ std::vector<RowRange> batch_ranges;
+ int64_t remaining = batch_size;
+ while (remaining > 0) {
+ auto range = TakeRange(remaining);
+ if (!range) {
+ break;
+ }
+ remaining -= range->end - range->start;
+ batch_ranges.push_back(*range);
+ }
+
+ auto result = BuildArray(batch_ranges, batch_size - remaining);
+ if (!result.has_value()) {
+ // Restore the consumed ranges in their original order.
+ ranges_.insert(ranges_.begin(), batch_ranges.begin(), batch_ranges.end());
+ }
+ return result;
+}
+
+std::optional<RowNumberReader::RowRange> RowNumberReader::TakeRange(int64_t
count) {
+ if (ranges_.empty()) {
+ return std::nullopt;
+ }
+
+ auto range = ranges_.front();
+ if (range.end - range.start <= count) {
+ ranges_.pop_front();
+ } else {
+ range.end = range.start + count;
+ ranges_.front().start = range.end;
+ }
+ return range;
+}
+
+Result<std::shared_ptr<::arrow::Int64Array>> RowNumberReader::BuildArray(
+ const std::vector<RowRange>& ranges, int64_t length) {
+ // Arrow's numeric builder converts the element count to bytes internally.
+ // Check the actual batch length, not the requested upper bound.
+ if (length >
+ std::numeric_limits<int64_t>::max() /
static_cast<int64_t>(sizeof(int64_t))) {
+ return InvalidArgument("Row number batch is too large");
+ }
+
+ ::arrow::Int64Builder builder(pool_);
+ ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(length));
+ for (const auto& range : ranges) {
+ for (int64_t row = range.start; row < range.end; ++row) {
+ builder.UnsafeAppend(row);
+ }
+ }
+ std::shared_ptr<::arrow::Int64Array> array;
+ ICEBERG_ARROW_RETURN_NOT_OK(builder.Finish(&array));
+
+ return array;
+}
+
+} // namespace iceberg::parquet
diff --git a/src/iceberg/parquet/row_number_reader_internal.h
b/src/iceberg/parquet/row_number_reader_internal.h
new file mode 100644
index 00000000..49068d57
--- /dev/null
+++ b/src/iceberg/parquet/row_number_reader_internal.h
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#pragma once
+
+/// \file iceberg/parquet/row_number_reader_internal.h
+/// Produces file row positions for the rows emitted by a Parquet record batch
reader.
+
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include <arrow/type_fwd.h>
+
+#include "iceberg/result.h"
+
+namespace parquet {
+class FileMetaData;
+} // namespace parquet
+
+namespace iceberg::parquet {
+
+/// \brief Reader of file row positions ("row numbers") for selected row
groups.
+///
+/// Row numbers are absolute positions in the file: row groups are numbered by
prefix
+/// sums of `num_rows` in footer order, and selected row groups are emitted in
the
+/// requested order, which may skip or reorder row groups.
+///
+/// Like parquet-rs RowNumberReader, this keeps ranges until a batch is
requested.
+/// Row positions are independent of Iceberg metadata column and lineage
semantics.
+class RowNumberReader {
+ public:
+ /// \brief Create a reader for `row_groups`, in the order they will be read.
+ static Result<std::unique_ptr<RowNumberReader>> Make(
+ const ::parquet::FileMetaData& metadata, const std::vector<int>&
row_groups,
+ ::arrow::MemoryPool* pool);
+
+ /// \brief Read up to `batch_size` row numbers, returning an empty array at
EOF.
+ /// If the call fails, no row numbers are consumed and the call can be
retried.
+ Result<std::shared_ptr<::arrow::Int64Array>> NextBatch(int64_t batch_size);
+
+ private:
+ /// \brief Half-open range [start, end) of file row positions.
+ struct RowRange {
+ int64_t start;
+ int64_t end;
+ };
+
+ RowNumberReader(std::deque<RowRange> ranges, ::arrow::MemoryPool* pool)
+ : ranges_(std::move(ranges)), pool_(pool) {}
+
+ // Take up to count rows from the first range, splitting it when necessary.
+ std::optional<RowRange> TakeRange(int64_t count);
+ Result<std::shared_ptr<::arrow::Int64Array>> BuildArray(
+ const std::vector<RowRange>& ranges, int64_t length);
+
+ std::deque<RowRange> ranges_;
+ ::arrow::MemoryPool* pool_;
+};
+
+} // namespace iceberg::parquet
diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt
index ac978c2c..292132bb 100644
--- a/src/iceberg/test/CMakeLists.txt
+++ b/src/iceberg/test/CMakeLists.txt
@@ -220,7 +220,8 @@ if(ICEBERG_BUILD_BUNDLE)
parquet_data_test.cc
parquet_metrics_test.cc
parquet_schema_test.cc
- parquet_test.cc)
+ parquet_test.cc
+ row_number_reader_test.cc)
if(ICEBERG_ENCRYPTION)
target_sources(parquet_test PRIVATE parquet_encryption_test.cc)
endif()
diff --git a/src/iceberg/test/arrow_test.cc b/src/iceberg/test/arrow_test.cc
index feecabf1..c6681480 100644
--- a/src/iceberg/test/arrow_test.cc
+++ b/src/iceberg/test/arrow_test.cc
@@ -20,6 +20,7 @@
#include <functional>
#include <memory>
#include <mutex>
+#include <optional>
#include <string>
#include <string_view>
#include <thread>
@@ -38,6 +39,7 @@
#include <gtest/gtest.h>
#include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow/metadata_column_util_internal.h"
#include "iceberg/constants.h"
#include "iceberg/result.h"
#include "iceberg/schema.h"
@@ -48,6 +50,15 @@
namespace iceberg {
+TEST(MetadataColumnUtilTest, RowIdsWithoutFirstRowId) {
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto row_ids,
+ arrow::MakeRowIdArray(std::nullopt, nullptr, 3,
::arrow::default_memory_pool()));
+ ASSERT_TRUE(row_ids->ValidateFull().ok());
+ EXPECT_EQ(row_ids->length(), 3);
+ EXPECT_EQ(row_ids->null_count(), row_ids->length());
+}
+
struct ToArrowSchemaParam {
std::shared_ptr<Type> iceberg_type;
bool optional = true;
diff --git a/src/iceberg/test/parquet_data_test.cc
b/src/iceberg/test/parquet_data_test.cc
index e300d69b..f7072324 100644
--- a/src/iceberg/test/parquet_data_test.cc
+++ b/src/iceberg/test/parquet_data_test.cc
@@ -29,6 +29,7 @@
#include <gtest/gtest.h>
#include "iceberg/expression/literal.h"
+#include "iceberg/metadata_columns.h"
#include "iceberg/parquet/parquet_data_util_internal.h"
#include "iceberg/schema.h"
#include "iceberg/schema_internal.h"
@@ -55,9 +56,11 @@ std::shared_ptr<::arrow::RecordBatch> RecordBatchFromJson(
struct_array->fields());
}
-void VerifyProjectRecordBatch(const Schema& projected_schema, const Schema&
source_schema,
- const std::string& input_json,
- const std::string& expected_json) {
+void VerifyProjectRecordBatch(
+ const Schema& projected_schema, const Schema& source_schema,
+ const std::string& input_json, const std::string& expected_json,
+ const arrow::MetadataColumnContext& metadata_context = {},
+ const std::shared_ptr<::arrow::Int64Array>& row_positions = nullptr) {
auto schema_projection_result =
Project(projected_schema, source_schema, /*prune_source=*/false);
ASSERT_THAT(schema_projection_result, IsOk());
@@ -75,7 +78,7 @@ void VerifyProjectRecordBatch(const Schema& projected_schema,
const Schema& sour
auto project_result = ProjectRecordBatch(
input_record_batch, projected_arrow_schema, projected_schema,
schema_projection,
- /*metadata_context=*/{}, ::arrow::default_memory_pool());
+ metadata_context, row_positions, ::arrow::default_memory_pool());
ASSERT_THAT(project_result, IsOk());
auto projected_record_batch = std::move(project_result.value());
@@ -88,6 +91,120 @@ void VerifyProjectRecordBatch(const Schema&
projected_schema, const Schema& sour
class ProjectRecordBatchTest : public
::testing::TestWithParam<ProjectRecordBatchParam> {
};
+TEST(ProjectRecordBatchTest, InheritedLineageUsesFilePositions) {
+ Schema source_schema({SchemaField::MakeRequired(1, "id", int32())});
+ Schema projected_schema({SchemaField::MakeRequired(1, "id", int32()),
+ MetadataColumns::kRowPosition,
MetadataColumns::kRowId,
+ MetadataColumns::kLastUpdatedSequenceNumber});
+ auto positions = std::static_pointer_cast<::arrow::Int64Array>(
+ ::arrow::json::ArrayFromJSONString(::arrow::int64(), "[12, 3,
30]").ValueOrDie());
+ ASSERT_NO_FATAL_FAILURE(VerifyProjectRecordBatch(
+ projected_schema, source_schema, "[[10], [20], [30]]",
+ "[[10, 12, 112, 7], [20, 3, 103, 7], [30, 30, 130, 7]]",
+ {.first_row_id = 100, .data_sequence_number = 7}, positions));
+}
+
+TEST(ProjectRecordBatchTest, PhysicalLineageUsesFilePositionsForNulls) {
+ Schema source_schema({SchemaField::MakeRequired(1, "id", int32()),
+ MetadataColumns::kRowId,
+ MetadataColumns::kLastUpdatedSequenceNumber});
+ Schema projected_schema({SchemaField::MakeRequired(1, "id", int32()),
+ MetadataColumns::kRowPosition,
MetadataColumns::kRowId,
+ MetadataColumns::kLastUpdatedSequenceNumber});
+ // A sliced array verifies that lineage uses the array's logical offset.
+ auto positions = std::static_pointer_cast<::arrow::Int64Array>(
+ ::arrow::json::ArrayFromJSONString(::arrow::int64(), "[999, 12, 3, 30]")
+ .ValueOrDie()
+ ->Slice(1));
+ ASSERT_NO_FATAL_FAILURE(VerifyProjectRecordBatch(
+ projected_schema, source_schema, "[[10, null, null], [20, 777, 8], [30,
null, 9]]",
+ "[[10, 12, 112, 7], [20, 3, 777, 8], [30, 30, 130, 9]]",
+ {.first_row_id = 100, .data_sequence_number = 7}, positions));
+}
+
+TEST(ProjectRecordBatchTest, RowPositionsMustMatchBatch) {
+ Schema source_schema(std::vector<SchemaField>{});
+ auto input = ::arrow::RecordBatch::Make(::arrow::schema({}), 3,
::arrow::ArrayVector{});
+ for (const auto& field : {MetadataColumns::kRowPosition,
MetadataColumns::kRowId}) {
+ SCOPED_TRACE(field.name());
+ Schema projected_schema({field});
+ ICEBERG_UNWRAP_OR_FAIL(auto projection, Project(projected_schema,
source_schema,
+ /*prune_source=*/false));
+ auto output_schema = ::arrow::schema(
+ {::arrow::field(std::string(field.name()), ::arrow::int64(),
field.optional())});
+
+ EXPECT_THAT(ProjectRecordBatch(input, output_schema, projected_schema,
projection,
+ {.first_row_id = 100}, nullptr,
+ ::arrow::default_memory_pool()),
+ IsError(ErrorKind::kInvalidArgument));
+ for (const auto& json : {"[12, 3]", "[12, null, 30]"}) {
+ auto positions = std::static_pointer_cast<::arrow::Int64Array>(
+ ::arrow::json::ArrayFromJSONString(::arrow::int64(),
json).ValueOrDie());
+ EXPECT_THAT(ProjectRecordBatch(input, output_schema, projected_schema,
projection,
+ {.first_row_id = 100}, positions,
+ ::arrow::default_memory_pool()),
+ IsError(ErrorKind::kInvalidArgument));
+ }
+ }
+}
+
+TEST(ProjectRecordBatchTest, RowPositionsRequiredForPhysicalRowIdInheritance) {
+ Schema schema({MetadataColumns::kRowId});
+ ICEBERG_UNWRAP_OR_FAIL(auto projection,
+ Project(schema, schema, /*prune_source=*/false));
+ auto arrow_schema = ::arrow::schema({::arrow::field("_row_id",
::arrow::int64())});
+ auto input = RecordBatchFromJson(arrow_schema, "[[null], [777], [null]]");
+ EXPECT_THAT(
+ ProjectRecordBatch(input, arrow_schema, 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()),
+ MetadataColumns::kRowId,
+ MetadataColumns::kLastUpdatedSequenceNumber});
+ ASSERT_NO_FATAL_FAILURE(VerifyProjectRecordBatch(
+ lineage_schema, data_schema, "[[10], [20], [30]]",
+ "[[10, null, 7], [20, null, 7], [30, null, 7]]", {.data_sequence_number
= 7}));
+ const std::string physical_values = "[[10, null, null], [20, 777, 8], [30,
null, 9]]";
+ ASSERT_NO_FATAL_FAILURE(VerifyProjectRecordBatch(
+ lineage_schema, lineage_schema, physical_values,
+ "[[10, null, 7], [20, 777, 8], [30, null, 9]]", {.data_sequence_number =
7}));
+ ASSERT_NO_FATAL_FAILURE(VerifyProjectRecordBatch(lineage_schema,
lineage_schema,
+ physical_values,
physical_values));
+}
+
+TEST(ProjectRecordBatchTest, RowPositionsAreIndependentOfPhysicalColumns) {
+ Schema source_schema({SchemaField::MakeRequired(1, "_pos", int64())});
+ Schema projected_schema({MetadataColumns::kRowPosition,
MetadataColumns::kRowId});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto projection, Project(projected_schema, source_schema,
/*prune_source=*/false));
+ auto positions = std::static_pointer_cast<::arrow::Int64Array>(
+ ::arrow::json::ArrayFromJSONString(::arrow::int64(), "[12, 3,
30]").ValueOrDie());
+ auto physical =
+ ::arrow::json::ArrayFromJSONString(::arrow::int64(), "[10, 20,
30]").ValueOrDie();
+ auto output_schema = ::arrow::schema({::arrow::field("_pos",
::arrow::int64(), false),
+ ::arrow::field("_row_id",
::arrow::int64())});
+ // A physical column named _pos must not be mistaken for the virtual column.
+ auto input_schema = ::arrow::schema({::arrow::field("_pos",
::arrow::int64())});
+ auto input = ::arrow::RecordBatch::Make(input_schema, 3, {physical});
+ EXPECT_THAT(
+ ProjectRecordBatch(input, output_schema, projected_schema, projection,
+ {.first_row_id = 100}, nullptr,
::arrow::default_memory_pool()),
+ IsError(ErrorKind::kInvalidArgument));
+
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto output, ProjectRecordBatch(input, output_schema, projected_schema,
projection,
+ {.first_row_id = 100}, positions,
+ ::arrow::default_memory_pool()));
+ EXPECT_TRUE(output->column(0)->Equals(*positions));
+ EXPECT_EQ(output->column(0)->data()->buffers[1],
positions->data()->buffers[1]);
+ EXPECT_TRUE(output->Equals(
+ *RecordBatchFromJson(output_schema, "[[12, 112], [3, 103], [30,
130]]")));
+}
+
TEST_P(ProjectRecordBatchTest, PrimitiveType) {
const auto& test_case = GetParam();
@@ -659,7 +776,7 @@ TEST(ProjectRecordBatchTest, LargeListOfIntegers) {
// Project the record batch
auto project_result = ProjectRecordBatch(
input_record_batch, output_arrow_schema, iceberg_schema,
schema_projection,
- /*metadata_context=*/{}, ::arrow::default_memory_pool());
+ /*metadata_context=*/{}, /*row_positions=*/nullptr,
::arrow::default_memory_pool());
ASSERT_THAT(project_result, IsOk());
auto projected_record_batch = std::move(project_result.value());
@@ -751,7 +868,7 @@ TEST(ProjectRecordBatchTest, LargeListOfStructs) {
// Project the record batch
auto project_result = ProjectRecordBatch(
input_record_batch, output_arrow_schema, iceberg_schema,
schema_projection,
- /*metadata_context=*/{}, ::arrow::default_memory_pool());
+ /*metadata_context=*/{}, /*row_positions=*/nullptr,
::arrow::default_memory_pool());
ASSERT_THAT(project_result, IsOk());
auto projected_record_batch = std::move(project_result.value());
diff --git a/src/iceberg/test/parquet_test.cc b/src/iceberg/test/parquet_test.cc
index ce43962c..6dedba44 100644
--- a/src/iceberg/test/parquet_test.cc
+++ b/src/iceberg/test/parquet_test.cc
@@ -18,6 +18,8 @@
*/
#include <array>
+#include <format>
+#include <limits>
#include <optional>
#include <string>
#include <utility>
@@ -1029,6 +1031,92 @@ TEST_F(ParquetReaderTest, ReadPhysicalLineage) {
[3, 102, null]])"));
}
+TEST_F(ParquetReaderTest, ReadPositionsAndLineageAcrossRowGroups) {
+ temp_parquet_file_ = "positions_across_row_groups.parquet";
+ constexpr int64_t kRowGroupRows = 100;
+ constexpr int64_t kNumRows = 3 * kRowGroupRows;
+ constexpr int64_t kFirstRowId = 1000;
+
+ // Odd rows keep a physical _row_id; even rows inherit `first_row_id + _pos`.
+ std::string input_json = "[";
+ for (int64_t i = 0; i < kNumRows; ++i) {
+ input_json += std::format("{}[{}, {}, null]", i == 0 ? "" : ",", i,
+ i % 2 == 0 ? "null" : std::to_string(5000 + i));
+ }
+ input_json += "]";
+
+ auto write_schema = RowLineageSchema();
+ ArrowSchema arrow_c_schema;
+ ASSERT_THAT(ToArrowSchema(*write_schema, &arrow_c_schema), IsOk());
+ auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie();
+ auto array = ::arrow::json::ArrayFromJSONString(arrow_type,
input_json).ValueOrDie();
+
+ WriterProperties writer_properties;
+ writer_properties.Set(WriterProperties::kParquetCompression,
+ std::string("uncompressed"));
+ writer_properties.Set(WriterProperties::kParquetMaxRowGroupRows,
kRowGroupRows);
+ ASSERT_THAT(WriteArray(array, {.path = temp_parquet_file_,
+ .schema = write_schema,
+ .io = file_io_,
+ .properties = std::move(writer_properties)}),
+ IsOk());
+
+ auto io = internal::checked_cast<arrow::ArrowFileSystemFileIO&>(*file_io_);
+ auto metadata =
+
::parquet::ReadMetaData(io.fs()->OpenInputFile(temp_parquet_file_).ValueOrDie());
+ ASSERT_EQ(metadata->num_row_groups(), 3);
+
+ auto read_schema = std::make_shared<Schema>(std::vector<SchemaField>{
+ SchemaField::MakeRequired(1, "id", int32()),
+ MetadataColumns::kRowPosition,
+ MetadataColumns::kRowId,
+ });
+
+ auto expected_json = [&](int64_t start, int64_t end) {
+ std::string json = "[";
+ for (int64_t pos = start; pos < end; ++pos) {
+ json += std::format("{}[{}, {}, {}]", pos == start ? "" : ",", pos, pos,
+ pos % 2 == 0 ? kFirstRowId + pos : 5000 + pos);
+ }
+ return json + "]";
+ };
+
+ ReaderProperties reader_properties;
+ reader_properties.Set(ReaderProperties::kBatchSize, int64_t{150});
+
+ // Skip the first row group; the first batch spans the second and third ones.
+ const auto second_offset =
static_cast<size_t>(metadata->RowGroup(1)->file_offset());
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ ReaderFactoryRegistry::Open(
+ FileFormatType::kParquet,
+ {.path = temp_parquet_file_,
+ .split = Split{.offset = second_offset,
+ .length = std::numeric_limits<size_t>::max() -
second_offset},
+ .io = file_io_,
+ .projection = read_schema,
+ .first_row_id = kFirstRowId,
+ .properties = reader_properties}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, expected_json(100, 250)));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, expected_json(250, 300)));
+ ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader));
+
+ // Read only the middle row group.
+ const auto third_offset =
static_cast<size_t>(metadata->RowGroup(2)->file_offset());
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto middle_reader,
+ ReaderFactoryRegistry::Open(FileFormatType::kParquet,
+ {.path = temp_parquet_file_,
+ .split = Split{.offset = second_offset,
+ .length = third_offset -
second_offset},
+ .io = file_io_,
+ .projection = read_schema,
+ .first_row_id = kFirstRowId,
+ .properties = reader_properties}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*middle_reader, expected_json(100,
200)));
+ ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*middle_reader));
+}
+
TEST_F(ParquetReaderTest, ReadNestedUnknownProjection) {
temp_parquet_file_ = "nested_unknown.parquet";
auto write_schema = std::make_shared<Schema>(std::vector<SchemaField>{
diff --git a/src/iceberg/test/row_number_reader_test.cc
b/src/iceberg/test/row_number_reader_test.cc
new file mode 100644
index 00000000..65d521a9
--- /dev/null
+++ b/src/iceberg/test/row_number_reader_test.cc
@@ -0,0 +1,209 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <limits>
+#include <memory>
+#include <numeric>
+#include <vector>
+
+#include <arrow/array.h>
+#include <arrow/buffer.h>
+#include <arrow/memory_pool.h>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <parquet/metadata.h>
+#include <parquet/properties.h>
+#include <parquet/schema.h>
+
+#include "iceberg/parquet/row_number_reader_internal.h"
+#include "iceberg/test/matchers.h"
+
+namespace iceberg::parquet {
+
+namespace {
+
+std::unique_ptr<::parquet::FileMetaData> MakeMetadata(
+ const std::vector<int64_t>& row_group_sizes) {
+ // Only row counts are needed; no column data or file I/O is involved.
+ ::parquet::SchemaDescriptor schema;
+ schema.Init(
+ ::parquet::schema::GroupNode::Make("schema",
::parquet::Repetition::REQUIRED, {}));
+ auto builder = ::parquet::FileMetaDataBuilder::Make(
+ &schema, ::parquet::default_writer_properties());
+ for (int64_t size : row_group_sizes) {
+ auto* row_group = builder->AppendRowGroup();
+ row_group->set_num_rows(size);
+ row_group->Finish(0);
+ }
+ return builder->Finish();
+}
+
+void VerifyNextBatch(RowNumberReader& reader, int64_t batch_size,
+ const std::vector<int64_t>& expected) {
+ ICEBERG_UNWRAP_OR_FAIL(auto array, reader.NextBatch(batch_size));
+ ASSERT_TRUE(array->ValidateFull().ok());
+ ASSERT_EQ(array->null_count(), 0);
+ ASSERT_EQ(array->length(), expected.size());
+ for (size_t i = 0; i < expected.size(); ++i) {
+ EXPECT_EQ(array->Value(i), expected[i]) << "Row " << i;
+ }
+}
+
+} // namespace
+
+TEST(RowNumberReaderTest, ReadsAllRowGroups) {
+ auto metadata = MakeMetadata({2, 3, 4});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ RowNumberReader::Make(*metadata, {0, 1, 2},
::arrow::default_memory_pool()));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 100, {0, 1, 2, 3, 4, 5, 6,
7, 8}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {}));
+}
+
+TEST(RowNumberReaderTest, SkippedRowGroupKeepsFilePositions) {
+ auto metadata = MakeMetadata({3, 3, 3});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ RowNumberReader::Make(*metadata, {0, 2},
::arrow::default_memory_pool()));
+
+ // The first batch spans the gap left by the skipped row group.
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 5, {0, 1, 2, 6, 7}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 5, {8}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {}));
+}
+
+TEST(RowNumberReaderTest, ReorderedAndRepeatedRowGroups) {
+ auto metadata = MakeMetadata({2, 2, 2});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ RowNumberReader::Make(*metadata, {2, 0, 2},
::arrow::default_memory_pool()));
+
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 3, {4, 5, 0}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 3, {1, 4, 5}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {}));
+}
+
+TEST(RowNumberReaderTest, BatchesNotAlignedWithRowGroups) {
+ auto metadata = MakeMetadata({3, 3, 3});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ RowNumberReader::Make(*metadata, {0, 1, 2},
::arrow::default_memory_pool()));
+
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 2, {0, 1}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 4, {2, 3, 4, 5}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 4, {6, 7, 8}));
+}
+
+TEST(RowNumberReaderTest, EmptyRowGroups) {
+ auto metadata = MakeMetadata({0, 2, 0, 1, 0});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ RowNumberReader::Make(*metadata, {0, 1, 2, 3, 4},
::arrow::default_memory_pool()));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 2, {0, 1}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 2, {2}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 2, {}));
+}
+
+TEST(RowNumberReaderTest, NoRowGroups) {
+ for (const auto& sizes : {std::vector<int64_t>{}, std::vector<int64_t>{5}}) {
+ auto metadata = MakeMetadata(sizes);
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader,
+ RowNumberReader::Make(*metadata, {}, ::arrow::default_memory_pool()));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 10, {}));
+ }
+}
+
+TEST(RowNumberReaderTest, ZeroAndNegativeBatchSizesDoNotAdvance) {
+ auto metadata = MakeMetadata({2});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader, RowNumberReader::Make(*metadata, {0},
::arrow::default_memory_pool()));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 0, {}));
+ EXPECT_THAT(reader->NextBatch(-1), IsError(ErrorKind::kInvalidArgument));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {0}));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 0, {}));
+ ASSERT_NO_FATAL_FAILURE(
+ VerifyNextBatch(*reader, std::numeric_limits<int64_t>::max(), {1}));
+}
+
+TEST(RowNumberReaderTest, RowGroupIndexOutOfBounds) {
+ auto metadata = MakeMetadata({5});
+ for (int index : {-1, 1}) {
+ EXPECT_THAT(RowNumberReader::Make(*metadata, {index},
::arrow::default_memory_pool()),
+ IsError(ErrorKind::kInvalidArgument));
+ }
+}
+
+TEST(RowNumberReaderTest, NegativeRowCount) {
+ auto metadata = MakeMetadata({-1});
+ EXPECT_THAT(RowNumberReader::Make(*metadata, {0},
::arrow::default_memory_pool()),
+ IsError(ErrorKind::kInvalidArgument));
+}
+
+TEST(RowNumberReaderTest, OversizedBatchDoesNotAdvance) {
+ auto metadata = MakeMetadata({std::numeric_limits<int64_t>::max()});
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto reader, RowNumberReader::Make(*metadata, {0},
::arrow::default_memory_pool()));
+ EXPECT_THAT(reader->NextBatch(std::numeric_limits<int64_t>::max()),
+ IsError(ErrorKind::kInvalidArgument));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 2, {0, 1}));
+}
+
+TEST(RowNumberReaderTest, AllocationFailureDoesNotAdvance) {
+ auto metadata = MakeMetadata({64, 64, 64, 64});
+ ::arrow::ProxyMemoryPool tracked_pool(::arrow::default_memory_pool());
+ ::arrow::CappedMemoryPool pool(&tracked_pool, 2048);
+ ICEBERG_UNWRAP_OR_FAIL(auto reader, RowNumberReader::Make(*metadata, {0, 2,
3}, &pool));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {0}));
+
+ // Leave room for the values buffer, but not the validity bitmap.
+ auto held_buffer = ::arrow::AllocateBuffer(1024, &pool).ValueOrDie();
+ EXPECT_FALSE(reader->NextBatch(128).has_value());
+ EXPECT_EQ(pool.bytes_allocated(), 1024);
+ EXPECT_FALSE(reader->NextBatch(128).has_value());
+ EXPECT_EQ(pool.bytes_allocated(), 1024);
+
+ // Retry the same request after freeing memory, including the row group gap.
+ held_buffer.reset();
+ EXPECT_EQ(pool.bytes_allocated(), 0);
+ std::vector<int64_t> expected(128);
+ std::iota(expected.begin(), expected.begin() + 63, 1);
+ std::iota(expected.begin() + 63, expected.end(), 128);
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 128, expected));
+ // The last row group was only partially consumed by the successful retry.
+ expected.resize(63);
+ std::ranges::iota(expected, 193);
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 128, expected));
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 1, {}));
+}
+
+TEST(RowNumberReaderTest, BatchesOwnBuffersFromMemoryPool) {
+ auto metadata = MakeMetadata({4});
+ ::arrow::ProxyMemoryPool pool(::arrow::default_memory_pool());
+ ICEBERG_UNWRAP_OR_FAIL(auto reader, RowNumberReader::Make(*metadata, {0},
&pool));
+ ICEBERG_UNWRAP_OR_FAIL(auto first_batch, reader->NextBatch(2));
+ EXPECT_GT(pool.bytes_allocated(), 0);
+ ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, 2, {2, 3}));
+ reader.reset();
+ EXPECT_EQ(first_batch->Value(0), 0);
+ EXPECT_EQ(first_batch->Value(1), 1);
+}
+
+} // namespace iceberg::parquet
