Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 27 additions & 9 deletions src/iceberg/data/file_scan_task_reader.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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));
Expand All @@ -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));

Expand Down
7 changes: 5 additions & 2 deletions src/iceberg/data/file_scan_task_reader.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
56 changes: 56 additions & 0 deletions src/iceberg/parquet/parquet_data_util.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down
15 changes: 15 additions & 0 deletions src/iceberg/parquet/parquet_data_util_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
81 changes: 54 additions & 27 deletions src/iceberg/parquet/parquet_reader.cc
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
#include "iceberg/parquet/parquet_reader.h"

#include <algorithm>
#include <limits>
#include <numeric>
#include <utility>
#include <variant>
#include <vector>

Expand Down Expand Up @@ -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
Expand All @@ -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_));
Expand Down Expand Up @@ -388,45 +390,70 @@ 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
// target of ProjectRecordBatch, so it must describe the projected schema rather than
// 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
Expand All @@ -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.
Expand Down
16 changes: 10 additions & 6 deletions src/iceberg/table_scan.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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);
}

Expand Down
Loading
Loading