diff --git a/src/iceberg/data/file_scan_task_reader.cc b/src/iceberg/data/file_scan_task_reader.cc index 7d3a8bc7e3..8c55e31b2b 100644 --- a/src/iceberg/data/file_scan_task_reader.cc +++ b/src/iceberg/data/file_scan_task_reader.cc @@ -39,16 +39,21 @@ namespace iceberg { namespace { -ReaderOptions MakeReaderOptions(const DataFile& data_file, std::shared_ptr io, +ReaderOptions MakeReaderOptions(const FileScanTask& task, std::shared_ptr io, std::shared_ptr projection, std::shared_ptr filter, std::shared_ptr name_mapping, ReaderProperties properties, std::optional first_row_id, std::optional data_sequence_number) { + const auto& data_file = *task.data_file(); return ReaderOptions{ .path = data_file.file_path, .length = static_cast(data_file.file_size_in_bytes), + .split = task.is_split() ? std::make_optional( + Split{.offset = static_cast(task.start()), + .length = static_cast(task.length())}) + : std::nullopt, .io = std::move(io), .projection = std::move(projection), .filter = std::move(filter), @@ -163,16 +168,29 @@ class FileScanTaskReader::Impl { ICEBERG_PRECHECK(data_file->file_size_in_bytes >= 0, "Data file size must not be negative: {}", data_file->file_size_in_bytes); + ICEBERG_PRECHECK(std::in_range(data_file->file_size_in_bytes), + "Data file size cannot be represented as size_t: {}", + data_file->file_size_in_bytes); if (task.is_split()) { - return NotSupported( - "Reading a partial FileScanTask is not supported yet: {} " - "(start={}, length={})", - data_file->file_path, task.start(), task.length()); + ICEBERG_PRECHECK(task.start() >= 0 && + task.start() <= data_file->file_size_in_bytes && + task.length() > 0 && + task.length() <= data_file->file_size_in_bytes - task.start(), + "Task range is outside file {} of size {} (start={}, length={})", + data_file->file_path, data_file->file_size_in_bytes, task.start(), + task.length()); + // The validated endpoints are bounded by the size_t-representable file size. + if (data_file->file_format != FileFormatType::kParquet) { + return NotSupported( + "Reading a partial FileScanTask requires Parquet: {} " + "(start={}, length={})", + data_file->file_path, task.start(), task.length()); + } } if (task.delete_files().empty()) { auto options = MakeReaderOptions( - *data_file, io_, projected_schema_, task.residual_filter(), name_mapping_, + task, io_, projected_schema_, task.residual_filter(), name_mapping_, properties_, data_file->first_row_id, data_file->data_sequence_number); ICEBERG_ASSIGN_OR_RAISE( auto reader, ReaderFactoryRegistry::Open(data_file->file_format, options)); @@ -197,9 +215,9 @@ class FileScanTaskReader::Impl { ProjectionContext::Make(*required_schema, *projected_schema_, project_batch_function)); - auto options = MakeReaderOptions( - *data_file, io_, required_schema, task.residual_filter(), name_mapping_, - properties_, data_file->first_row_id, data_file->data_sequence_number); + auto options = MakeReaderOptions(task, io_, required_schema, task.residual_filter(), + name_mapping_, properties_, data_file->first_row_id, + data_file->data_sequence_number); ICEBERG_ASSIGN_OR_RAISE(auto reader, ReaderFactoryRegistry::Open(data_file->file_format, options)); diff --git a/src/iceberg/data/file_scan_task_reader.h b/src/iceberg/data/file_scan_task_reader.h index c34f3e7aa8..33a4b82fff 100644 --- a/src/iceberg/data/file_scan_task_reader.h +++ b/src/iceberg/data/file_scan_task_reader.h @@ -68,8 +68,11 @@ class ICEBERG_DATA_EXPORT FileScanTaskReader { /// \brief Open a task and return an Arrow C stream for its projected live rows. /// - /// Partial tasks are rejected until range-aware reading is implemented. The task and - /// its shared data/delete-file metadata must not be mutated while the stream uses them. + /// Partial tasks are supported for Parquet, with each row group assigned to the + /// half-open range containing its midpoint, as in Java Iceberg. Row positions and + /// inherited row IDs remain relative to the whole file. Other formats require + /// whole-file tasks. The task and its shared data/delete-file metadata must not be + /// mutated while the stream uses them. Result Open(const FileScanTask& task); FileScanTaskReader(const FileScanTaskReader&) = delete; diff --git a/src/iceberg/parquet/parquet_data_util.cc b/src/iceberg/parquet/parquet_data_util.cc index 449657df84..73597cbfb8 100644 --- a/src/iceberg/parquet/parquet_data_util.cc +++ b/src/iceberg/parquet/parquet_data_util.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow/literal_util_internal.h" @@ -37,6 +38,61 @@ namespace iceberg::parquet { +Result RowGroupStart(const ::parquet::RowGroupMetaData& row_group, + int64_t file_size) { + if (row_group.num_columns() == 0) { + return Invalid("Cannot split a Parquet row group without columns"); + } + auto column = row_group.ColumnChunk(0); + int64_t start = column->data_page_offset(); + // A dictionary page comes before the data pages. Like parquet-java, ignore it + // unless its offset is positive and earlier. + if (column->has_dictionary_page()) { + const int64_t dictionary_start = column->dictionary_page_offset(); + if (dictionary_start > 0 && dictionary_start < start) { + start = dictionary_start; + } + } + // Byte 0 is the "PAR1" magic, so no page can start there. + if (start <= 0 || start >= file_size) { + return Invalid( + "Cannot split a Parquet row group with missing or out-of-file " + "starting page {} (file size {})", + start, file_size); + } + return start; +} + +Result RowGroupMidpoint(const ::parquet::RowGroupMetaData& row_group, + int64_t file_size) { + ICEBERG_ASSIGN_OR_RAISE(auto start, RowGroupStart(row_group, file_size)); + // The row group's pages are contiguous from its start, so its size is bounded by the + // rest of the file. This bound also keeps the sum below from overflowing. + const int64_t max_size = file_size - start; + // Arrow reports an omitted row-group size as 0. Like parquet-java, fall back to the + // sum of the column chunk sizes. + int64_t size = row_group.total_compressed_size(); + if (size == 0) { + for (int i = 0; i < row_group.num_columns(); ++i) { + const int64_t column_size = row_group.ColumnChunk(i)->total_compressed_size(); + if (column_size < 0 || column_size > max_size - size) { + return Invalid( + "Cannot split a Parquet row group with invalid column {} compressed size {} " + "(start {}, file size {})", + i, column_size, start, file_size); + } + size += column_size; + } + } + if (size < 0 || size > max_size) { + return Invalid( + "Cannot split a Parquet row group with invalid compressed size {} " + "(start {}, file size {})", + size, start, file_size); + } + return start + size / 2; +} + namespace { // Forward declaration diff --git a/src/iceberg/parquet/parquet_data_util_internal.h b/src/iceberg/parquet/parquet_data_util_internal.h index ad70c223e0..afd0ea397a 100644 --- a/src/iceberg/parquet/parquet_data_util_internal.h +++ b/src/iceberg/parquet/parquet_data_util_internal.h @@ -24,12 +24,27 @@ #include "iceberg/arrow/metadata_column_util_internal.h" #include "iceberg/schema_util.h" +namespace parquet { +class RowGroupMetaData; +} // namespace parquet + namespace iceberg::parquet { /// \brief Whether direct fields need positions for _pos or inherited _row_id. bool NeedsRowPositions(const StructType& struct_type, const arrow::MetadataColumnContext& metadata_context); +/// \brief Return the first column's starting page, matching parquet-java's +/// BlockMetaData.getStartingPos(). Reject missing or out-of-file starts. +Result RowGroupStart(const ::parquet::RowGroupMetaData& row_group, + int64_t file_size); + +/// \brief Return the byte offset that assigns a row group to a split: its start plus +/// half its compressed size, matching parquet-java's range filter used by Java Iceberg. +/// Reject sizes that are negative or extend past the end of the file. +Result RowGroupMidpoint(const ::parquet::RowGroupMetaData& row_group, + int64_t file_size); + /// \brief Convert record batch read from a Parquet file to projected Iceberg Schema. /// /// \param record_batch The record batch to convert. diff --git a/src/iceberg/parquet/parquet_reader.cc b/src/iceberg/parquet/parquet_reader.cc index 90fa064658..12da76a9ee 100644 --- a/src/iceberg/parquet/parquet_reader.cc +++ b/src/iceberg/parquet/parquet_reader.cc @@ -20,7 +20,9 @@ #include "iceberg/parquet/parquet_reader.h" #include +#include #include +#include #include #include @@ -276,7 +278,6 @@ class ParquetReader::Impl { return InvalidArgument("Projected schema is required by Parquet reader"); } - split_ = options.split; read_schema_ = options.projection; // Prepare reader properties @@ -297,6 +298,7 @@ class ParquetReader::Impl { ICEBERG_ARROW_ASSIGN_OR_RETURN( reader_, ::parquet::arrow::FileReader::Make(pool_, std::move(file_reader), arrow_reader_properties)); + ICEBERG_RETURN_UNEXPECTED(PlanRowGroups(options)); // Project read schema onto the Parquet file schema ICEBERG_ASSIGN_OR_RAISE(projection_, BuildProjection(reader_.get(), *read_schema_)); @@ -388,37 +390,62 @@ class ParquetReader::Impl { } private: - Status InitReadContext() { - context_ = std::make_unique(); - - // Row group pruning based on the split + Status PlanRowGroups(const ReaderOptions& options) { // 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()) { + if (options.split.has_value()) { + const auto& split = *options.split; + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto file_size, input_stream_->GetSize()); + int64_t row_count = 0; 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 && - row_group_offset < split_->offset + split_->length) { - row_group_indices.push_back(i); - } else if (row_group_offset >= split_->offset + split_->length) { - break; + auto row_group = metadata->RowGroup(i); + // Empty row groups have no pages to assign to a split. + if (row_group->num_rows() == 0) { + continue; + } + // Like Java Iceberg, a split reads the row groups whose midpoints it contains. + auto midpoint = RowGroupMidpoint(*row_group, file_size); + if (!midpoint.has_value()) { + return Invalid( + "Cannot read Parquet split {} (start={}, length={}), row " + "group {}: {}", + options.path, split.offset, split.length, i, midpoint.error().message); + } + // Subtraction avoids overflow for a caller-supplied unbounded split length. + if (std::cmp_greater_equal(*midpoint, split.offset) && + static_cast(*midpoint) - split.offset < split.length) { + row_group_indices_.push_back(i); + } + if (row_group->num_rows() < 0 || + row_group->num_rows() > std::numeric_limits::max() - row_count) { + return Invalid("Invalid Parquet row count in file {}, row group {}", + options.path, i); } + row_count += row_group->num_rows(); } - } else { - row_group_indices.resize(metadata->num_row_groups()); - std::iota(row_group_indices.begin(), row_group_indices.end(), 0); // NOLINT + } else if (metadata->num_row_groups() > 0) { + row_group_indices_.resize(metadata->num_row_groups()); + std::iota(row_group_indices_.begin(), row_group_indices_.end(), 0); // NOLINT } + return {}; + } + Status InitReadContext() { + // Build locally so initialization can be retried without consuming any rows. + auto context = std::make_unique(); // Create the record batch reader, including for empty splits to preserve the schema. - auto column_indices = SelectedColumnIndices(projection_); ICEBERG_ARROW_ASSIGN_OR_RETURN( - context_->record_batch_reader_, - reader_->GetRecordBatchReader(row_group_indices, column_indices)); + context->record_batch_reader_, + reader_->GetRecordBatchReader(row_group_indices_, + SelectedColumnIndices(projection_))); if (NeedsRowPositions(*read_schema_, metadata_context_)) { - ICEBERG_ASSIGN_OR_RAISE(context_->row_number_reader_, - RowNumberReader::Make(*metadata, row_group_indices, pool_)); + auto metadata = reader_->parquet_reader()->metadata(); + // Share the batch reader's selection so positions remain file-relative even + // when a batch spans nonconsecutive row groups. + 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 @@ -426,7 +453,7 @@ class ParquetReader::Impl { // the schema of the file. ArrowSchema arrow_schema; ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*read_schema_, &arrow_schema)); - ICEBERG_ARROW_ASSIGN_OR_RETURN(context_->output_arrow_schema_, + ICEBERG_ARROW_ASSIGN_OR_RETURN(context->output_arrow_schema_, ::arrow::ImportSchema(&arrow_schema)); // Align the output schema with the arrays the reader actually produces. The reader's @@ -437,18 +464,18 @@ class ParquetReader::Impl { // 3. Mixed list and large_list types in files with stored schemas // For each projected field, we use the reader's actual type. For missing fields // (columns not in the file), we apply the configured use_large_list preference. - context_->output_arrow_schema_ = AlignOutputSchemaToReaderSchema( - context_->output_arrow_schema_, context_->record_batch_reader_->schema(), + context->output_arrow_schema_ = AlignOutputSchemaToReaderSchema( + context->output_arrow_schema_, context->record_batch_reader_->schema(), projection_, use_large_list_); - + context_ = std::move(context); return {}; } private: // TODO(gangwu): make memory pool configurable ::arrow::MemoryPool* pool_ = ::arrow::default_memory_pool(); - // The split to read from the Parquet file. - std::optional split_; + // Row groups in footer order, shared by the batch and row-number readers. + std::vector row_group_indices_; // Whether to read list columns as large_list (64-bit offsets). bool use_large_list_ = false; // Schema to read from the Parquet file. diff --git a/src/iceberg/table_scan.cc b/src/iceberg/table_scan.cc index 25548bc769..d04e3385b6 100644 --- a/src/iceberg/table_scan.cc +++ b/src/iceberg/table_scan.cc @@ -224,7 +224,7 @@ Result> FileScanTask::MakeSplit( return std::shared_ptr( new FileScanTask(std::move(data_file), std::move(delete_files), std::move(filter), - Range{.start = start, .length = length, .file_size = file_size})); + Range{.start = start, .length = length})); } int64_t FileScanTask::length() const { @@ -241,16 +241,20 @@ int64_t FileScanTask::estimated_row_count() const { } const int64_t record_count = data_file_->record_count; - if (record_count <= 0) { + const int64_t file_size = data_file_->file_size_in_bytes; + if (record_count <= 0 || file_size <= 0 || range_->start > file_size || + range_->length > file_size - range_->start) { + // Keep arithmetic safe even for tasks invalidated by metadata mutation. return 0; } - // floor(end * rows / size) - floor(start * rows / size), so splits that cover - // the file sum exactly to record_count. int128 avoids overflow. + // floor(end * rows / size) - floor(start * rows / size). For non-overlapping + // splits covering a file with unchanged metadata, estimates sum to record_count. + // int128 avoids overflow. const auto end = static_cast(range_->start) + range_->length; - const auto count_at_end = end * record_count / range_->file_size; + const auto count_at_end = end * record_count / file_size; const auto count_at_start = - static_cast(range_->start) * record_count / range_->file_size; + static_cast(range_->start) * record_count / file_size; return static_cast(count_at_end - count_at_start); } diff --git a/src/iceberg/table_scan.h b/src/iceberg/table_scan.h index 7afaa79d3b..1307089e39 100644 --- a/src/iceberg/table_scan.h +++ b/src/iceberg/table_scan.h @@ -72,8 +72,8 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask { /// \param data_file The data file to read. /// \param delete_files Delete files that apply to this data file. /// \param filter Optional residual filter to apply after reading. - /// \note The data file, delete files, and residual filter must not be mutated while - /// a planning or reading stream uses this task. + /// \note The shared data file, delete files, and residual filter must remain + /// unchanged for this task's lifetime. Metadata changes require rebuilding tasks. explicit FileScanTask(std::shared_ptr data_file, std::vector> delete_files = {}, std::shared_ptr filter = nullptr); @@ -82,8 +82,8 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask { /// /// A range covering the entire file is represented as a whole-file task. An empty /// range is valid only for an empty whole file. The data file, delete files, and - /// residual filter are shared with the resulting task and must not be mutated - /// while a planning or reading stream uses it. + /// residual filter are shared with the resulting task and must remain unchanged + /// for its lifetime. Metadata changes require rebuilding tasks. static Result> MakeSplit( std::shared_ptr data_file, int64_t start, int64_t length, std::vector> delete_files = {}, @@ -92,7 +92,7 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask { /// \brief The inclusive byte offset where this task starts, or zero for a whole file. int64_t start() const { return range_ ? range_->start : 0; } - /// \brief The byte length of this task, or the current size of a whole file. + /// \brief The byte length of this task, or the size of a whole file. int64_t length() const; /// \brief Whether this task reads a proper subset of its data file. @@ -112,6 +112,10 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask { Kind kind() const override { return Kind::kFileScanTask; } int64_t size_bytes() const override; int32_t files_count() const override; + + /// \brief The estimated number of rows in this task, prorated by byte range. + /// For valid, non-overlapping splits covering a file with unchanged metadata, + /// the estimates sum to the file's record_count. int64_t estimated_row_count() const override; private: @@ -119,8 +123,6 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask { struct Range { int64_t start; int64_t length; - // Size when the range was validated. Always positive, so safe to divide by. - int64_t file_size; }; FileScanTask(std::shared_ptr data_file, diff --git a/src/iceberg/test/file_scan_task_reader_test.cc b/src/iceberg/test/file_scan_task_reader_test.cc index 11366f72d9..05416bd7d7 100644 --- a/src/iceberg/test/file_scan_task_reader_test.cc +++ b/src/iceberg/test/file_scan_task_reader_test.cc @@ -19,9 +19,12 @@ #include "iceberg/data/file_scan_task_reader.h" +#include +#include #include #include #include +#include #include #include @@ -31,11 +34,13 @@ #include #include #include +#include #include #include #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_register.h" +#include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow_c_data.h" #include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" @@ -56,6 +61,7 @@ #include "iceberg/schema_internal.h" #include "iceberg/table_scan.h" #include "iceberg/test/matchers.h" +#include "iceberg/test/std_io.h" #include "iceberg/test/temp_file_test_base.h" #include "iceberg/type.h" #include "iceberg/util/macros.h" @@ -121,6 +127,31 @@ class FileScanTaskReaderTest : public TempFileTestBase { table_schema_->schema_id()); } + using RowWithPositionAndLineage = std::array; + + Result> ReadRows(FileScanTaskReader& reader, + const FileScanTask& task) { + ICEBERG_ASSIGN_OR_RAISE(auto stream, reader.Open(task)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batches, + ::arrow::ImportRecordBatchReader(&stream)); + std::vector rows; + while (true) { + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, batches->Next()); + if (batch == nullptr) { + return rows; + } + ICEBERG_ARROW_RETURN_NOT_OK(batch->ValidateFull()); + const auto& ids = static_cast(*batch->column(0)); + const auto& positions = static_cast(*batch->column(1)); + const auto& row_ids = static_cast(*batch->column(2)); + const auto& sequences = static_cast(*batch->column(3)); + for (int64_t i = 0; i < batch->num_rows(); ++i) { + rows.push_back( + {ids.Value(i), positions.Value(i), row_ids.Value(i), sequences.Value(i)}); + } + } + } + Result MakeBatch(const Schema& schema, const std::string& json_data) const { ICEBERG_ASSIGN_OR_RAISE(auto arrow_schema, MakeArrowSchema(schema)); @@ -409,7 +440,7 @@ TEST_F(FileScanTaskReaderTest, OpenWithoutDeletesReadsProjectedSchema) { VerifyStream(&stream, R"([[1, "Foo"], [2, "Bar"], [3, "Baz"]])")); } -TEST_F(FileScanTaskReaderTest, RejectsPartialTaskUntilRangeReadsAreSupported) { +TEST_F(FileScanTaskReaderTest, RejectsPartialTasksForUnsupportedFormats) { ICEBERG_UNWRAP_OR_FAIL( auto data_file, MakeDataFile(table_schema_, @@ -423,11 +454,147 @@ TEST_F(FileScanTaskReaderTest, RejectsPartialTaskUntilRangeReadsAreSupported) { .projected_schema = projected_schema_, }; ICEBERG_UNWRAP_OR_FAIL(auto reader, FileScanTaskReader::Make(std::move(options))); - auto result = reader->Open(*task); - EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); - EXPECT_THAT(result, HasErrorMessage("Reading a partial FileScanTask")); + for (auto format : {FileFormatType::kAvro, FileFormatType::kOrc}) { + data_file->file_format = format; + auto result = reader->Open(*task); + EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); + EXPECT_THAT(result, + HasErrorMessage("Reading a partial FileScanTask requires Parquet")); + } +} + +TEST_F(FileScanTaskReaderTest, RechecksSplitBoundsAfterFileSizeChanges) { + ICEBERG_UNWRAP_OR_FAIL(auto data_file, + MakeDataFile(table_schema_, R"([[1, "a", "b"]])", 1)); + const int64_t file_size = data_file->file_size_in_bytes; + ICEBERG_UNWRAP_OR_FAIL(auto task, FileScanTask::MakeSplit(data_file, 4, file_size - 4)); + ICEBERG_UNWRAP_OR_FAIL(auto reader, FileScanTaskReader::Make({ + .io = file_io_, + .table_schema = table_schema_, + .projected_schema = projected_schema_, + })); + for (int64_t size : {file_size - 1, int64_t{3}, int64_t{0}}) { + data_file->file_size_in_bytes = size; + auto result = reader->Open(*task); + EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(result, HasErrorMessage("Task range is outside file")); + } + data_file->file_size_in_bytes = -1; + EXPECT_THAT(reader->Open(*task), + HasErrorMessage("Data file size must not be negative")); +} + +TEST_F(FileScanTaskReaderTest, EmptyParquetFileAndSplitReturnNoRows) { + ICEBERG_UNWRAP_OR_FAIL(auto data_file, MakeDataFile(table_schema_, "[]", 0)); + ICEBERG_UNWRAP_OR_FAIL(auto reader, FileScanTaskReader::Make({ + .io = file_io_, + .table_schema = table_schema_, + .projected_schema = projected_schema_, + })); + ICEBERG_UNWRAP_OR_FAIL(auto split, FileScanTask::MakeSplit(data_file, 0, 1)); + for (const auto& task : {std::make_shared(data_file), split}) { + ICEBERG_UNWRAP_OR_FAIL(auto stream, reader->Open(*task)); + auto batches = ::arrow::ImportRecordBatchReader(&stream).ValueOrDie(); + auto batch = batches->Next(); + ASSERT_TRUE(batch.ok()) << batch.status(); + ASSERT_EQ(batch.ValueOrDie(), nullptr); + } +} + +class FileScanTaskReaderSplitTest + : public FileScanTaskReaderTest, + public ::testing::WithParamInterface> {}; + +TEST_P(FileScanTaskReaderSplitTest, SplitsMatchWholeFileWithAbsolutePositionsAndLineage) { + const auto& [delete_kind, batch_size] = GetParam(); + ICEBERG_UNWRAP_OR_FAIL( + auto data_file, + MakeDataFile(table_schema_, + std::vector{R"([[1, "a", "keep"], [2, "b", "delete"]])", + R"([[3, "c", "delete"], [4, "d", "delete"]])", + R"([[5, "e", "keep"], [6, "f", "delete"]])"}, + /*record_count=*/6, /*max_row_group_length=*/2)); + data_file->first_row_id = 100; + data_file->data_sequence_number = 7; + std::vector> deletes; + if (delete_kind == "position" || delete_kind == "mixed") { + ICEBERG_UNWRAP_OR_FAIL( + auto deletion, MakePositionDeleteFile(CreateNewTempFilePathWithSuffix(".parquet"), + {1, 2, 3, 5}, data_file->file_path)); + deletes.push_back(std::move(deletion)); + } + if (delete_kind == "equality" || delete_kind == "mixed") { + // The equality field is absent from the output projection. + ICEBERG_UNWRAP_OR_FAIL( + auto deletion, + MakeEqualityDeleteFile(CreateNewTempFilePathWithSuffix(".parquet"), table_schema_, + R"([[0, "", "delete"]])", {3})); + deletes.push_back(std::move(deletion)); + } + if (delete_kind == "dv") { + ICEBERG_UNWRAP_OR_FAIL( + auto deletion, MakeDeletionVectorFile(CreateNewTempFilePathWithSuffix(".puffin"), + {1, 2, 3, 5}, data_file->file_path)); + deletes.push_back(std::move(deletion)); + } + auto projection = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), MetadataColumns::kRowPosition, + MetadataColumns::kRowId, MetadataColumns::kLastUpdatedSequenceNumber}); + // This FileIO honors the supplied length hint, so a split length passed as the + // file length would make the footer unreadable. + ICEBERG_UNWRAP_OR_FAIL( + auto reader, FileScanTaskReader::Make({ + .io = std::make_shared(), + .table_schema = table_schema_, + .projected_schema = projection, + .properties = {{"read.batch-size", std::to_string(batch_size)}}, + })); + + std::vector expected; + for (int64_t id = 1; id <= 6; ++id) { + if (delete_kind == "none" || id == 1 || id == 5) { + expected.push_back({id, id - 1, 100 + id - 1, 7}); + } + } + ICEBERG_UNWRAP_OR_FAIL(auto whole_rows, + ReadRows(*reader, FileScanTask(data_file, deletes))); + ASSERT_EQ(whole_rows, expected); + + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, data_file->file_path)); + auto metadata = ::parquet::ReadMetaData(input); + ASSERT_EQ(metadata->num_row_groups(), 3); + std::vector boundaries{0, 1, data_file->file_size_in_bytes}; + for (int i = 0; i < metadata->num_row_groups(); ++i) { + auto column = metadata->RowGroup(i)->ColumnChunk(0); + ASSERT_TRUE(column->has_dictionary_page()); + ASSERT_LT(column->dictionary_page_offset(), column->data_page_offset()); + // Include exact starts, ranges containing no starts, and boundaries inside groups. + boundaries.push_back(column->dictionary_page_offset()); + boundaries.push_back(column->dictionary_page_offset() + 1); + boundaries.push_back(column->data_page_offset()); + } + std::ranges::sort(boundaries); + auto duplicates = std::ranges::unique(boundaries); + boundaries.erase(duplicates.begin(), duplicates.end()); + std::vector split_rows; + for (size_t i = 1; i < boundaries.size(); ++i) { + SCOPED_TRACE(i); + ICEBERG_UNWRAP_OR_FAIL( + auto split, FileScanTask::MakeSplit(data_file, boundaries[i - 1], + boundaries[i] - boundaries[i - 1], deletes)); + ICEBERG_UNWRAP_OR_FAIL(auto rows, ReadRows(*reader, *split)); + split_rows.insert(split_rows.end(), rows.begin(), rows.end()); + } + EXPECT_EQ(split_rows, whole_rows); + EXPECT_EQ(split_rows, expected); } +INSTANTIATE_TEST_SUITE_P(DeleteKindsAndBatchSizes, FileScanTaskReaderSplitTest, + ::testing::Combine(::testing::Values("none", "position", + "equality", "dv", "mixed"), + ::testing::Values(int64_t{1}, int64_t{100}))); + TEST_F(FileScanTaskReaderTest, ReadLastUpdatedFromDataSeq) { ICEBERG_UNWRAP_OR_FAIL( auto data_file, diff --git a/src/iceberg/test/file_scan_task_test.cc b/src/iceberg/test/file_scan_task_test.cc index 6394c96fed..41861caec9 100644 --- a/src/iceberg/test/file_scan_task_test.cc +++ b/src/iceberg/test/file_scan_task_test.cc @@ -64,9 +64,10 @@ TEST(FileScanTaskRangeTest, WholeFileAndEmptyFile) { EXPECT_EQ(explicit_whole->length(), 100); EXPECT_EQ(explicit_whole->estimated_row_count(), 7); - data_file->file_size_in_bytes = 0; - data_file->record_count = 0; - ICEBERG_UNWRAP_OR_FAIL(auto empty, FileScanTask::MakeSplit(data_file, 0, 0)); + auto empty_file = std::make_shared(); + empty_file->file_size_in_bytes = 0; + empty_file->record_count = 0; + ICEBERG_UNWRAP_OR_FAIL(auto empty, FileScanTask::MakeSplit(empty_file, 0, 0)); EXPECT_FALSE(empty->is_split()); EXPECT_EQ(empty->start(), 0); EXPECT_EQ(empty->length(), 0); @@ -155,6 +156,34 @@ TEST(FileScanTaskRangeTest, RowEstimateDoesNotOverflow) { EXPECT_EQ(last->estimated_row_count(), 1); } +TEST(FileScanTaskRangeTest, RowEstimateIsBoundedForInvalidatedRanges) { + auto data_file = std::make_shared(); + data_file->file_path = "test.parquet"; + data_file->file_size_in_bytes = 100; + data_file->record_count = std::numeric_limits::max(); + ICEBERG_UNWRAP_OR_FAIL(auto split, FileScanTask::MakeSplit(data_file, 40, 60)); + + // Mutation violates the task's lifetime contract, but the estimate should still + // avoid overflow and division by zero when the range is invalidated. + for (int64_t file_size : {99, 40, 39, 1, 0, -1}) { + SCOPED_TRACE(file_size); + data_file->file_size_in_bytes = file_size; + EXPECT_EQ(split->estimated_row_count(), 0); + } +} + +TEST(FileScanTaskRangeTest, RowEstimatesHandleNonPositiveRecordCounts) { + auto data_file = std::make_shared(); + data_file->file_size_in_bytes = 100; + + for (int64_t record_count : {0, -1}) { + SCOPED_TRACE(record_count); + data_file->record_count = record_count; + ICEBERG_UNWRAP_OR_FAIL(auto split, FileScanTask::MakeSplit(data_file, 0, 40)); + EXPECT_EQ(split->estimated_row_count(), 0); + } +} + class FileScanTaskTest : public TempFileTestBase { protected: static void SetUpTestSuite() { parquet::RegisterAll(); } diff --git a/src/iceberg/test/parquet_test.cc b/src/iceberg/test/parquet_test.cc index 9fbeae8789..1327f03f5f 100644 --- a/src/iceberg/test/parquet_test.cc +++ b/src/iceberg/test/parquet_test.cc @@ -39,7 +39,9 @@ #include #include #include +#include #include +#include #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_status_internal.h" @@ -47,6 +49,7 @@ #include "iceberg/file_reader.h" #include "iceberg/file_writer.h" #include "iceberg/metadata_columns.h" +#include "iceberg/parquet/parquet_data_util_internal.h" #include "iceberg/parquet/parquet_register.h" #include "iceberg/result.h" #include "iceberg/schema.h" @@ -55,7 +58,6 @@ #include "iceberg/test/matchers.h" #include "iceberg/test/std_io.h" #include "iceberg/test/temp_file_test_base.h" -#include "iceberg/test/test_resource.h" #include "iceberg/type.h" #include "iceberg/util/checked_cast.h" #include "iceberg/util/macros.h" @@ -158,6 +160,31 @@ std::optional FirstUnavailableParquetCodec() { return std::nullopt; } +std::unique_ptr<::parquet::FileMetaData> MakeMetadataWithPageOffsets( + const std::vector>>& page_offsets) { + ::parquet::schema::NodeVector fields; + for (size_t i = 0; i < page_offsets.size(); ++i) { + fields.push_back(::parquet::schema::PrimitiveNode::Make( + "col" + std::to_string(i), ::parquet::Repetition::REQUIRED, nullptr, + ::parquet::Type::INT32, /*primitive_length=*/-1, /*field_id=*/i + 1)); + } + ::parquet::SchemaDescriptor schema; + schema.Init(::parquet::schema::GroupNode::Make( + "schema", ::parquet::Repetition::REQUIRED, fields)); + auto builder = ::parquet::FileMetaDataBuilder::Make( + &schema, ::parquet::default_writer_properties()); + auto row_group = builder->AppendRowGroup(); + row_group->set_num_rows(1); + for (const auto& [data_offset, dictionary_offset] : page_offsets) { + row_group->NextColumnChunk()->Finish( + /*num_values=*/1, dictionary_offset.value_or(0), /*index_page_offset=*/0, + data_offset, /*compressed_size=*/10, /*uncompressed_size=*/10, + dictionary_offset.has_value(), /*dictionary_fallback=*/false, {}, {}); + } + row_group->Finish(10 * page_offsets.size()); + return builder->Finish(); +} + constexpr std::array kUuidBytes1 = { 0x12, 0x3e, 0x45, 0x67, 0xe8, 0x9b, 0x12, 0xd3, 0xa4, 0x56, 0x42, 0x66, 0x14, 0x17, 0x40, 0x00}; @@ -167,6 +194,45 @@ constexpr std::array kUuidBytes2 = { } // namespace +TEST(ParquetSplitTest, CanonicalStartUsesFirstColumnPages) { + for (auto dictionary_offset : + {std::optional{}, std::optional{0}, std::optional{80}, + std::optional{100}, std::optional{120}}) { + // The second column starts earlier and must not be used as the assignment key. + auto metadata = MakeMetadataWithPageOffsets({{100, dictionary_offset}, {50, 40}}); + const int64_t expected = dictionary_offset == 80 ? 80 : 100; + EXPECT_THAT(RowGroupStart(*metadata->RowGroup(0), 200), + HasValue(::testing::Eq(expected))); + } +} + +TEST(ParquetSplitTest, CanonicalStartRejectsMissingOrOutOfFilePages) { + for (int64_t offset : {-1, 0, 200, 201}) { + auto metadata = MakeMetadataWithPageOffsets({{offset, std::nullopt}}); + EXPECT_THAT(RowGroupStart(*metadata->RowGroup(0), 200), IsError(ErrorKind::kInvalid)); + } + auto no_columns = MakeMetadataWithPageOffsets({}); + EXPECT_THAT(RowGroupStart(*no_columns->RowGroup(0), 200), + HasErrorMessage("without columns")); +} + +TEST(ParquetSplitTest, MidpointAddsHalfTheRowGroupSize) { + // Each column chunk is 10 bytes, so the row group is 20 bytes from its start at 80. + auto metadata = MakeMetadataWithPageOffsets({{100, 80}, {50, 40}}); + ASSERT_EQ(metadata->RowGroup(0)->total_compressed_size(), 20); + EXPECT_THAT(RowGroupMidpoint(*metadata->RowGroup(0), 200), HasValue(::testing::Eq(90))); + // The size may reach, but not pass, the end of the file. + EXPECT_THAT(RowGroupMidpoint(*metadata->RowGroup(0), 100), HasValue(::testing::Eq(90))); + EXPECT_THAT(RowGroupMidpoint(*metadata->RowGroup(0), 99), + HasErrorMessage("invalid compressed size 20")); +} + +TEST(ParquetSplitTest, MidpointRejectsInvalidStart) { + auto metadata = MakeMetadataWithPageOffsets({{0, std::nullopt}}); + EXPECT_THAT(RowGroupMidpoint(*metadata->RowGroup(0), 200), + HasErrorMessage("starting page 0")); +} + class ParquetReaderTest : public TempFileTestBase { protected: static void SetUpTestSuite() { parquet::RegisterAll(); } @@ -251,7 +317,7 @@ class ParquetReaderTest : public TempFileTestBase { .properties = std::move(writer_properties)})); } - void CreateSplitParquetFile() { + void CreateSplitParquetFile(int64_t row_group_size = 2, bool dictionary = true) { const std::string kParquetFieldIdKey = "PARQUET:field_id"; auto arrow_schema = ::arrow::schema( {::arrow::field("id", ::arrow::int32(), /*nullable=*/false, @@ -270,12 +336,30 @@ class ParquetReaderTest : public TempFileTestBase { auto io = internal::checked_cast(*file_io_); auto outfile = io.fs()->OpenOutputStream(temp_parquet_file_).ValueOrDie(); + ::parquet::WriterProperties::Builder properties; + if (!dictionary) { + properties.disable_dictionary(); + } ASSERT_TRUE(::parquet::arrow::WriteTable(*table, ::arrow::default_memory_pool(), - outfile, /*chunk_size=*/2) + outfile, row_group_size, properties.build()) .ok()); ASSERT_TRUE(outfile->Close().ok()); } + void ReorderRowGroups(const std::vector& row_groups) { + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto metadata = ::parquet::ReadMetaData(input); + auto file_size = input->GetSize().ValueOrDie(); + auto pages = input->ReadAt(0, file_size - metadata->size() - 8).ValueOrDie(); + ASSERT_TRUE(input->Close().ok()); + ICEBERG_UNWRAP_OR_FAIL(auto output, + arrow::OpenArrowOutputStream(file_io_, temp_parquet_file_)); + ASSERT_TRUE(output->Write(pages).ok()); + ::parquet::WriteFileMetaData(*metadata->Subset(row_groups), output.get()); + ASSERT_TRUE(output->Close().ok()); + } + std::shared_ptr RowLineageSchema() { return std::make_shared(std::vector{ SchemaField::MakeRequired(1, "id", int32()), @@ -1009,15 +1093,24 @@ TEST_F(ParquetReaderTest, EmptySplitKeepsPrunedNestedListTypes) { TEST_F(ParquetReaderTest, ReadSplit) { CreateSplitParquetFile(); - // Read split offsets + // Read row-group starts and midpoints auto io = internal::checked_cast(*file_io_); auto input_stream = io.fs()->OpenInputFile(temp_parquet_file_).ValueOrDie(); auto metadata = ::parquet::ReadMetaData(input_stream); - std::vector split_offsets; + const int64_t file_size = input_stream->GetSize().ValueOrDie(); + std::vector starts; + std::vector midpoints; for (int i = 0; i < metadata->num_row_groups(); ++i) { - split_offsets.push_back(static_cast(metadata->RowGroup(i)->file_offset())); + auto column = metadata->RowGroup(i)->ColumnChunk(0); + ASSERT_TRUE(column->has_dictionary_page()); + ASSERT_LT(column->dictionary_page_offset(), column->data_page_offset()); + starts.push_back(static_cast(column->dictionary_page_offset())); + ICEBERG_UNWRAP_OR_FAIL(auto midpoint, + RowGroupMidpoint(*metadata->RowGroup(i), file_size)); + ASSERT_GT(static_cast(midpoint), starts.back()); + midpoints.push_back(static_cast(midpoint)); } - ASSERT_EQ(split_offsets.size(), 2); + ASSERT_EQ(starts.size(), 2); auto schema = std::make_shared(std::vector{ SchemaField::MakeRequired(1, "id", int32()), @@ -1026,13 +1119,26 @@ TEST_F(ParquetReaderTest, ReadSplit) { std::vector splits = { {.offset = 0, .length = std::numeric_limits::max()}, - {.offset = split_offsets[0], .length = split_offsets[1] - split_offsets[0]}, - {.offset = split_offsets[1], .length = 1}, - {.offset = split_offsets[1] + 1, .length = std::numeric_limits::max()}, - {.offset = 0, .length = split_offsets[0]}, + {.offset = starts[0], .length = starts[1] - starts[0]}, + {.offset = midpoints[1], .length = 1}, + {.offset = midpoints[1] + 1, .length = std::numeric_limits::max()}, + {.offset = 0, .length = midpoints[0]}, + {.offset = midpoints[0] + 1, .length = midpoints[1] - midpoints[0]}, + {.offset = starts[0], .length = 0}, + // A row group belongs to the split containing its midpoint, not its start. + {.offset = starts[1], .length = 1}, + {.offset = starts[0] + 1, .length = midpoints[0] - starts[0]}, }; std::vector expected_json = { - R"([[1, 0], [2, 1], [3, 2]])", R"([[1, 0], [2, 1]])", R"([[3, 2]])", "", "", + R"([[1, 0], [2, 1], [3, 2]])", + R"([[1, 0], [2, 1]])", + R"([[3, 2]])", + "", + "", + R"([[3, 2]])", + "", + "", + R"([[1, 0], [2, 1]])", }; ReaderProperties reader_properties; @@ -1056,6 +1162,222 @@ TEST_F(ParquetReaderTest, ReadSplit) { } } +TEST_F(ParquetReaderTest, ReadSplitWithoutDictionaryPages) { + ASSERT_NO_FATAL_FAILURE( + CreateSplitParquetFile(/*row_group_size=*/2, /*dictionary=*/false)); + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto metadata = ::parquet::ReadMetaData(input); + auto column = metadata->RowGroup(1)->ColumnChunk(0); + ASSERT_FALSE(column->has_dictionary_page()); + const int64_t file_size = input->GetSize().ValueOrDie(); + EXPECT_THAT(RowGroupStart(*metadata->RowGroup(1), file_size), + HasValue(::testing::Eq(column->data_page_offset()))); + ICEBERG_UNWRAP_OR_FAIL(auto midpoint, + RowGroupMidpoint(*metadata->RowGroup(1), file_size)); + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), MetadataColumns::kRowPosition}); + ICEBERG_UNWRAP_OR_FAIL( + auto reader, + ReaderFactoryRegistry::Open( + FileFormatType::kParquet, + {.path = temp_parquet_file_, + .split = Split{.offset = static_cast(midpoint), .length = 1}, + .io = file_io_, + .projection = schema})); + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[3, 2]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); +} + +TEST_F(ParquetReaderTest, ReadSplitsWithOmittedRowGroupOffsets) { + auto arrow_schema = ::arrow::schema( + {::arrow::field("id", ::arrow::int32(), /*nullable=*/false, + ::arrow::KeyValueMetadata::Make({"PARQUET:field_id"}, {"1"}))}); + auto table = ::arrow::Table::Make( + arrow_schema, + {::arrow::json::ArrayFromJSONString(::arrow::int32(), "[1, 2, 3, 4, 5, 6]") + .ValueOrDie()}); + ::parquet::WriterProperties::Builder properties_builder; + properties_builder.version(::parquet::ParquetVersion::PARQUET_2_6) + ->compression(::parquet::Compression::UNCOMPRESSED); + auto properties = properties_builder.build(); + ICEBERG_UNWRAP_OR_FAIL(auto original_output, + arrow::OpenArrowOutputStream(file_io_, temp_parquet_file_)); + ASSERT_TRUE(::parquet::arrow::WriteTable(*table, ::arrow::default_memory_pool(), + original_output, /*chunk_size=*/2, properties) + .ok()); + ASSERT_TRUE(original_output->Close().ok()); + ICEBERG_UNWRAP_OR_FAIL(auto original_input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto original_metadata = ::parquet::ReadMetaData(original_input); + auto original_size = original_input->GetSize().ValueOrDie(); + auto pages = original_input->ReadAt(0, original_size - original_metadata->size() - 8) + .ValueOrDie(); + ASSERT_TRUE(original_input->Close().ok()); + + auto builder = + ::parquet::FileMetaDataBuilder::Make(original_metadata->schema(), properties); + for (int i = 0; i < original_metadata->num_row_groups(); ++i) { + auto original_row_group = original_metadata->RowGroup(i); + auto row_group = builder->AppendRowGroup(); + row_group->set_num_rows(original_row_group->num_rows()); + for (int j = 0; j < original_row_group->num_columns(); ++j) { + auto column = original_row_group->ColumnChunk(j); + row_group->NextColumnChunk()->Finish( + column->num_values(), column->dictionary_page_offset(), + /*index_page_offset=*/0, column->data_page_offset(), + column->total_compressed_size(), column->total_uncompressed_size(), + column->has_dictionary_page(), /*dictionary_fallback=*/false, {}, {}); + } + // RowGroupMetaDataBuilder::Finish fills in the optional file_offset and + // total_compressed_size. Leave them unset to exercise legacy footer handling. + } + auto rebuilt_metadata = builder->Finish(original_metadata->key_value_metadata()); + ICEBERG_UNWRAP_OR_FAIL(auto output, + arrow::OpenArrowOutputStream(file_io_, temp_parquet_file_)); + ASSERT_TRUE(output->Write(pages).ok()); + ::parquet::WriteFileMetaData(*rebuilt_metadata, output.get()); + ASSERT_TRUE(output->Close().ok()); + + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto metadata = ::parquet::ReadMetaData(input); + ASSERT_EQ(metadata->num_columns(), 1); + ASSERT_EQ(metadata->num_rows(), 6); + ASSERT_EQ(metadata->num_row_groups(), 3); + for (int i = 0; i < metadata->num_row_groups(); ++i) { + auto row_group = metadata->RowGroup(i); + ASSERT_EQ(row_group->num_rows(), 2); + ASSERT_EQ(row_group->file_offset(), 0); + ASSERT_EQ(row_group->total_compressed_size(), 0); + auto column = row_group->ColumnChunk(0); + ASSERT_TRUE(column->has_dictionary_page()); + ASSERT_LT(column->dictionary_page_offset(), column->data_page_offset()); + } + auto boundary = static_cast( + metadata->RowGroup(1)->ColumnChunk(0)->dictionary_page_offset()); + auto file_size = static_cast(input->GetSize().ValueOrDie()); + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), MetadataColumns::kRowPosition, + MetadataColumns::kRowId, MetadataColumns::kLastUpdatedSequenceNumber}); + const std::vector> splits{ + std::nullopt, Split{.offset = 0, .length = boundary}, + Split{.offset = boundary, .length = file_size - boundary}}; + const std::vector expected{ + R"([[1, 0, 100, 7], [2, 1, 101, 7], [3, 2, 102, 7], [4, 3, 103, 7], [5, 4, 104, 7], [6, 5, 105, 7]])", + R"([[1, 0, 100, 7], [2, 1, 101, 7]])", + R"([[3, 2, 102, 7], [4, 3, 103, 7], [5, 4, 104, 7], [6, 5, 105, 7]])"}; + for (size_t i = 0; i < splits.size(); ++i) { + SCOPED_TRACE(i); + ICEBERG_UNWRAP_OR_FAIL(auto reader, + ReaderFactoryRegistry::Open(FileFormatType::kParquet, + {.path = temp_parquet_file_, + .length = file_size, + .split = splits[i], + .io = file_io_, + .projection = schema, + .first_row_id = 100, + .data_sequence_number = 7})); + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, expected[i])); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); + } +} + +TEST_F(ParquetReaderTest, RejectsSplitWithInvalidPageOffsetsBeforeReading) { + auto schema = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "col0", int32())}); + for (int64_t offset : {int64_t{-1}, int64_t{0}, std::numeric_limits::max()}) { + auto metadata = MakeMetadataWithPageOffsets({{offset, std::nullopt}}); + ICEBERG_UNWRAP_OR_FAIL(auto output, + arrow::OpenArrowOutputStream(file_io_, temp_parquet_file_)); + ::parquet::WriteMetaDataFile(*metadata, output.get()); + ASSERT_TRUE(output->Close().ok()); + auto reader = ReaderFactoryRegistry::Open(FileFormatType::kParquet, + {.path = temp_parquet_file_, + .split = Split{.offset = 0, .length = 1}, + .io = file_io_, + .projection = schema}); + EXPECT_THAT(reader, IsError(ErrorKind::kInvalid)); + EXPECT_THAT(reader, HasErrorMessage("row group 0")); + EXPECT_THAT(reader, HasErrorMessage(temp_parquet_file_)); + } +} + +TEST_F(ParquetReaderTest, ReadSplitWithNonmonotonicRowGroupStarts) { + ASSERT_NO_FATAL_FAILURE(CreateSplitParquetFile(/*row_group_size=*/1)); + // Logical row order becomes [1, 3, 2], while the pages stay in physical order [1, 2, + // 3]. + ASSERT_NO_FATAL_FAILURE(ReorderRowGroups({0, 2, 1})); + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto metadata = ::parquet::ReadMetaData(input); + auto last_start = static_cast( + metadata->RowGroup(1)->ColumnChunk(0)->dictionary_page_offset()); + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), MetadataColumns::kRowPosition, + MetadataColumns::kRowId, MetadataColumns::kLastUpdatedSequenceNumber}); + + ICEBERG_UNWRAP_OR_FAIL( + auto reader, + ReaderFactoryRegistry::Open(FileFormatType::kParquet, + {.path = temp_parquet_file_, + .split = Split{.offset = 0, .length = last_start}, + .io = file_io_, + .projection = schema, + .first_row_id = 100, + .data_sequence_number = 7})); + // A batch can span skipped row groups while keeping absolute positions and lineage. + ASSERT_NO_FATAL_FAILURE( + VerifyNextBatch(*reader, R"([[1, 0, 100, 7], [2, 2, 102, 7]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); + + ICEBERG_UNWRAP_OR_FAIL( + auto tail, ReaderFactoryRegistry::Open( + FileFormatType::kParquet, + {.path = temp_parquet_file_, + .split = Split{.offset = last_start, + .length = std::numeric_limits::max()}, + .io = file_io_, + .projection = schema, + .first_row_id = 100, + .data_sequence_number = 7})); + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*tail, R"([[3, 1, 101, 7]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*tail)); +} + +TEST_F(ParquetReaderTest, ReadMetadataOnlySplitWithNonmonotonicRowGroups) { + ASSERT_NO_FATAL_FAILURE(CreateSplitParquetFile(/*row_group_size=*/1)); + ASSERT_NO_FATAL_FAILURE(ReorderRowGroups({0, 2, 1})); + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto metadata = ::parquet::ReadMetaData(input); + auto last_start = static_cast( + metadata->RowGroup(1)->ColumnChunk(0)->dictionary_page_offset()); + auto schema = std::make_shared( + std::vector{MetadataColumns::kRowPosition, MetadataColumns::kRowId, + MetadataColumns::kLastUpdatedSequenceNumber}); + + for (int64_t batch_size : {1, 100}) { + SCOPED_TRACE(batch_size); + ReaderProperties properties; + properties.Set(ReaderProperties::kBatchSize, batch_size); + ICEBERG_UNWRAP_OR_FAIL( + auto reader, + ReaderFactoryRegistry::Open(FileFormatType::kParquet, + {.path = temp_parquet_file_, + .split = Split{.offset = 0, .length = last_start}, + .io = file_io_, + .projection = schema, + .first_row_id = 100, + .data_sequence_number = 7, + .properties = properties})); + // Arrow emits one batch per row group when no physical columns are read. + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[0, 100, 7]])")); + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[2, 102, 7]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); + } +} + TEST_F(ParquetReaderTest, ReadMetadataColumns) { temp_parquet_file_ = "metadata_columns.parquet"; CreateSimpleParquetFile();