diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 6939b2bef..717f1da4b 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 4b81d7b1b..0009bf37b 100644 --- a/src/iceberg/arrow/metadata_column_util.cc +++ b/src/iceberg/arrow/metadata_column_util.cc @@ -69,21 +69,9 @@ Result> MakeFilePathArray(const std::string& fil return array; } -Result> 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> MakeRowIdArray( - std::optional first_row_id, int64_t start_position, int64_t num_rows, + std::optional 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> 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 580244e7d..a08ea9bfb 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 first_row_id; // First row ID used to populate _row_id. std::optional data_sequence_number; // Data sequence number for lineage. }; @@ -66,27 +67,18 @@ Result> 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> 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> MakeRowIdArray( - std::optional first_row_id, int64_t start_position, int64_t num_rows, + std::optional 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 052d03b55..77524aac6 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 + #include #include #include @@ -42,7 +44,8 @@ Result> ProjectNestedArray( const std::shared_ptr<::arrow::Array>& array, const std::shared_ptr<::arrow::DataType>& output_arrow_type, const NestedType& nested_type, std::span 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> MakeNullArray( @@ -72,10 +75,12 @@ Result> ProjectPrimitiveArray( Result> 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> 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> 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> 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> ProjectStructArray( const std::shared_ptr<::arrow::StructArray>& struct_array, const std::shared_ptr<::arrow::StructType>& output_struct_type, const StructType& struct_type, std::span 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> ProjectStructArray( if (projected_field.type()->is_nested()) { const auto& nested_type = internal::checked_cast(*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> 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> ProjectListArrayImpl( } else if (element_field.type()->is_nested()) { const auto& nested_type = internal::checked_cast(*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> ProjectMapArray( ToString(key_projection.kind)); } else if (key_type->is_nested()) { const auto& nested_type = internal::checked_cast(*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> ProjectMapArray( ToString(value_projection.kind)); } else if (value_type->is_nested()) { const auto& nested_type = internal::checked_cast(*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> ProjectNestedArray( const std::shared_ptr<::arrow::Array>& array, const std::shared_ptr<::arrow::DataType>& output_arrow_type, const NestedType& nested_type, std::span 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> ProjectNestedArray( internal::checked_pointer_cast<::arrow::StructType>(output_arrow_type); const auto& struct_type = internal::checked_cast(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(nested_type); @@ -393,18 +419,39 @@ Result> 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> 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 5aad68f8f..77bb5645c 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> 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 1f8146107..3f15ab7a5 100644 --- a/src/iceberg/parquet/parquet_reader.cc +++ b/src/iceberg/parquet/parquet_reader.cc @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -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 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 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 context_; diff --git a/src/iceberg/parquet/row_number_reader.cc b/src/iceberg/parquet/row_number_reader.cc new file mode 100644 index 000000000..3ba5a424c --- /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 + +#include +#include +#include +#include + +#include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/parquet/row_number_reader_internal.h" + +namespace iceberg::parquet { + +Result> RowNumberReader::Make( + const ::parquet::FileMetaData& metadata, const std::vector& 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 offsets(static_cast(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::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 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(new RowNumberReader(std::move(ranges), pool)); +} + +Result> RowNumberReader::NextBatch( + int64_t batch_size) { + if (batch_size < 0) { + return InvalidArgument("Negative batch size: {}", batch_size); + } + + std::vector 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::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> RowNumberReader::BuildArray( + const std::vector& 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::max() / static_cast(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 000000000..49068d571 --- /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 +#include +#include +#include +#include +#include + +#include + +#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> Make( + const ::parquet::FileMetaData& metadata, const std::vector& 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> 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 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 TakeRange(int64_t count); + Result> BuildArray( + const std::vector& ranges, int64_t length); + + std::deque ranges_; + ::arrow::MemoryPool* pool_; +}; + +} // namespace iceberg::parquet diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 3a38db35d..fea09c034 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 feecabf13..c66814809 100644 --- a/src/iceberg/test/arrow_test.cc +++ b/src/iceberg/test/arrow_test.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -38,6 +39,7 @@ #include #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 iceberg_type; bool optional = true; diff --git a/src/iceberg/test/parquet_data_test.cc b/src/iceberg/test/parquet_data_test.cc index e300d69b4..f70723244 100644 --- a/src/iceberg/test/parquet_data_test.cc +++ b/src/iceberg/test/parquet_data_test.cc @@ -29,6 +29,7 @@ #include #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 { }; +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{}); + 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 ce43962c6..6dedba443 100644 --- a/src/iceberg/test/parquet_test.cc +++ b/src/iceberg/test/parquet_test.cc @@ -18,6 +18,8 @@ */ #include +#include +#include #include #include #include @@ -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(*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(std::vector{ + 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(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::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(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(std::vector{ 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 000000000..65d521a94 --- /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 +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#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& 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& 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{}, std::vector{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::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::max()}); + ICEBERG_UNWRAP_OR_FAIL( + auto reader, RowNumberReader::Make(*metadata, {0}, ::arrow::default_memory_pool())); + EXPECT_THAT(reader->NextBatch(std::numeric_limits::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 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