feat(scan): read Parquet files by task range (#1002)

Author: kamcheungting-db

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));
 

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;

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

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.

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,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
@@ -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.

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);
 }