This is an automated email from the ASF dual-hosted git repository.
wgtmac pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-cpp.git
The following commit(s) were added to refs/heads/main by this push:
new 7f8009f8 feat(scan): read Parquet files by task range (#1002)
7f8009f8 is described below
commit 7f8009f88372a6adc950fb0ae6457e50273c2aa1
Author: kamcheungting-db <[email protected]>
AuthorDate: Tue Oct 20 20:22:01 2026 -0700
feat(scan): read Parquet files by task range (#1002)
---
src/iceberg/data/file_scan_task_reader.cc | 36 ++-
src/iceberg/data/file_scan_task_reader.h | 7 +-
src/iceberg/parquet/parquet_data_util.cc | 56 ++++
src/iceberg/parquet/parquet_data_util_internal.h | 15 +
src/iceberg/parquet/parquet_reader.cc | 81 ++++--
src/iceberg/table_scan.cc | 16 +-
src/iceberg/table_scan.h | 16 +-
src/iceberg/test/file_scan_task_reader_test.cc | 175 +++++++++++-
src/iceberg/test/file_scan_task_test.cc | 35 ++-
src/iceberg/test/parquet_test.cc | 346 ++++++++++++++++++++++-
10 files changed, 713 insertions(+), 70 deletions(-)
diff --git a/src/iceberg/data/file_scan_task_reader.cc
b/src/iceberg/data/file_scan_task_reader.cc
index 7d3a8bc7..8c55e31b 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<FileIO> io,
+ReaderOptions MakeReaderOptions(const FileScanTask& task,
std::shared_ptr<FileIO> io,
std::shared_ptr<Schema> projection,
std::shared_ptr<Expression> filter,
std::shared_ptr<NameMapping> name_mapping,
ReaderProperties properties,
std::optional<int64_t> first_row_id,
std::optional<int64_t> data_sequence_number) {
+ const auto& data_file = *task.data_file();
return ReaderOptions{
.path = data_file.file_path,
.length = static_cast<size_t>(data_file.file_size_in_bytes),
+ .split = task.is_split() ? std::make_optional(
+ Split{.offset =
static_cast<size_t>(task.start()),
+ .length =
static_cast<size_t>(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<size_t>(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 c34f3e7a..33a4b82f 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<ArrowArrayStream> 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 449657df..73597cbf 100644
--- a/src/iceberg/parquet/parquet_data_util.cc
+++ b/src/iceberg/parquet/parquet_data_util.cc
@@ -24,6 +24,7 @@
#include <arrow/compute/api.h>
#include <arrow/record_batch.h>
#include <arrow/type.h>
+#include <parquet/metadata.h>
#include "iceberg/arrow/arrow_status_internal.h"
#include "iceberg/arrow/literal_util_internal.h"
@@ -37,6 +38,61 @@
namespace iceberg::parquet {
+Result<int64_t> 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<int64_t> 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 ad70c223..afd0ea39 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<int64_t> 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<int64_t> 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 90fa0646..12da76a9 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 <algorithm>
+#include <limits>
#include <numeric>
+#include <utility>
#include <variant>
#include <vector>
@@ -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<ReadContext>();
-
- // 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<int> 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<uint64_t>(*midpoint) - split.offset < split.length) {
+ row_group_indices_.push_back(i);
+ }
+ if (row_group->num_rows() < 0 ||
+ row_group->num_rows() > std::numeric_limits<int64_t>::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<ReadContext>();
// 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> split_;
+ // Row groups in footer order, shared by the batch and row-number readers.
+ std::vector<int> 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 25548bc7..d04e3385 100644
--- a/src/iceberg/table_scan.cc
+++ b/src/iceberg/table_scan.cc
@@ -224,7 +224,7 @@ Result<std::shared_ptr<FileScanTask>>
FileScanTask::MakeSplit(
return std::shared_ptr<FileScanTask>(
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<int128_t>(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<int128_t>(range_->start) * record_count / range_->file_size;
+ static_cast<int128_t>(range_->start) * record_count / file_size;
return static_cast<int64_t>(count_at_end - count_at_start);
}
diff --git a/src/iceberg/table_scan.h b/src/iceberg/table_scan.h
index 7afaa79d..1307089e 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<DataFile> data_file,
std::vector<std::shared_ptr<DataFile>> delete_files =
{},
std::shared_ptr<Expression> 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<std::shared_ptr<FileScanTask>> MakeSplit(
std::shared_ptr<DataFile> data_file, int64_t start, int64_t length,
std::vector<std::shared_ptr<DataFile>> 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<DataFile> 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 11366f72..05416bd7 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 <algorithm>
+#include <array>
#include <memory>
#include <optional>
#include <string>
+#include <tuple>
#include <utility>
#include <vector>
@@ -31,11 +34,13 @@
#include <arrow/record_batch.h>
#include <arrow/table.h>
#include <parquet/arrow/writer.h>
+#include <parquet/file_reader.h>
#include <parquet/metadata.h>
#include <parquet/properties.h>
#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<int64_t, 4>;
+
+ Result<std::vector<RowWithPositionAndLineage>> 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<RowWithPositionAndLineage> 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<const
::arrow::Int32Array&>(*batch->column(0));
+ const auto& positions = static_cast<const
::arrow::Int64Array&>(*batch->column(1));
+ const auto& row_ids = static_cast<const
::arrow::Int64Array&>(*batch->column(2));
+ const auto& sequences = static_cast<const
::arrow::Int64Array&>(*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<ExportedBatch> 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<FileScanTask>(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<std::tuple<std::string, int64_t>>
{};
+
+TEST_P(FileScanTaskReaderSplitTest,
SplitsMatchWholeFileWithAbsolutePositionsAndLineage) {
+ const auto& [delete_kind, batch_size] = GetParam();
+ ICEBERG_UNWRAP_OR_FAIL(
+ auto data_file,
+ MakeDataFile(table_schema_,
+ std::vector<std::string>{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<std::shared_ptr<DataFile>> 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<Schema>(std::vector<SchemaField>{
+ 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<test::StdFileIO>(),
+ .table_schema = table_schema_,
+ .projected_schema = projection,
+ .properties = {{"read.batch-size",
std::to_string(batch_size)}},
+ }));
+
+ std::vector<RowWithPositionAndLineage> 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<int64_t> 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<RowWithPositionAndLineage> 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 6394c96f..41861cae 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<DataFile>();
+ 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<DataFile>();
+ data_file->file_path = "test.parquet";
+ data_file->file_size_in_bytes = 100;
+ data_file->record_count = std::numeric_limits<int64_t>::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<DataFile>();
+ 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 9fbeae87..1327f03f 100644
--- a/src/iceberg/test/parquet_test.cc
+++ b/src/iceberg/test/parquet_test.cc
@@ -39,7 +39,9 @@
#include <parquet/arrow/reader.h>
#include <parquet/arrow/writer.h>
#include <parquet/file_reader.h>
+#include <parquet/file_writer.h>
#include <parquet/metadata.h>
+#include <parquet/schema.h>
#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<ParquetCodec> FirstUnavailableParquetCodec()
{
return std::nullopt;
}
+std::unique_ptr<::parquet::FileMetaData> MakeMetadataWithPageOffsets(
+ const std::vector<std::pair<int64_t, std::optional<int64_t>>>&
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<uint8_t, Uuid::kLength> kUuidBytes1 = {
0x12, 0x3e, 0x45, 0x67, 0xe8, 0x9b, 0x12, 0xd3,
0xa4, 0x56, 0x42, 0x66, 0x14, 0x17, 0x40, 0x00};
@@ -167,6 +194,45 @@ constexpr std::array<uint8_t, Uuid::kLength> kUuidBytes2 =
{
} // namespace
+TEST(ParquetSplitTest, CanonicalStartUsesFirstColumnPages) {
+ for (auto dictionary_offset :
+ {std::optional<int64_t>{}, std::optional<int64_t>{0},
std::optional<int64_t>{80},
+ std::optional<int64_t>{100}, std::optional<int64_t>{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<arrow::ArrowFileSystemFileIO&>(*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<int>& 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<Schema> RowLineageSchema() {
return std::make_shared<Schema>(std::vector<SchemaField>{
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<arrow::ArrowFileSystemFileIO&>(*file_io_);
auto input_stream = io.fs()->OpenInputFile(temp_parquet_file_).ValueOrDie();
auto metadata = ::parquet::ReadMetaData(input_stream);
- std::vector<size_t> split_offsets;
+ const int64_t file_size = input_stream->GetSize().ValueOrDie();
+ std::vector<size_t> starts;
+ std::vector<size_t> midpoints;
for (int i = 0; i < metadata->num_row_groups(); ++i) {
-
split_offsets.push_back(static_cast<size_t>(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<size_t>(column->dictionary_page_offset()));
+ ICEBERG_UNWRAP_OR_FAIL(auto midpoint,
+ RowGroupMidpoint(*metadata->RowGroup(i),
file_size));
+ ASSERT_GT(static_cast<size_t>(midpoint), starts.back());
+ midpoints.push_back(static_cast<size_t>(midpoint));
}
- ASSERT_EQ(split_offsets.size(), 2);
+ ASSERT_EQ(starts.size(), 2);
auto schema = std::make_shared<Schema>(std::vector<SchemaField>{
SchemaField::MakeRequired(1, "id", int32()),
@@ -1026,13 +1119,26 @@ TEST_F(ParquetReaderTest, ReadSplit) {
std::vector<Split> splits = {
{.offset = 0, .length = std::numeric_limits<size_t>::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<size_t>::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<size_t>::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<std::string> 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<Schema>(std::vector<SchemaField>{
+ 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<size_t>(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<size_t>(
+ metadata->RowGroup(1)->ColumnChunk(0)->dictionary_page_offset());
+ auto file_size = static_cast<size_t>(input->GetSize().ValueOrDie());
+ auto schema = std::make_shared<Schema>(std::vector<SchemaField>{
+ SchemaField::MakeRequired(1, "id", int32()),
MetadataColumns::kRowPosition,
+ MetadataColumns::kRowId, MetadataColumns::kLastUpdatedSequenceNumber});
+ const std::vector<std::optional<Split>> splits{
+ std::nullopt, Split{.offset = 0, .length = boundary},
+ Split{.offset = boundary, .length = file_size - boundary}};
+ const std::vector<std::string> 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<Schema>(
+ std::vector<SchemaField>{SchemaField::MakeRequired(1, "col0", int32())});
+ for (int64_t offset : {int64_t{-1}, int64_t{0},
std::numeric_limits<int64_t>::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<size_t>(
+ metadata->RowGroup(1)->ColumnChunk(0)->dictionary_page_offset());
+ auto schema = std::make_shared<Schema>(std::vector<SchemaField>{
+ 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<size_t>::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<size_t>(
+ metadata->RowGroup(1)->ColumnChunk(0)->dictionary_page_offset());
+ auto schema = std::make_shared<Schema>(
+ std::vector<SchemaField>{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();
