diff --git a/src/iceberg/catalog/rest/json_serde.cc b/src/iceberg/catalog/rest/json_serde.cc index ba348dec3f..54feb17263 100644 --- a/src/iceberg/catalog/rest/json_serde.cc +++ b/src/iceberg/catalog/rest/json_serde.cc @@ -110,6 +110,8 @@ constexpr std::string_view kDeleteFileReferences = "delete-file-references"; constexpr std::string_view kResidualFilter = "residual-filter"; constexpr std::string_view kStart = "start"; constexpr std::string_view kLength = "length"; +constexpr std::string_view kMapKeys = "keys"; +constexpr std::string_view kMapValues = "values"; Result StorageCredentialToJson(const StorageCredential& credential) { ICEBERG_RETURN_UNEXPECTED(credential.Validate()); @@ -168,14 +170,19 @@ Result>> FileScanTasksFromJson( ICEBERG_ASSIGN_OR_RAISE(auto data_file_json, GetJsonValue(task_json, kDataFile)); ICEBERG_ASSIGN_OR_RAISE( - auto data_file, - iceberg::DataFileFromJson(data_file_json, partition_spec_by_id, schema)); - ICEBERG_ASSIGN_OR_RAISE(auto start, GetJsonIntegerOrDefault(task_json, kStart, 0)); - ICEBERG_ASSIGN_OR_RAISE( - auto length, - GetJsonIntegerOrDefault(task_json, kLength, data_file.file_size_in_bytes)); - ICEBERG_RETURN_UNEXPECTED( - CheckFileScanTaskNotSplit(start, length, data_file.file_size_in_bytes)); + auto data_file, DataFileFromJson(data_file_json, partition_spec_by_id, schema)); + // Neither field means the whole file. Only one field is ambiguous. + const bool has_start = task_json.contains(kStart); + const bool has_length = task_json.contains(kLength); + if (has_start != has_length) { + return JsonParseError("'start' and 'length' must be provided together"); + } + int64_t start = 0; + int64_t length = data_file.file_size_in_bytes; + if (has_start) { + ICEBERG_ASSIGN_OR_RAISE(start, GetJsonInteger(task_json, kStart)); + ICEBERG_ASSIGN_OR_RAISE(length, GetJsonInteger(task_json, kLength)); + } // FIXME: REST scan-task DataFile JSON currently carries first-row-id, // but not the manifest-entry data sequence number. Until the REST API exposes // it, REST-planned tasks cannot inherit _last_updated_sequence_number. @@ -203,9 +210,13 @@ Result>> FileScanTasksFromJson( ICEBERG_ASSIGN_OR_RAISE(residual_filter, ExpressionFromJson(filter_json)); } - file_scan_tasks.push_back(std::make_shared( - std::make_shared(std::move(data_file)), std::move(task_delete_files), - std::move(residual_filter))); + auto task = FileScanTask::MakeSplit(std::make_shared(std::move(data_file)), + start, length, std::move(task_delete_files), + std::move(residual_filter)); + if (!task) { + return JsonParseError("Invalid FileScanTask range: {}", task.error().message); + } + file_scan_tasks.push_back(std::move(task.value())); } return file_scan_tasks; } @@ -253,10 +264,19 @@ Result ScanTaskFieldsToJson( if (!task) continue; nlohmann::json task_json; if (task->data_file()) { + // The shared data file may have changed size; don't write an invalid range. + auto range_result = + FileScanTask::MakeSplit(task->data_file(), task->start(), task->length()); + if (!range_result) { + return ValidationFailed("Invalid FileScanTask range: {}", + range_result.error().message); + } ICEBERG_ASSIGN_OR_RAISE( auto data_file_json, iceberg::ToJson(*task->data_file(), partition_specs_by_id, schema)); task_json[kDataFile] = std::move(data_file_json); + task_json[kStart] = task->start(); + task_json[kLength] = task->length(); } if (!task->delete_files().empty()) { std::vector refs; diff --git a/src/iceberg/catalog/rest/types.cc b/src/iceberg/catalog/rest/types.cc index 84fba9a7ca..ebc2c2876f 100644 --- a/src/iceberg/catalog/rest/types.cc +++ b/src/iceberg/catalog/rest/types.cc @@ -192,6 +192,7 @@ bool FileScanTaskEqual(const std::shared_ptr& lhs, return false; } return SharedPtrEqual(lhs->data_file(), rhs->data_file()) && + lhs->start() == rhs->start() && lhs->length() == rhs->length() && SharedPtrVectorEqual(lhs->delete_files(), rhs->delete_files()) && ExpressionPtrEqual(lhs->residual_filter(), rhs->residual_filter()); } diff --git a/src/iceberg/data/file_scan_task_reader.cc b/src/iceberg/data/file_scan_task_reader.cc index 0c41f25eda..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,10 +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()) { + 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)); @@ -191,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 a71ef5f84f..33a4b82fff 100644 --- a/src/iceberg/data/file_scan_task_reader.h +++ b/src/iceberg/data/file_scan_task_reader.h @@ -67,6 +67,12 @@ class ICEBERG_DATA_EXPORT FileScanTaskReader { ~FileScanTaskReader(); /// \brief Open a task and return an Arrow C stream for its projected live rows. + /// + /// 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/json_serde.cc b/src/iceberg/json_serde.cc index c1c3ffb73f..99b2b5e7fb 100644 --- a/src/iceberg/json_serde.cc +++ b/src/iceberg/json_serde.cc @@ -2423,8 +2423,8 @@ Result ToJson( } json[kSpec] = ToJson(*spec_it->second); json[kDataFile] = std::move(data_file_json); - json[kStart] = 0; - json[kLength] = task.data_file()->file_size_in_bytes; + json[kStart] = task.start(); + json[kLength] = task.length(); nlohmann::json delete_files_json = nlohmann::json::array(); for (const auto& delete_file : task.delete_files()) { @@ -2457,17 +2457,6 @@ Result ToJson( return json; } -Status CheckFileScanTaskNotSplit(int64_t start, int64_t length, - int64_t file_size_in_bytes) { - if (start != 0 || length != file_size_in_bytes) { - return NotSupported( - "Split FileScanTask is not supported: start={}, length={}, " - "file-size-in-bytes={}", - start, length, file_size_in_bytes); - } - return {}; -} - Result> FileScanTaskFromJson(const nlohmann::json& json) { if (!json.is_object()) { return JsonParseError("Cannot parse file scan task from a non-object: {}", @@ -2510,8 +2499,6 @@ Result> FileScanTaskFromJson(const nlohmann::json& ICEBERG_ASSIGN_OR_RAISE(auto start, GetJsonInteger(json, kStart)); ICEBERG_ASSIGN_OR_RAISE(auto length, GetJsonInteger(json, kLength)); - ICEBERG_RETURN_UNEXPECTED( - CheckFileScanTaskNotSplit(start, length, data_file.file_size_in_bytes)); std::vector> delete_files; if (json.contains(kDeleteFiles)) { @@ -2539,9 +2526,13 @@ Result> FileScanTaskFromJson(const nlohmann::json& ICEBERG_ASSIGN_OR_RAISE(residual_filter, ExpressionFromJson(filter_json)); } - return std::make_shared(std::make_shared(std::move(data_file)), - std::move(delete_files), - std::move(residual_filter)); + auto task = FileScanTask::MakeSplit(std::make_shared(std::move(data_file)), + start, length, std::move(delete_files), + std::move(residual_filter)); + if (!task) { + return JsonParseError("Invalid FileScanTask range: {}", task.error().message); + } + return std::move(task.value()); } } // namespace iceberg diff --git a/src/iceberg/json_serde_internal.h b/src/iceberg/json_serde_internal.h index 2465df1404..e37367826e 100644 --- a/src/iceberg/json_serde_internal.h +++ b/src/iceberg/json_serde_internal.h @@ -455,8 +455,8 @@ ICEBERG_EXPORT Result DataFileFromJson( /// - `schema` (required): schema used by this task /// - `spec` (required): partition spec used by this task /// - `data-file` (required): ContentFile JSON -/// - `start` (required): always 0 because split tasks are unsupported -/// - `length` (required): always the data file size +/// - `start` (required): inclusive byte offset where the task starts +/// - `length` (required): number of bytes in the task range /// - `delete-files` (optional on input; omitted when the task has no delete files): /// array of ContentFile JSON /// - `residual-filter` (optional): Expression JSON @@ -469,13 +469,10 @@ ICEBERG_EXPORT Result ToJson( /// \brief Deserializes a self-contained FileScanTask JSON object. /// /// Reads the embedded schema and partition spec. Rejects REST -/// `delete-file-references` and byte-range splits (`start` != 0 or `length` != file -/// size). Both range fields are required, matching Java core. +/// `delete-file-references`. Both range fields are required, matching Java core. +/// Valid ranges are preserved. C++ additionally rejects negative, out-of-file, +/// and empty partial ranges; an empty range is valid only for a zero-byte file. ICEBERG_EXPORT Result> FileScanTaskFromJson( const nlohmann::json& json); -/// Rejects `start != 0 || length != file_size_in_bytes`. -ICEBERG_EXPORT Status CheckFileScanTaskNotSplit(int64_t start, int64_t length, - int64_t file_size_in_bytes); - } // namespace iceberg diff --git a/src/iceberg/parquet/parquet_data_util.cc b/src/iceberg/parquet/parquet_data_util.cc index 052d03b550..d8b44bb1b9 100644 --- a/src/iceberg/parquet/parquet_data_util.cc +++ b/src/iceberg/parquet/parquet_data_util.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow/literal_util_internal.h" @@ -35,6 +36,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-mr, 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 5aad68f8f2..efe4831229 100644 --- a/src/iceberg/parquet/parquet_data_util_internal.h +++ b/src/iceberg/parquet/parquet_data_util_internal.h @@ -24,8 +24,23 @@ #include "iceberg/arrow/metadata_column_util_internal.h" #include "iceberg/schema_util.h" +namespace parquet { +class RowGroupMetaData; +} // namespace parquet + namespace iceberg::parquet { +/// \brief Return the first column's starting page, matching parquet-mr'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 1f81461078..292181fd6c 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 @@ -270,6 +272,13 @@ struct ReadContext { std::unique_ptr<::arrow::RecordBatchReader> record_batch_reader_; }; +// Consecutive row groups can share a batch reader without introducing gaps in _pos. +struct RowGroupRange { + int begin; + int end; + int64_t first_row_position; +}; + // TODO(gangwu): list of work items // 1. Make the memory pool configurable // 2. Catch ParquetException and convert to Status/Result @@ -283,7 +292,6 @@ class ParquetReader::Impl { return InvalidArgument("Projected schema is required by Parquet reader"); } - split_ = options.split; read_schema_ = options.projection; // Prepare reader properties @@ -304,6 +312,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(PlanRowGroupRanges(options)); // Project read schema onto the Parquet file schema ICEBERG_ASSIGN_OR_RAISE(projection_, BuildProjection(reader_.get(), *read_schema_)); @@ -321,9 +330,19 @@ class ParquetReader::Impl { ICEBERG_RETURN_UNEXPECTED(InitReadContext()); } - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, context_->record_batch_reader_->Next()); - if (!batch) { - return std::nullopt; + // Each row-group range has its own batch reader. When one runs out, open the next. + std::shared_ptr<::arrow::RecordBatch> batch; + while (true) { + ICEBERG_ARROW_ASSIGN_OR_RETURN(batch, context_->record_batch_reader_->Next()); + if (batch) { + break; + } + if (next_row_group_range_ == row_group_ranges_.size()) { + return std::nullopt; + } + ICEBERG_ARROW_RETURN_NOT_OK(context_->record_batch_reader_->Close()); + ICEBERG_ASSIGN_OR_RAISE(context_->record_batch_reader_, OpenNextRowGroupRange()); + ConsumeRowGroupRange(); } ICEBERG_ASSIGN_OR_RAISE( @@ -387,47 +406,89 @@ class ParquetReader::Impl { } private: - Status InitReadContext() { - context_ = std::make_unique(); - - // Row group pruning based on the split + Status PlanRowGroupRanges(const ReaderOptions& options) { // TODO(gangwu): add row group filtering based on zone map, bloom filter, etc. - std::vector row_group_indices; - if (split_.has_value()) { - auto metadata = reader_->parquet_reader()->metadata(); + auto metadata = reader_->parquet_reader()->metadata(); + if (options.split.has_value()) { + const auto& split = *options.split; + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto file_size, input_stream_->GetSize()); + // Rows in all earlier row groups, selected or not, so _pos is file-relative. + int64_t row_position = 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; - } else { - metadata_context_.next_file_pos += metadata->RowGroup(i)->num_rows(); + 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) { + // Extend the last range if this row group follows it; otherwise start a + // new range so its _pos can jump past the skipped row groups. + if (!row_group_ranges_.empty() && row_group_ranges_.back().end == i) { + row_group_ranges_.back().end = i + 1; + } else { + row_group_ranges_.push_back( + {.begin = i, .end = i + 1, .first_row_position = row_position}); + } + } + if (row_group->num_rows() < 0 || + row_group->num_rows() > std::numeric_limits::max() - row_position) { + return Invalid("Invalid Parquet row count in file {}, row group {}", + options.path, i); + } + row_position += row_group->num_rows(); } - } else { - row_group_indices.resize(reader_->parquet_reader()->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_ranges_.push_back( + {.begin = 0, .end = metadata->num_row_groups(), .first_row_position = 0}); } + return {}; + } - // Create the record batch reader - if (row_group_indices.empty()) { - // None of the row groups are selected, return an empty record batch reader - context_->record_batch_reader_ = std::make_unique(); - } else { - auto column_indices = SelectedColumnIndices(projection_); - ICEBERG_ARROW_ASSIGN_OR_RETURN( - context_->record_batch_reader_, - reader_->GetRecordBatchReader(row_group_indices, column_indices)); + // Open the next range without consuming it, so a caller that fails afterwards can + // retry it. Call ConsumeRowGroupRange() once the reader is installed. + Result> OpenNextRowGroupRange() { + if (next_row_group_range_ == row_group_ranges_.size()) { + return std::make_unique(); } + const auto& range = row_group_ranges_[next_row_group_range_]; + std::vector row_group_indices(range.end - range.begin); + std::iota(row_group_indices.begin(), row_group_indices.end(), range.begin); // NOLINT + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto batch_reader, reader_->GetRecordBatchReader( + row_group_indices, SelectedColumnIndices(projection_))); + return batch_reader; + } + + void ConsumeRowGroupRange() { + if (next_row_group_range_ == row_group_ranges_.size()) { + return; + } + metadata_context_.next_file_pos = + row_group_ranges_[next_row_group_range_].first_row_position; + ++next_row_group_range_; + } + + Status InitReadContext() { + // Build locally so a failure leaves context_ unset and the first range unconsumed. + auto context = std::make_unique(); + ICEBERG_ASSIGN_OR_RAISE(context->record_batch_reader_, OpenNextRowGroupRange()); // 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 @@ -438,18 +499,20 @@ 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); + ConsumeRowGroupRange(); 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_; + // Selected row groups in file row order, separated at gaps to preserve absolute _pos. + std::vector row_group_ranges_; + size_t next_row_group_range_ = 0; // 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 46f235fbe2..b83599e2c4 100644 --- a/src/iceberg/table_scan.cc +++ b/src/iceberg/table_scan.cc @@ -37,6 +37,8 @@ #include "iceberg/snapshot.h" #include "iceberg/table.h" #include "iceberg/table_metadata.h" +#include "iceberg/util/content_file_util.h" +#include "iceberg/util/int128.h" #include "iceberg/util/macros.h" #include "iceberg/util/snapshot_util.h" #include "iceberg/util/timepoint.h" @@ -179,11 +181,78 @@ FileScanTask::FileScanTask(std::shared_ptr data_file, ICEBERG_DCHECK(data_file_ != nullptr, "Data file cannot be null for FileScanTask"); } -int64_t FileScanTask::size_bytes() const { return data_file_->file_size_in_bytes; } +FileScanTask::FileScanTask(std::shared_ptr data_file, + std::vector> delete_files, + std::shared_ptr residual_filter, Range range) + : FileScanTask(std::move(data_file), std::move(delete_files), + std::move(residual_filter)) { + range_ = range; +} + +Result> FileScanTask::MakeSplit( + std::shared_ptr data_file, int64_t start, int64_t length, + std::vector> delete_files, + std::shared_ptr filter) { + if (data_file == nullptr) { + return InvalidArgument("Cannot create a file scan task without a data file"); + } + + const int64_t file_size = data_file->file_size_in_bytes; + if (file_size < 0) { + return InvalidArgument("Cannot split file {} with negative size {}", + data_file->file_path, file_size); + } + if (start < 0 || start > file_size) { + return InvalidArgument("Split start {} is outside file {} of size {}", start, + data_file->file_path, file_size); + } + // Subtraction is safe after validating start, even when file_size is INT64_MAX. + if (length < 0 || length > file_size - start) { + return InvalidArgument("Split length {} from start {} is outside file {} of size {}", + length, start, data_file->file_path, file_size); + } + // A full-file range is a whole-file task (is_split() is false). This is also the + // only valid empty range: [0, 0) of an empty file. + if (start == 0 && length == file_size) { + return std::make_shared(std::move(data_file), std::move(delete_files), + std::move(filter)); + } + if (length == 0) { + return InvalidArgument("A partial split of file {} must have positive length", + data_file->file_path); + } + + 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})); +} + +int64_t FileScanTask::length() const { + return range_ ? range_->length : data_file_->file_size_in_bytes; +} + +int64_t FileScanTask::size_bytes() const { return length(); } int32_t FileScanTask::files_count() const { return 1; } -int64_t FileScanTask::estimated_row_count() const { return data_file_->record_count; } +int64_t FileScanTask::estimated_row_count() const { + if (!range_) { + return data_file_->record_count; + } + + const int64_t record_count = data_file_->record_count; + if (record_count <= 0) { + return 0; + } + + // floor(end * rows / size) - floor(start * rows / size), so splits that cover + // the file sum exactly 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_start = + static_cast(range_->start) * record_count / range_->file_size; + return static_cast(count_at_end - count_at_start); +} // ChangelogScanTask implementation diff --git a/src/iceberg/table_scan.h b/src/iceberg/table_scan.h index 6ec29ad9f6..3d4893e758 100644 --- a/src/iceberg/table_scan.h +++ b/src/iceberg/table_scan.h @@ -22,6 +22,7 @@ /// \file iceberg/table_scan.h /// \brief Define table scan APIs and scan task types. +#include #include #include #include @@ -71,10 +72,32 @@ 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. explicit FileScanTask(std::shared_ptr data_file, std::vector> delete_files = {}, std::shared_ptr filter = nullptr); + /// \brief Make a task for a validated half-open byte range [start, start + length). + /// + /// 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. + static Result> MakeSplit( + std::shared_ptr data_file, int64_t start, int64_t length, + std::vector> delete_files = {}, + std::shared_ptr filter = nullptr); + + /// \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. + int64_t length() const; + + /// \brief Whether this task reads a proper subset of its data file. + bool is_split() const { return range_.has_value(); } + /// \brief The data file that should be read by this scan task. const std::shared_ptr& data_file() const { return data_file_; } @@ -92,9 +115,22 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask { int64_t estimated_row_count() const override; private: + // Only set for a partial file; whole-file tasks leave range_ empty. + 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, + std::vector> delete_files, + std::shared_ptr filter, Range range); + std::shared_ptr data_file_; std::vector> delete_files_; std::shared_ptr residual_filter_; + std::optional range_; }; /// \brief Stream of file scan tasks. diff --git a/src/iceberg/test/file_scan_task_reader_test.cc b/src/iceberg/test/file_scan_task_reader_test.cc index cca13ce0df..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,6 +440,161 @@ TEST_F(FileScanTaskReaderTest, OpenWithoutDeletesReadsProjectedSchema) { VerifyStream(&stream, R"([[1, "Foo"], [2, "Bar"], [3, "Baz"]])")); } +TEST_F(FileScanTaskReaderTest, RejectsPartialTasksForUnsupportedFormats) { + ICEBERG_UNWRAP_OR_FAIL( + auto data_file, + MakeDataFile(table_schema_, + R"([[1, "Foo", "blue"], [2, "Bar", "red"], [3, "Baz", "green"]])")); + ICEBERG_UNWRAP_OR_FAIL(auto task, FileScanTask::MakeSplit(data_file, 0, 1)); + + FileScanTaskReader::Options options{ + .io = file_io_, + .table_schema = table_schema_, + .schemas = {table_schema_}, + .projected_schema = projected_schema_, + }; + ICEBERG_UNWRAP_OR_FAIL(auto reader, FileScanTaskReader::Make(std::move(options))); + 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 cb945ba1fb..6394c96fed 100644 --- a/src/iceberg/test/file_scan_task_test.cc +++ b/src/iceberg/test/file_scan_task_test.cc @@ -17,6 +17,10 @@ * under the License. */ +#include +#include +#include + #include #include #include @@ -29,6 +33,7 @@ #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/data/file_scan_task_reader.h" +#include "iceberg/expression/expressions.h" #include "iceberg/file_format.h" #include "iceberg/manifest/manifest_entry.h" #include "iceberg/parquet/parquet_register.h" @@ -40,6 +45,116 @@ namespace iceberg { +TEST(FileScanTaskRangeTest, WholeFileAndEmptyFile) { + auto data_file = std::make_shared(); + data_file->file_path = "test.parquet"; + data_file->file_size_in_bytes = 100; + data_file->record_count = 7; + + FileScanTask whole(data_file); + EXPECT_FALSE(whole.is_split()); + EXPECT_EQ(whole.start(), 0); + EXPECT_EQ(whole.length(), 100); + EXPECT_EQ(whole.size_bytes(), 100); + EXPECT_EQ(whole.estimated_row_count(), 7); + + ICEBERG_UNWRAP_OR_FAIL(auto explicit_whole, FileScanTask::MakeSplit(data_file, 0, 100)); + EXPECT_FALSE(explicit_whole->is_split()); + EXPECT_EQ(explicit_whole->start(), 0); + 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)); + EXPECT_FALSE(empty->is_split()); + EXPECT_EQ(empty->start(), 0); + EXPECT_EQ(empty->length(), 0); + EXPECT_EQ(empty->size_bytes(), 0); + EXPECT_EQ(empty->estimated_row_count(), 0); +} + +TEST(FileScanTaskRangeTest, DistinctSplitsShareMetadataAndEstimateRowsOnce) { + auto data_file = std::make_shared(); + data_file->file_path = "test.parquet"; + data_file->file_size_in_bytes = 100; + data_file->record_count = 7; + data_file->first_row_id = 1000; + data_file->data_sequence_number = 22; + auto delete_file = std::make_shared(); + auto residual = Expressions::AlwaysTrue(); + + ICEBERG_UNWRAP_OR_FAIL( + auto first, FileScanTask::MakeSplit(data_file, 0, 40, {delete_file}, residual)); + ICEBERG_UNWRAP_OR_FAIL( + auto second, FileScanTask::MakeSplit(data_file, 40, 60, {delete_file}, residual)); + + EXPECT_TRUE(first->is_split()); + EXPECT_TRUE(second->is_split()); + EXPECT_EQ(first->start(), 0); + EXPECT_EQ(first->length(), 40); + EXPECT_EQ(second->start(), 40); + EXPECT_EQ(second->length(), 60); + EXPECT_EQ(first->size_bytes() + second->size_bytes(), 100); + EXPECT_EQ(first->estimated_row_count(), 2); + EXPECT_EQ(second->estimated_row_count(), 5); + EXPECT_EQ(first->estimated_row_count() + second->estimated_row_count(), 7); + EXPECT_EQ(first->data_file(), data_file); + EXPECT_EQ(second->data_file(), data_file); + ASSERT_EQ(first->delete_files().size(), 1); + ASSERT_EQ(second->delete_files().size(), 1); + EXPECT_EQ(first->delete_files()[0], delete_file); + EXPECT_EQ(second->delete_files()[0], delete_file); + EXPECT_EQ(first->residual_filter(), residual); + EXPECT_EQ(second->residual_filter(), residual); + EXPECT_EQ(first->data_file()->first_row_id, 1000); + EXPECT_EQ(first->data_file()->data_sequence_number, 22); + EXPECT_EQ(second->data_file()->first_row_id, 1000); + EXPECT_EQ(second->data_file()->data_sequence_number, 22); +} + +TEST(FileScanTaskRangeTest, RejectsInvalidRanges) { + auto data_file = std::make_shared(); + data_file->file_path = "test.parquet"; + data_file->file_size_in_bytes = 100; + + const auto expect_invalid = [&](int64_t start, int64_t length) { + auto result = FileScanTask::MakeSplit(data_file, start, length); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().kind, ErrorKind::kInvalidArgument); + }; + expect_invalid(-1, 1); + expect_invalid(0, -1); + expect_invalid(0, 0); + expect_invalid(100, 0); + expect_invalid(101, 1); + expect_invalid(99, 2); + + auto null_file = FileScanTask::MakeSplit(nullptr, 0, 1); + ASSERT_FALSE(null_file.has_value()); + EXPECT_EQ(null_file.error().kind, ErrorKind::kInvalidArgument); + + data_file->file_size_in_bytes = -1; + expect_invalid(0, 1); + + data_file->file_size_in_bytes = std::numeric_limits::max(); + expect_invalid(std::numeric_limits::max() - 1, 2); +} + +TEST(FileScanTaskRangeTest, RowEstimateDoesNotOverflow) { + auto data_file = std::make_shared(); + data_file->file_size_in_bytes = std::numeric_limits::max(); + data_file->record_count = std::numeric_limits::max(); + + ICEBERG_UNWRAP_OR_FAIL( + auto first, + FileScanTask::MakeSplit(data_file, 0, data_file->file_size_in_bytes - 1)); + ICEBERG_UNWRAP_OR_FAIL(auto last, FileScanTask::MakeSplit( + data_file, data_file->file_size_in_bytes - 1, 1)); + EXPECT_EQ(first->estimated_row_count(), data_file->record_count - 1); + EXPECT_EQ(last->estimated_row_count(), 1); +} + class FileScanTaskTest : public TempFileTestBase { protected: static void SetUpTestSuite() { parquet::RegisterAll(); } diff --git a/src/iceberg/test/json_serde_test.cc b/src/iceberg/test/json_serde_test.cc index 51bc2af11c..efcf961515 100644 --- a/src/iceberg/test/json_serde_test.cc +++ b/src/iceberg/test/json_serde_test.cc @@ -18,6 +18,7 @@ */ #include +#include #include #include #include @@ -1803,31 +1804,103 @@ TEST_F(FileScanTaskJsonTest, RequiresJavaCoreFields) { } } -TEST_F(FileScanTaskJsonTest, RejectsSplitTasksUsingExactOrCondition) { +TEST_F(FileScanTaskJsonTest, RoundTripsValidRanges) { struct Case { int64_t start; int64_t length; - bool supported; + bool is_split; }; - for (const auto& test_case : {Case{.start = 0, .length = 10, .supported = true}, - Case{.start = 1, .length = 10, .supported = false}, - Case{.start = 0, .length = 9, .supported = false}}) { + for (const auto& test_case : {Case{.start = 0, .length = 10, .is_split = false}, + Case{.start = 0, .length = 9, .is_split = true}, + Case{.start = 4, .length = 6, .is_split = true}}) { SCOPED_TRACE(testing::Message() << "start=" << test_case.start << ", length=" << test_case.length); auto json = JavaGolden(); json["start"] = test_case.start; json["length"] = test_case.length; + json["residual-filter"] = R"({"type":"gt","term":"id","value":5})"_json; + + ICEBERG_UNWRAP_OR_FAIL(auto task, FileScanTaskFromJson(json)); + EXPECT_EQ(task->start(), test_case.start); + EXPECT_EQ(task->length(), test_case.length); + EXPECT_EQ(task->is_split(), test_case.is_split); + ASSERT_NE(task->data_file(), nullptr); + EXPECT_EQ(*task->data_file(), DataFileForTask()); + ASSERT_EQ(task->delete_files().size(), 1U); + ASSERT_NE(task->delete_files()[0], nullptr); + EXPECT_EQ(*task->delete_files()[0], DeleteFileForTask()); + ASSERT_NE(task->residual_filter(), nullptr); + ICEBERG_UNWRAP_OR_FAIL(auto residual, ToJson(*task->residual_filter())); + EXPECT_EQ(residual, json["residual-filter"]); + + ICEBERG_UNWRAP_OR_FAIL(auto serialized, ToJson(*task, specs_, schema_)); + EXPECT_EQ(serialized, json); + ICEBERG_UNWRAP_OR_FAIL(auto reparsed, FileScanTaskFromJson(serialized)); + EXPECT_EQ(reparsed->start(), test_case.start); + EXPECT_EQ(reparsed->length(), test_case.length); + EXPECT_EQ(reparsed->is_split(), test_case.is_split); + } +} + +TEST_F(FileScanTaskJsonTest, RejectsInvalidRanges) { + struct Case { + int64_t file_size; + int64_t start; + int64_t length; + }; + constexpr auto kMax = std::numeric_limits::max(); + for (const auto& test_case : + {// Extends beyond the file. + Case{.file_size = 10, .start = 1, .length = 10}, + // Starts beyond the file. + Case{.file_size = 10, .start = 11, .length = 1}, + // Negative start. + Case{.file_size = 10, .start = -1, .length = 1}, + // Negative length. + Case{.file_size = 10, .start = 0, .length = -1}, + // Empty range in a nonempty file. + Case{.file_size = 10, .start = 0, .length = 0}, + // Empty range at EOF. + Case{.file_size = 10, .start = 10, .length = 0}, + // Nonempty range in an empty file. + Case{.file_size = 0, .start = 0, .length = 1}, + // Negative file size. + Case{.file_size = -1, .start = 0, .length = 0}, + // End would overflow int64_t. + Case{.file_size = kMax, .start = kMax - 1, .length = 2}}) { + SCOPED_TRACE(testing::Message() + << "file_size=" << test_case.file_size << ", start=" << test_case.start + << ", length=" << test_case.length); + auto json = JavaGolden(); + json["data-file"]["file-size-in-bytes"] = test_case.file_size; + json["start"] = test_case.start; + json["length"] = test_case.length; auto result = FileScanTaskFromJson(json); - if (test_case.supported) { - EXPECT_THAT(result, IsOk()); - } else { - EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); - EXPECT_THAT(result, HasErrorMessage("Split FileScanTask is not supported")); - } + EXPECT_THAT(result, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(result, HasErrorMessage("Invalid FileScanTask range")); } } +TEST_F(FileScanTaskJsonTest, RoundTripsEmptyWholeFile) { + auto json = JavaGolden(); + json["data-file"]["file-size-in-bytes"] = 0; + json["data-file"]["record-count"] = 0; + json["data-file"].erase("lower-bounds"); + json["data-file"].erase("upper-bounds"); + json.erase("delete-files"); + json["start"] = 0; + json["length"] = 0; + + ICEBERG_UNWRAP_OR_FAIL(auto task, FileScanTaskFromJson(json)); + EXPECT_EQ(task->start(), 0); + EXPECT_EQ(task->length(), 0); + EXPECT_FALSE(task->is_split()); + + ICEBERG_UNWRAP_OR_FAIL(auto serialized, ToJson(*task, specs_, schema_)); + EXPECT_EQ(serialized, json); +} + TEST_F(FileScanTaskJsonTest, RejectsFractionalIntegerFields) { for (std::string_view field : {"start", "length", "spec-id", "file-size-in-bytes", "record-count", diff --git a/src/iceberg/test/parquet_test.cc b/src/iceberg/test/parquet_test.cc index ce43962c6a..acc2f07ab4 100644 --- a/src/iceberg/test/parquet_test.cc +++ b/src/iceberg/test/parquet_test.cc @@ -37,7 +37,9 @@ #include #include #include +#include #include +#include #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_status_internal.h" @@ -45,6 +47,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" @@ -156,6 +159,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}; @@ -165,6 +193,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(); } @@ -249,7 +316,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, @@ -268,12 +335,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()), @@ -859,15 +944,24 @@ TEST_F(ParquetReaderTest, ReadRenamedLargeListColumnWithArrowSchema) { 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()), @@ -876,13 +970,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; @@ -906,6 +1013,134 @@ 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) { + file_io_ = arrow::ArrowFileSystemFileIO::MakeLocalFileIO(); + temp_parquet_file_ = GetResourcePath("split-without-row-group-offsets.parquet"); + ICEBERG_UNWRAP_OR_FAIL(auto input, + arrow::OpenArrowInputStream(file_io_, temp_parquet_file_)); + auto metadata = ::parquet::ReadMetaData(input); + ASSERT_EQ(metadata->num_row_groups(), 3); + for (int i = 0; i < metadata->num_row_groups(); ++i) { + ASSERT_EQ(metadata->RowGroup(i)->file_offset(), 0); + } + 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 must not span the skipped row group, or generated positions would be wrong. + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[1, 0, 100, 7]])")); + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[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, ReadMetadataColumns) { temp_parquet_file_ = "metadata_columns.parquet"; CreateSimpleParquetFile(); diff --git a/src/iceberg/test/resources/split-without-row-group-offsets.md b/src/iceberg/test/resources/split-without-row-group-offsets.md new file mode 100644 index 0000000000..c6ad010961 --- /dev/null +++ b/src/iceberg/test/resources/split-without-row-group-offsets.md @@ -0,0 +1,45 @@ + + +# Parquet split fixture with omitted row-group offsets + +`split-without-row-group-offsets.parquet` contains a required `int32` column +named `id`, with Iceberg field ID 1 and values 1 through 6. Its three row groups +contain two rows each, with dictionary pages before their first data pages. +The optional `RowGroup.file_offset` field is absent from every row group. + +The fixture was generated with PyArrow 12.0.1 using the following write options: + +```python +schema = pa.schema([ + pa.field("id", pa.int32(), nullable=False, metadata={b"PARQUET:field_id": b"1"}) +]) +table = pa.Table.from_arrays([pa.array([1, 2, 3, 4, 5, 6], type=pa.int32())], schema=schema) +pq.write_table(table, path, row_group_size=2, compression="NONE", + use_dictionary=True, version="2.6", store_schema=False) +``` + +The footer was then decoded with Thrift's compact protocol, field 5 +(`file_offset`) was removed from each entry of `FileMetaData.row_groups`, and +the footer was serialized again. The original data pages were retained and +the four-byte footer length was updated. A PyArrow read of the resulting file +was checked against the original table. + +This fixture tests that split selection uses the first column's page offsets +even when the optional row-group offset is unavailable. diff --git a/src/iceberg/test/resources/split-without-row-group-offsets.parquet b/src/iceberg/test/resources/split-without-row-group-offsets.parquet new file mode 100644 index 0000000000..fc84a7c3c9 Binary files /dev/null and b/src/iceberg/test/resources/split-without-row-group-offsets.parquet differ diff --git a/src/iceberg/test/rest_json_serde_test.cc b/src/iceberg/test/rest_json_serde_test.cc index 34ca4a45b3..119561daf7 100644 --- a/src/iceberg/test/rest_json_serde_test.cc +++ b/src/iceberg/test/rest_json_serde_test.cc @@ -17,6 +17,7 @@ * under the License. */ +#include #include #include #include @@ -2127,6 +2128,9 @@ TEST(FileScanTasksFromJsonTest, SingleTaskNoDeleteFiles) { const auto& task = result.value()[0]; ASSERT_NE(task->data_file(), nullptr); EXPECT_EQ(task->data_file()->file_path, "s3://bucket/data/file.parquet"); + EXPECT_EQ(task->start(), 0); + EXPECT_EQ(task->length(), 12345); + EXPECT_FALSE(task->is_split()); EXPECT_TRUE(task->delete_files().empty()); EXPECT_EQ(task->residual_filter(), nullptr); } @@ -2146,11 +2150,15 @@ TEST(FileScanTasksFromJsonTest, AcceptsWholeFileStartAndLength) { "length": 12345 }])"_json; - EXPECT_THAT(FileScanTasksFromJson(json, {}, UnpartitionedSpecs(), Schema({}, 0)), - IsOk()); + auto result = FileScanTasksFromJson(json, {}, UnpartitionedSpecs(), Schema({}, 0)); + ASSERT_THAT(result, IsOk()); + ASSERT_EQ(result.value().size(), 1U); + EXPECT_EQ(result.value()[0]->start(), 0); + EXPECT_EQ(result.value()[0]->length(), 12345); + EXPECT_FALSE(result.value()[0]->is_split()); } -TEST(FileScanTasksFromJsonTest, NullStartAndLengthUseWholeFileDefaults) { +TEST(FileScanTasksFromJsonTest, RejectsNullOrMissingRangeEndpoint) { auto json = R"([{ "data-file": { "content": "data", @@ -2161,12 +2169,25 @@ TEST(FileScanTasksFromJsonTest, NullStartAndLengthUseWholeFileDefaults) { "file-size-in-bytes": 12345, "record-count": 100 }, - "start": null, - "length": null + "start": 0, + "length": 12345 }])"_json; - EXPECT_THAT(FileScanTasksFromJson(json, {}, UnpartitionedSpecs(), Schema({}, 0)), - IsOk()); + for (std::string_view missing_field : {"start", "length"}) { + SCOPED_TRACE(missing_field); + auto incomplete_json = json; + incomplete_json[0].erase(std::string(missing_field)); + auto result = + FileScanTasksFromJson(incomplete_json, {}, UnpartitionedSpecs(), Schema({}, 0)); + EXPECT_THAT(result, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(result, HasErrorMessage("must be provided together")); + } + + json[0]["start"] = nullptr; + json[0]["length"] = nullptr; + auto null_result = FileScanTasksFromJson(json, {}, UnpartitionedSpecs(), Schema({}, 0)); + EXPECT_THAT(null_result, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(null_result, HasErrorMessage("Missing 'start'")); } TEST(FileScanTasksFromJsonTest, RejectsInvalidRangeIntegers) { @@ -2202,7 +2223,7 @@ TEST(FileScanTasksFromJsonTest, RejectsInvalidRangeIntegers) { } } -TEST(FileScanTasksFromJsonTest, RejectsSplitTaskWhenEitherRangeFieldDiffers) { +TEST(FileScanTasksFromJsonTest, AcceptsValidPartialRanges) { auto json = R"([{ "data-file": { "content": "data", @@ -2217,17 +2238,85 @@ TEST(FileScanTasksFromJsonTest, RejectsSplitTaskWhenEitherRangeFieldDiffers) { "length": 12345 }])"_json; - for (const auto& [field, value] : - {std::pair{"start", 100}, - std::pair{"length", 200}}) { - SCOPED_TRACE(field); + for (const auto& [start, length] : + {std::pair{0, 200}, {100, 200}, {12245, 100}}) { + SCOPED_TRACE(testing::Message() << "start=" << start << ", length=" << length); auto split_json = json; - split_json[0][field] = value; + split_json[0]["start"] = start; + split_json[0]["length"] = length; auto result = FileScanTasksFromJson(split_json, {}, UnpartitionedSpecs(), Schema({}, 0)); - EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); - EXPECT_THAT(result, HasErrorMessage("Split FileScanTask is not supported")); + ASSERT_THAT(result, IsOk()); + ASSERT_EQ(result.value().size(), 1U); + EXPECT_EQ(result.value()[0]->start(), start); + EXPECT_EQ(result.value()[0]->length(), length); + EXPECT_TRUE(result.value()[0]->is_split()); + } +} + +TEST(FileScanTasksFromJsonTest, RejectsInvalidRanges) { + auto json = R"([{ + "data-file": { + "content": "data", + "file-path": "s3://bucket/data/file.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": [], + "file-size-in-bytes": 12345, + "record-count": 100 + }, + "start": 0, + "length": 12345 + }])"_json; + + for (const auto& [start, length] : {std::pair{-1, 100}, + {0, -1}, + {0, 0}, + {12345, 0}, + {12345, 1}, + {12000, 500}}) { + SCOPED_TRACE(testing::Message() << "start=" << start << ", length=" << length); + auto invalid_json = json; + invalid_json[0]["start"] = start; + invalid_json[0]["length"] = length; + auto result = + FileScanTasksFromJson(invalid_json, {}, UnpartitionedSpecs(), Schema({}, 0)); + EXPECT_THAT(result, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(result, HasErrorMessage("Invalid FileScanTask range")); } + + auto overflow_json = json; + overflow_json[0]["data-file"]["file-size-in-bytes"] = + std::numeric_limits::max(); + overflow_json[0]["start"] = std::numeric_limits::max() - 1; + overflow_json[0]["length"] = 2; + auto overflow_result = + FileScanTasksFromJson(overflow_json, {}, UnpartitionedSpecs(), Schema({}, 0)); + EXPECT_THAT(overflow_result, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(overflow_result, HasErrorMessage("Invalid FileScanTask range")); +} + +TEST(FileScanTasksFromJsonTest, AcceptsEmptyWholeFile) { + auto json = R"([{ + "data-file": { + "content": "data", + "file-path": "s3://bucket/data/empty.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": [], + "file-size-in-bytes": 0, + "record-count": 0 + }, + "start": 0, + "length": 0 + }])"_json; + + auto result = FileScanTasksFromJson(json, {}, UnpartitionedSpecs(), Schema({}, 0)); + ASSERT_THAT(result, IsOk()); + ASSERT_EQ(result.value().size(), 1U); + EXPECT_EQ(result.value()[0]->start(), 0); + EXPECT_EQ(result.value()[0]->length(), 0); + EXPECT_FALSE(result.value()[0]->is_split()); } TEST(FileScanTasksFromJsonTest, RowLineageSequence) { @@ -2400,6 +2489,80 @@ TEST(FetchScanTasksResponseRoundtripTest, WithFileScanTasksAndDeleteFiles) { EXPECT_EQ(*result, *result2); } +TEST(FetchScanTasksResponseRoundtripTest, PreservesDistinctSplitRanges) { + auto json = nlohmann::json::parse(R"({ + "file-scan-tasks": [ + { + "data-file": { + "content": "data", + "file-path": "s3://bucket/data/file.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": [], + "file-size-in-bytes": 12345, + "record-count": 100 + }, + "start": 0, + "length": 5000 + }, + { + "data-file": { + "content": "data", + "file-path": "s3://bucket/data/file.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": [], + "file-size-in-bytes": 12345, + "record-count": 100 + }, + "start": 5000, + "length": 7345 + } + ] + })"); + + auto result = FetchScanTasksResponseFromJson(json, UnpartitionedSpecs(), EmptySchema()); + ASSERT_THAT(result, IsOk()); + ASSERT_TRUE(result.value().file_scan_tasks.has_value()); + ASSERT_EQ(result.value().file_scan_tasks->size(), 2U); + EXPECT_EQ(result.value().file_scan_tasks->at(0)->start(), 0); + EXPECT_EQ(result.value().file_scan_tasks->at(0)->length(), 5000); + EXPECT_EQ(result.value().file_scan_tasks->at(1)->start(), 5000); + EXPECT_EQ(result.value().file_scan_tasks->at(1)->length(), 7345); + + ICEBERG_UNWRAP_OR_FAIL(auto roundtrip_json, + ToJson(result.value(), UnpartitionedSpecs(), EmptySchema())); + EXPECT_EQ(roundtrip_json["file-scan-tasks"][0]["start"], 0); + EXPECT_EQ(roundtrip_json["file-scan-tasks"][0]["length"], 5000); + EXPECT_EQ(roundtrip_json["file-scan-tasks"][1]["start"], 5000); + EXPECT_EQ(roundtrip_json["file-scan-tasks"][1]["length"], 7345); + + auto result2 = + FetchScanTasksResponseFromJson(roundtrip_json, UnpartitionedSpecs(), EmptySchema()); + ASSERT_THAT(result2, IsOk()); + EXPECT_EQ(result.value(), result2.value()); +} + +TEST(FetchScanTasksResponseRoundtripTest, RejectsSplitOutsideChangedDataFile) { + auto data_file = std::make_shared(); + data_file->content = DataFile::Content::kData; + data_file->file_path = "s3://bucket/data/file.parquet"; + data_file->file_format = FileFormatType::kParquet; + data_file->partition_spec_id = PartitionSpec::kInitialSpecId; + data_file->partition = PartitionValues{}; + data_file->file_size_in_bytes = 12345; + data_file->record_count = 100; + ICEBERG_UNWRAP_OR_FAIL(auto split, FileScanTask::MakeSplit(data_file, 5000, 7345)); + + FetchScanTasksResponse response; + response.file_scan_tasks = std::vector>{split}; + data_file->file_size_in_bytes = 12000; + + auto result = ToJson(response, UnpartitionedSpecs(), EmptySchema()); + EXPECT_THAT(result, IsError(ErrorKind::kValidationFailed)); + EXPECT_THAT(result, HasErrorMessage("Invalid FileScanTask range")); +} + TEST(FetchScanTasksResponseRoundtripTest, PreservesResidualFilter) { auto json = nlohmann::json::parse(R"({ "file-scan-tasks": [