diff --git a/src/iceberg/arrow/metadata_column_util.cc b/src/iceberg/arrow/metadata_column_util.cc index 2176d897a..2ceb3933e 100644 --- a/src/iceberg/arrow/metadata_column_util.cc +++ b/src/iceberg/arrow/metadata_column_util.cc @@ -24,9 +24,18 @@ #include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow/metadata_column_util_internal.h" +#include "iceberg/metadata_columns.h" namespace iceberg::arrow { +bool HasRowLineageValue(int32_t field_id, const MetadataColumnContext& metadata_context) { + if (!metadata_context.first_row_id.has_value()) { + return false; + } + return field_id != MetadataColumns::kLastUpdatedSequenceNumberColumnId || + metadata_context.data_sequence_number.has_value(); +} + Result> MakeFilePathArray(const std::string& file_path, int64_t num_rows, ::arrow::MemoryPool* pool) { @@ -53,4 +62,41 @@ Result> MakeRowPositionArray(int64_t start_posit return array; } +Result> MakeRowIdArray( + std::optional first_row_id, int64_t start_position, int64_t num_rows, + ::arrow::MemoryPool* pool) { + ::arrow::Int64Builder builder(pool); + ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(num_rows)); + if (!first_row_id.has_value()) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.AppendNulls(num_rows)); + } else { + for (int64_t row_index = 0; row_index < num_rows; ++row_index) { + ICEBERG_ARROW_RETURN_NOT_OK( + builder.Append(first_row_id.value() + start_position + row_index)); + } + } + + std::shared_ptr<::arrow::Array> array; + ICEBERG_ARROW_RETURN_NOT_OK(builder.Finish(&array)); + return array; +} + +Result> MakeLastUpdatedSequenceNumberArray( + std::optional first_row_id, std::optional data_sequence_number, + int64_t num_rows, ::arrow::MemoryPool* pool) { + ::arrow::Int64Builder builder(pool); + ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(num_rows)); + if (!first_row_id.has_value() || !data_sequence_number.has_value()) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.AppendNulls(num_rows)); + } else { + for (int64_t row_index = 0; row_index < num_rows; ++row_index) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(data_sequence_number.value())); + } + } + + std::shared_ptr<::arrow::Array> array; + ICEBERG_ARROW_RETURN_NOT_OK(builder.Finish(&array)); + return array; +} + } // namespace iceberg::arrow diff --git a/src/iceberg/arrow/metadata_column_util_internal.h b/src/iceberg/arrow/metadata_column_util_internal.h index ca21700ac..49d3f6eac 100644 --- a/src/iceberg/arrow/metadata_column_util_internal.h +++ b/src/iceberg/arrow/metadata_column_util_internal.h @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -36,8 +37,13 @@ namespace iceberg::arrow { struct MetadataColumnContext { std::string file_path; // The file path to populate _file column. int64_t next_file_pos = 0; // The next row position for populating _pos column. + std::optional first_row_id; // First row ID used to populate _row_id. + std::optional data_sequence_number; // Data sequence number for lineage. }; +/// \brief Check if row lineage metadata can be inherited for the column. +bool HasRowLineageValue(int32_t field_id, const MetadataColumnContext& metadata_context); + /// \brief Create a constant string array for _file column. /// /// Creates an Arrow StringArray where all values are the same file path string. @@ -62,4 +68,26 @@ Result> MakeRowPositionArray(int64_t start_posit int64_t num_rows, ::arrow::MemoryPool* pool); +/// \brief Create a row ID array for _row_id column. +/// +/// \param first_row_id First row ID of the data file. +/// \param start_position Starting row position (inclusive). +/// \param num_rows Number of rows in the batch. +/// \param pool Arrow memory pool. +/// \return Arrow Int64Array for _row_id. +Result> MakeRowIdArray( + std::optional first_row_id, int64_t start_position, int64_t num_rows, + ::arrow::MemoryPool* pool); + +/// \brief Create a sequence number array for _last_updated_sequence_number column. +/// +/// \param first_row_id First row ID of the data file. +/// \param data_sequence_number Data sequence number of the data file. +/// \param num_rows Number of rows in the batch. +/// \param pool Arrow memory pool. +/// \return Arrow Int64Array for _last_updated_sequence_number. +Result> MakeLastUpdatedSequenceNumberArray( + std::optional first_row_id, std::optional data_sequence_number, + int64_t num_rows, ::arrow::MemoryPool* pool); + } // namespace iceberg::arrow diff --git a/src/iceberg/avro/avro_data_util.cc b/src/iceberg/avro/avro_data_util.cc index b44e4b110..756f9be44 100644 --- a/src/iceberg/avro/avro_data_util.cc +++ b/src/iceberg/avro/avro_data_util.cc @@ -54,6 +54,22 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node, const arrow::MetadataColumnContext& metadata_context, ::arrow::ArrayBuilder* array_builder); +Status AppendRowLineageValue(int32_t field_id, + const arrow::MetadataColumnContext& metadata_context, + ::arrow::ArrayBuilder* array_builder) { + auto* int_builder = internal::checked_cast<::arrow::Int64Builder*>(array_builder); + if (!arrow::HasRowLineageValue(field_id, metadata_context)) { + ICEBERG_ARROW_RETURN_NOT_OK(int_builder->AppendNull()); + } else if (field_id == MetadataColumns::kRowIdColumnId) { + ICEBERG_ARROW_RETURN_NOT_OK(int_builder->Append( + metadata_context.first_row_id.value() + metadata_context.next_file_pos)); + } else { + ICEBERG_ARROW_RETURN_NOT_OK( + int_builder->Append(metadata_context.data_sequence_number.value())); + } + return {}; +} + /// \brief Append Avro record data to Arrow struct builder. Status AppendStructToBuilder(const ::avro::NodePtr& avro_node, const ::avro::GenericDatum& avro_datum, @@ -97,6 +113,9 @@ Status AppendStructToBuilder(const ::avro::NodePtr& avro_node, } else if (field_id == MetadataColumns::kFilePositionColumnId) { auto int_builder = internal::checked_cast<::arrow::Int64Builder*>(field_builder); ICEBERG_ARROW_RETURN_NOT_OK(int_builder->Append(metadata_context.next_file_pos)); + } else if (MetadataColumns::IsRowLineageColumn(field_id)) { + ICEBERG_RETURN_UNEXPECTED( + AppendRowLineageValue(field_id, metadata_context, field_builder)); } else { return NotSupported("Unsupported metadata column field id: {}", field_id); } @@ -463,9 +482,21 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node, return {}; } + const bool is_row_lineage = + MetadataColumns::IsRowLineageColumn(projected_field.field_id()); + if (is_row_lineage && + !arrow::HasRowLineageValue(projected_field.field_id(), metadata_context)) { + ICEBERG_ARROW_RETURN_NOT_OK(array_builder->AppendNull()); + return {}; + } + if (avro_node->type() == ::avro::AVRO_UNION) { size_t branch = avro_datum.unionBranch(); if (avro_node->leafAt(branch)->type() == ::avro::AVRO_NULL) { + if (is_row_lineage) { + return AppendRowLineageValue(projected_field.field_id(), metadata_context, + array_builder); + } ICEBERG_ARROW_RETURN_NOT_OK(array_builder->AppendNull()); return {}; } else { diff --git a/src/iceberg/avro/avro_direct_decoder.cc b/src/iceberg/avro/avro_direct_decoder.cc index 583b9ac5c..1db2f3699 100644 --- a/src/iceberg/avro/avro_direct_decoder.cc +++ b/src/iceberg/avro/avro_direct_decoder.cc @@ -144,6 +144,22 @@ Status SkipAvroValue(const ::avro::NodePtr& avro_node, ::avro::Decoder& decoder) } } +Status AppendRowLineageValue(int32_t field_id, + const arrow::MetadataColumnContext& metadata_context, + ::arrow::ArrayBuilder* array_builder) { + auto* int_builder = internal::checked_cast<::arrow::Int64Builder*>(array_builder); + if (!arrow::HasRowLineageValue(field_id, metadata_context)) { + ICEBERG_ARROW_RETURN_NOT_OK(int_builder->AppendNull()); + } else if (field_id == MetadataColumns::kRowIdColumnId) { + ICEBERG_ARROW_RETURN_NOT_OK(int_builder->Append( + metadata_context.first_row_id.value() + metadata_context.next_file_pos)); + } else { + ICEBERG_ARROW_RETURN_NOT_OK( + int_builder->Append(metadata_context.data_sequence_number.value())); + } + return {}; +} + /// \brief Decode Avro record directly to Arrow struct builder. Status DecodeStructToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& decoder, const std::span& projections, @@ -218,6 +234,9 @@ Status DecodeStructToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& } else if (field_id == MetadataColumns::kFilePositionColumnId) { auto int_builder = internal::checked_cast<::arrow::Int64Builder*>(field_builder); ICEBERG_ARROW_RETURN_NOT_OK(int_builder->Append(metadata_context.next_file_pos)); + } else if (MetadataColumns::IsRowLineageColumn(field_id)) { + ICEBERG_RETURN_UNEXPECTED( + AppendRowLineageValue(field_id, metadata_context, field_builder)); } else { return NotSupported("Unsupported metadata column field id: {}", field_id); } @@ -594,6 +613,15 @@ Status DecodeFieldToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& d return {}; } + const bool is_row_lineage = + MetadataColumns::IsRowLineageColumn(projected_field.field_id()); + if (is_row_lineage && + !arrow::HasRowLineageValue(projected_field.field_id(), metadata_context)) { + ICEBERG_RETURN_UNEXPECTED(SkipAvroValue(avro_node, decoder)); + ICEBERG_ARROW_RETURN_NOT_OK(array_builder->AppendNull()); + return {}; + } + if (avro_node->type() == ::avro::AVRO_UNION) { const size_t branch_index = decoder.decodeUnionIndex(); @@ -606,6 +634,10 @@ Status DecodeFieldToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& d const auto& branch_node = avro_node->leafAt(branch_index); if (branch_node->type() == ::avro::AVRO_NULL) { + if (is_row_lineage) { + return AppendRowLineageValue(projected_field.field_id(), metadata_context, + array_builder); + } ICEBERG_ARROW_RETURN_NOT_OK(array_builder->AppendNull()); return {}; } else { diff --git a/src/iceberg/avro/avro_reader.cc b/src/iceberg/avro/avro_reader.cc index 95218ccbd..b0977bcdb 100644 --- a/src/iceberg/avro/avro_reader.cc +++ b/src/iceberg/avro/avro_reader.cc @@ -68,6 +68,15 @@ bool HasRowPositionColumn(const Schema& schema) { return false; } +bool HasRowIdColumn(const Schema& schema) { + for (const auto& field : schema.fields()) { + if (field.field_id() == MetadataColumns::kRowIdColumnId) { + return true; + } + } + return false; +} + // Abstract base class for Avro read backends. class AvroReadBackend { public: @@ -332,9 +341,17 @@ class AvroReader::Impl { return NotSupported( "Reading '_pos' metadata column with split is not supported for Avro files."); } + if (options.split->offset != 0 && HasRowIdColumn(*read_schema_)) { + return NotSupported( + "Reading '_row_id' metadata column with split is not supported for Avro " + "files."); + } } - metadata_context_ = {.file_path = options.path, .next_file_pos = 0}; + metadata_context_ = {.file_path = options.path, + .next_file_pos = 0, + .first_row_id = options.first_row_id, + .data_sequence_number = options.data_sequence_number}; return {}; } diff --git a/src/iceberg/data/file_scan_task_reader.cc b/src/iceberg/data/file_scan_task_reader.cc index 7076486ac..0c41f25ed 100644 --- a/src/iceberg/data/file_scan_task_reader.cc +++ b/src/iceberg/data/file_scan_task_reader.cc @@ -43,7 +43,9 @@ ReaderOptions MakeReaderOptions(const DataFile& data_file, std::shared_ptr projection, std::shared_ptr filter, std::shared_ptr name_mapping, - ReaderProperties properties) { + ReaderProperties properties, + std::optional first_row_id, + std::optional data_sequence_number) { return ReaderOptions{ .path = data_file.file_path, .length = static_cast(data_file.file_size_in_bytes), @@ -51,6 +53,8 @@ ReaderOptions MakeReaderOptions(const DataFile& data_file, std::shared_ptrfile_size_in_bytes); if (task.delete_files().empty()) { - auto options = - MakeReaderOptions(*data_file, io_, projected_schema_, task.residual_filter(), - name_mapping_, properties_); + auto options = MakeReaderOptions( + *data_file, 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)); return MakeArrowArrayStream(std::move(reader)); @@ -187,8 +191,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_); + auto options = MakeReaderOptions( + *data_file, 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/file_reader.h b/src/iceberg/file_reader.h index c76c10093..cefc688c0 100644 --- a/src/iceberg/file_reader.h +++ b/src/iceberg/file_reader.h @@ -22,6 +22,7 @@ /// \file iceberg/file_reader.h /// Reader interface for file formats like Parquet, Avro and ORC. +#include #include #include #include @@ -105,6 +106,10 @@ struct ICEBERG_EXPORT ReaderOptions { /// \brief Name mapping for schema evolution compatibility. Used when reading files /// that may have different field names than the current schema. std::shared_ptr name_mapping; + /// \brief First row ID inherited by the data file, if assigned. + std::optional first_row_id; + /// \brief Data sequence number inherited by the manifest entry, if assigned. + std::optional data_sequence_number; /// \brief Format-specific or implementation-specific properties. ReaderProperties properties; }; diff --git a/src/iceberg/inheritable_metadata.cc b/src/iceberg/inheritable_metadata.cc index 7ff2ddbcb..4e4592d7b 100644 --- a/src/iceberg/inheritable_metadata.cc +++ b/src/iceberg/inheritable_metadata.cc @@ -64,6 +64,7 @@ Status BaseInheritableMetadata::Apply(ManifestEntry& entry) { if (entry.data_file) { entry.data_file->partition_spec_id = spec_id_; + entry.data_file->data_sequence_number = entry.sequence_number; } else { return InvalidManifest("Manifest entry has no data file"); } diff --git a/src/iceberg/manifest/manifest_entry.h b/src/iceberg/manifest/manifest_entry.h index e2dd5ea61..d3dc2f7a2 100644 --- a/src/iceberg/manifest/manifest_entry.h +++ b/src/iceberg/manifest/manifest_entry.h @@ -174,11 +174,18 @@ struct ICEBERG_EXPORT DataFile { /// present std::optional content_size_in_bytes; + /// \note Fields defined below are inherited and not persisted to the DataFile struct. + /// \brief Partition spec id for this data file. - /// \note This field is for internal use only and will not be persisted to manifest - /// entry. + /// \note This inherited field is for internal use only and will not be persisted to + /// the DataFile struct. std::optional partition_spec_id; + /// \brief Data sequence number for this data file. + /// \note This inherited field is for internal use only and will not be persisted to + /// the DataFile struct. + std::optional data_sequence_number; + static constexpr int32_t kContentFieldId = 134; inline static const SchemaField kContent = SchemaField::MakeOptional( kContentFieldId, "content", int32(), diff --git a/src/iceberg/metadata_columns.cc b/src/iceberg/metadata_columns.cc index 0c2521133..e00bc1661 100644 --- a/src/iceberg/metadata_columns.cc +++ b/src/iceberg/metadata_columns.cc @@ -66,6 +66,10 @@ bool MetadataColumns::IsMetadataColumn(int32_t id) { return GetMetadataFieldIdSet().find(id) != GetMetadataFieldIdSet().end(); } +bool MetadataColumns::IsRowLineageColumn(int32_t id) { + return id == kRowIdColumnId || id == kLastUpdatedSequenceNumberColumnId; +} + Result MetadataColumns::MetadataColumn(std::string_view name) { const auto& metadata_column_map = GetMetadataColumnMap(); const auto it = metadata_column_map.find(name); diff --git a/src/iceberg/metadata_columns.h b/src/iceberg/metadata_columns.h index b390a50e8..dedc725a2 100644 --- a/src/iceberg/metadata_columns.h +++ b/src/iceberg/metadata_columns.h @@ -113,6 +113,9 @@ struct ICEBERG_EXPORT MetadataColumns { /// \brief Check if a column ID is a metadata column. static bool IsMetadataColumn(int32_t id); + /// \brief Check if a column ID is a row lineage column. + static bool IsRowLineageColumn(int32_t id); + /// \brief Get a metadata column by name. /// /// \param name The name of the metadata column. diff --git a/src/iceberg/parquet/parquet_data_util.cc b/src/iceberg/parquet/parquet_data_util.cc index b6473edcc..54733598e 100644 --- a/src/iceberg/parquet/parquet_data_util.cc +++ b/src/iceberg/parquet/parquet_data_util.cc @@ -70,6 +70,53 @@ Result> ProjectPrimitiveArray( return cast_result.make_array(); } +Result> MakeInheritedRowLineageArray( + int32_t field_id, int64_t num_rows, + const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) { + if (field_id == MetadataColumns::kRowIdColumnId) { + return arrow::MakeRowIdArray(metadata_context.first_row_id, + metadata_context.next_file_pos, num_rows, pool); + } + if (field_id == MetadataColumns::kLastUpdatedSequenceNumberColumnId) { + return arrow::MakeLastUpdatedSequenceNumberArray( + metadata_context.first_row_id, metadata_context.data_sequence_number, num_rows, + pool); + } + return NotSupported("Unsupported row lineage field id: {}", field_id); +} + +Result> ProjectRowLineageArray( + const std::shared_ptr<::arrow::Array>& physical_array, int32_t field_id, + const arrow::MetadataColumnContext& metadata_context, ::arrow::MemoryPool* pool) { + if (!arrow::HasRowLineageValue(field_id, metadata_context)) { + return MakeInheritedRowLineageArray(field_id, physical_array->length(), + metadata_context, pool); + } + if (physical_array->null_count() == 0) { + return physical_array; + } + + auto values = internal::checked_pointer_cast<::arrow::Int64Array>(physical_array); + ::arrow::Int64Builder builder(pool); + ICEBERG_ARROW_RETURN_NOT_OK(builder.Reserve(values->length())); + for (int64_t row_index = 0; row_index < values->length(); ++row_index) { + if (!values->IsNull(row_index)) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(values->Value(row_index))); + } else if (field_id == MetadataColumns::kRowIdColumnId) { + ICEBERG_ARROW_RETURN_NOT_OK(builder.Append(metadata_context.first_row_id.value() + + metadata_context.next_file_pos + + row_index)); + } else { + ICEBERG_ARROW_RETURN_NOT_OK( + builder.Append(metadata_context.data_sequence_number.value())); + } + } + + std::shared_ptr<::arrow::Array> output; + ICEBERG_ARROW_RETURN_NOT_OK(builder.Finish(&output)); + return output; +} + Result> ProjectStructArray( const std::shared_ptr<::arrow::StructArray>& struct_array, const std::shared_ptr<::arrow::StructType>& output_struct_type, @@ -111,6 +158,11 @@ Result> ProjectStructArray( projected_array, ProjectNestedArray(parquet_array, output_arrow_type, nested_type, field_projection.children, metadata_context, pool)); + } else if (MetadataColumns::IsRowLineageColumn(projected_field.field_id())) { + ICEBERG_ASSIGN_OR_RAISE( + projected_array, + ProjectRowLineageArray(parquet_array, projected_field.field_id(), + metadata_context, pool)); } else { ICEBERG_ASSIGN_OR_RAISE( projected_array, @@ -135,6 +187,10 @@ Result> ProjectStructArray( ICEBERG_ASSIGN_OR_RAISE( projected_array, arrow::MakeRowPositionArray(metadata_context.next_file_pos, struct_array->length(), pool)); + } else if (MetadataColumns::IsRowLineageColumn(field_id)) { + ICEBERG_ASSIGN_OR_RAISE(projected_array, MakeInheritedRowLineageArray( + field_id, struct_array->length(), + metadata_context, pool)); } else { return NotSupported("Unsupported metadata field id: {}", field_id); } diff --git a/src/iceberg/parquet/parquet_reader.cc b/src/iceberg/parquet/parquet_reader.cc index 775644a94..f0df07ac2 100644 --- a/src/iceberg/parquet/parquet_reader.cc +++ b/src/iceberg/parquet/parquet_reader.cc @@ -129,7 +129,10 @@ class ParquetReader::Impl { // Project read schema onto the Parquet file schema ICEBERG_ASSIGN_OR_RAISE(projection_, BuildProjection(reader_.get(), *read_schema_)); - metadata_context_ = {.file_path = options.path, .next_file_pos = 0}; + metadata_context_ = {.file_path = options.path, + .next_file_pos = 0, + .first_row_id = options.first_row_id, + .data_sequence_number = options.data_sequence_number}; return {}; } diff --git a/src/iceberg/table_scan.cc b/src/iceberg/table_scan.cc index d0a214d07..5f7d03648 100644 --- a/src/iceberg/table_scan.cc +++ b/src/iceberg/table_scan.cc @@ -182,6 +182,18 @@ int64_t FileScanTask::estimated_row_count() const { return data_file_->record_co // ChangelogScanTask implementation +ChangelogScanTask::ChangelogScanTask(int32_t change_ordinal, int64_t commit_snapshot_id, + std::shared_ptr data_file, + std::vector> delete_files, + std::shared_ptr residual_filter) + : change_ordinal_(change_ordinal), + commit_snapshot_id_(commit_snapshot_id), + data_file_(std::move(data_file)), + delete_files_(std::move(delete_files)), + residual_filter_(std::move(residual_filter)) { + ICEBERG_DCHECK(data_file_ != nullptr, "Data file cannot be null for ChangelogScanTask"); +} + int64_t ChangelogScanTask::size_bytes() const { int64_t total_size = data_file_->file_size_in_bytes; for (const auto& delete_file : delete_files_) { @@ -790,7 +802,6 @@ IncrementalChangelogScan::PlanFiles(std::optional from_snapshot_id_excl ICEBERG_ASSIGN_OR_RAISE(auto residual, ctx.residuals->ResidualFor(entry.data_file->partition)); - switch (entry.status) { case ManifestStatus::kAdded: tasks.push_back(std::make_shared( diff --git a/src/iceberg/table_scan.h b/src/iceberg/table_scan.h index e55330f30..f47ba8a38 100644 --- a/src/iceberg/table_scan.h +++ b/src/iceberg/table_scan.h @@ -113,12 +113,7 @@ class ICEBERG_EXPORT ChangelogScanTask : public ScanTask { ChangelogScanTask(int32_t change_ordinal, int64_t commit_snapshot_id, std::shared_ptr data_file, std::vector> delete_files = {}, - std::shared_ptr residual_filter = nullptr) - : change_ordinal_(change_ordinal), - commit_snapshot_id_(commit_snapshot_id), - data_file_(std::move(data_file)), - delete_files_(std::move(delete_files)), - residual_filter_(std::move(residual_filter)) {} + std::shared_ptr residual_filter = nullptr); Kind kind() const override { return Kind::kChangelogScanTask; } diff --git a/src/iceberg/test/avro_test.cc b/src/iceberg/test/avro_test.cc index 3ae1696d0..808aac8fb 100644 --- a/src/iceberg/test/avro_test.cc +++ b/src/iceberg/test/avro_test.cc @@ -162,6 +162,29 @@ class AvroReaderTest : public TempFileTestBase { ASSERT_THAT(writer->Close(), IsOk()); } + std::shared_ptr RowLineageSchema() { + return std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + MetadataColumns::kRowId, + MetadataColumns::kLastUpdatedSequenceNumber, + }); + } + + Result> OpenRowLineageReader( + const std::shared_ptr& schema, + std::optional first_row_id = std::nullopt, + std::optional data_sequence_number = std::nullopt) { + ReaderProperties properties; + properties.Set(ReaderProperties::kAvroSkipDatum, skip_datum_); + return ReaderFactoryRegistry::Open(FileFormatType::kAvro, + {.path = temp_avro_file_, + .io = file_io_, + .projection = schema, + .first_row_id = first_row_id, + .data_sequence_number = data_sequence_number, + .properties = std::move(properties)}); + } + void VerifyNextBatch(Reader& reader, std::string_view expected_json) { // Boilerplate to get Arrow schema auto schema_result = reader.Schema(); @@ -694,6 +717,30 @@ TEST_P(AvroReaderParameterizedTest, ReadMetadataOnlyProjection) { VerifyNextBatch(*reader, expected_string); } +TEST_P(AvroReaderParameterizedTest, ReadRowLineage) { + temp_avro_file_ = "avro_row_lineage.avro"; + CreateSimpleAvroFile(); + + ICEBERG_UNWRAP_OR_FAIL(auto reader, OpenRowLineageReader(RowLineageSchema(), 100L, 7L)); + + VerifyNextBatch(*reader, R"([[1, 100, 7], [2, 101, 7], [3, 102, 7]])"); +} + +TEST_P(AvroReaderParameterizedTest, ReadPartialLineage) { + temp_avro_file_ = "avro_partial_lineage.avro"; + CreateSimpleAvroFile(); + + auto schema = RowLineageSchema(); + ICEBERG_UNWRAP_OR_FAIL(auto reader, OpenRowLineageReader(schema)); + + VerifyNextBatch(*reader, R"([[1, null, null], [2, null, null], [3, null, null]])"); + + ICEBERG_UNWRAP_OR_FAIL(auto reader_with_row_id, OpenRowLineageReader(schema, 100L)); + + VerifyNextBatch(*reader_with_row_id, + R"([[1, 100, null], [2, 101, null], [3, 102, null]])"); +} + TEST_P(AvroReaderParameterizedTest, SplitWithRowPositionNotSupported) { CreateSimpleAvroFile(); @@ -714,6 +761,29 @@ TEST_P(AvroReaderParameterizedTest, SplitWithRowPositionNotSupported) { HasErrorMessage("'_pos' metadata column with split is not supported")); } +TEST_P(AvroReaderParameterizedTest, SplitWithRowIdNotSupported) { + CreateSimpleAvroFile(); + + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + MetadataColumns::kRowId, + }); + + ReaderProperties properties; + properties.Set(ReaderProperties::kAvroSkipDatum, skip_datum_); + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kAvro, {.path = temp_avro_file_, + .split = Split{.offset = 100, .length = 200}, + .io = file_io_, + .projection = schema, + .first_row_id = 100L, + .properties = std::move(properties)}); + + ASSERT_THAT(reader_result, IsError(ErrorKind::kNotSupported)); + EXPECT_THAT(reader_result, + HasErrorMessage("'_row_id' metadata column with split is not supported")); +} + INSTANTIATE_TEST_SUITE_P(DirectDecoderModes, AvroReaderParameterizedTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { diff --git a/src/iceberg/test/data_writer_test.cc b/src/iceberg/test/data_writer_test.cc index 14b7bf628..ca1986429 100644 --- a/src/iceberg/test/data_writer_test.cc +++ b/src/iceberg/test/data_writer_test.cc @@ -30,6 +30,7 @@ #include "iceberg/data/equality_delete_writer.h" #include "iceberg/data/position_delete_writer.h" #include "iceberg/file_format.h" +#include "iceberg/file_reader.h" #include "iceberg/manifest/manifest_entry.h" #include "iceberg/metadata_columns.h" #include "iceberg/parquet/parquet_register.h" @@ -76,17 +77,34 @@ class DataWriterTest : public ::testing::Test { }; } - std::shared_ptr<::arrow::Array> CreateTestData() { + std::shared_ptr<::arrow::Array> CreateArray(const Schema& schema, + std::string_view json) { ArrowSchema arrow_c_schema; - ICEBERG_THROW_NOT_OK(ToArrowSchema(*schema_, &arrow_c_schema)); + ICEBERG_THROW_NOT_OK(ToArrowSchema(schema, &arrow_c_schema)); auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); - return ::arrow::json::ArrayFromJSONString( - ::arrow::struct_(arrow_schema->fields()), - R"([[1, "Alice"], [2, "Bob"], [3, "Charlie"]])") + return ::arrow::json::ArrayFromJSONString(::arrow::struct_(arrow_schema->fields()), + std::string(json)) .ValueOrDie(); } + std::shared_ptr<::arrow::Array> CreateTestData() { + return CreateArray(*schema_, R"([[1, "Alice"], [2, "Bob"], [3, "Charlie"]])"); + } + + std::shared_ptr RowLineageSchema() { + return std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), MetadataColumns::kRowId, + MetadataColumns::kLastUpdatedSequenceNumber}); + } + + std::unordered_map FormatProperties(FileFormatType format) { + if (format == FileFormatType::kParquet) { + return {{"write.parquet.compression-codec", "uncompressed"}}; + } + return {}; + } + void WriteTestDataToWriter(DataWriter* writer) { auto test_data = CreateTestData(); ArrowArray arrow_array; @@ -94,6 +112,20 @@ class DataWriterTest : public ::testing::Test { ASSERT_THAT(writer->Write(&arrow_array), IsOk()); } + void VerifyNextBatch(Reader& reader, const Schema& schema, + std::string_view expected_json) { + auto data = reader.Next(); + ASSERT_THAT(data, IsOk()) << "Reader.Next() failed: " << data.error().message; + ASSERT_TRUE(data.value().has_value()) << "Reader.Next() returned no data"; + + auto expected_array = CreateArray(schema, expected_json); + auto actual_array = + ::arrow::ImportArray(&data.value().value(), expected_array->type()).ValueOrDie(); + ASSERT_TRUE(actual_array->Equals(expected_array)) + << "Expected: " << expected_array->ToString() + << "\nActual: " << actual_array->ToString(); + } + std::shared_ptr file_io_; std::shared_ptr schema_; std::shared_ptr partition_spec_; @@ -112,12 +144,7 @@ TEST_P(DataWriterFormatTest, CreateWithFormat) { .partition = PartitionValues{}, .format = format, .io = file_io_, - .properties = - format == FileFormatType::kParquet - ? std::unordered_map{{"write.parquet.compression-codec", - "uncompressed"}} - : std::unordered_map{}, + .properties = FormatProperties(format), }; auto writer_result = DataWriter::Make(options); @@ -126,6 +153,41 @@ TEST_P(DataWriterFormatTest, CreateWithFormat) { ASSERT_NE(writer, nullptr); } +TEST_P(DataWriterFormatTest, WriteRowLineage) { + auto [format, path] = GetParam(); + auto schema = RowLineageSchema(); + DataWriterOptions options{ + .path = path, + .schema = schema, + .spec = partition_spec_, + .partition = PartitionValues{}, + .format = format, + .io = file_io_, + .properties = FormatProperties(format), + }; + + ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options)); + + auto array = CreateArray(*schema, R"([[1, null, null], + [2, 777, 8], + [3, null, null]])"); + ArrowArray arrow_array; + ASSERT_TRUE(::arrow::ExportArray(*array, &arrow_array).ok()); + ASSERT_THAT(writer->Write(&arrow_array), IsOk()); + ASSERT_THAT(writer->Close(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL( + auto reader, ReaderFactoryRegistry::Open(format, {.path = path, + .io = file_io_, + .projection = schema, + .first_row_id = 100L, + .data_sequence_number = 7L})); + + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, *schema, R"([[1, 100, 7], + [2, 777, 8], + [3, 102, 7]])")); +} + INSTANTIATE_TEST_SUITE_P( FormatTypes, DataWriterFormatTest, ::testing::Values(std::make_pair(FileFormatType::kParquet, "test_data.parquet"), diff --git a/src/iceberg/test/file_scan_task_reader_test.cc b/src/iceberg/test/file_scan_task_reader_test.cc index 1630a108a..024d72f3c 100644 --- a/src/iceberg/test/file_scan_task_reader_test.cc +++ b/src/iceberg/test/file_scan_task_reader_test.cc @@ -43,6 +43,7 @@ #include "iceberg/file_format.h" #include "iceberg/file_reader.h" #include "iceberg/manifest/manifest_entry.h" +#include "iceberg/metadata_columns.h" #include "iceberg/parquet/parquet_register.h" #include "iceberg/partition_spec.h" #include "iceberg/row/partition_values.h" @@ -377,6 +378,36 @@ TEST_F(FileScanTaskReaderTest, OpenWithoutDeletesReadsProjectedSchema) { VerifyStream(&stream, R"([[1, "Foo"], [2, "Bar"], [3, "Baz"]])")); } +TEST_F(FileScanTaskReaderTest, ReadLastUpdatedFromDataSeq) { + ICEBERG_UNWRAP_OR_FAIL( + auto data_file, + MakeDataFile(table_schema_, + R"([[1, "Foo", "blue"], [2, "Bar", "red"], [3, "Baz", "green"]])")); + data_file->first_row_id = 100L; + data_file->data_sequence_number = 5L; + FileScanTask task(data_file); + auto projected_schema = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int32()), + MetadataColumns::kRowId, + MetadataColumns::kLastUpdatedSequenceNumber}, + table_schema_->schema_id()); + + 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))); + auto stream_result = reader->Open(task); + ASSERT_THAT(stream_result, IsOk()); + auto stream = std::move(stream_result.value()); + + ASSERT_NO_FATAL_FAILURE(VerifyStream(&stream, R"([[1, 100, 5], + [2, 101, 5], + [3, 102, 5]])")); +} + TEST_F(FileScanTaskReaderTest, OpenWithPositionDeletesFiltersRowsAndPrunesPos) { ICEBERG_UNWRAP_OR_FAIL( auto data_file, diff --git a/src/iceberg/test/parquet_test.cc b/src/iceberg/test/parquet_test.cc index e163c150b..9adaf4af1 100644 --- a/src/iceberg/test/parquet_test.cc +++ b/src/iceberg/test/parquet_test.cc @@ -223,6 +223,26 @@ class ParquetReaderTest : public TempFileTestBase { ASSERT_TRUE(outfile->Close().ok()); } + std::shared_ptr RowLineageSchema() { + return std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + MetadataColumns::kRowId, + MetadataColumns::kLastUpdatedSequenceNumber, + }); + } + + Result> OpenRowLineageReader( + const std::shared_ptr& schema, + std::optional first_row_id = std::nullopt, + std::optional data_sequence_number = std::nullopt) { + return ReaderFactoryRegistry::Open(FileFormatType::kParquet, + {.path = temp_parquet_file_, + .io = file_io_, + .projection = schema, + .first_row_id = first_row_id, + .data_sequence_number = data_sequence_number}); + } + void VerifyNextBatch(Reader& reader, std::string_view expected_json) { // Boilerplate to get Arrow schema auto schema_result = reader.Schema(); @@ -525,6 +545,36 @@ TEST_F(ParquetReaderTest, ReadMetadataOnlyProjection) { ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, kExpectedJson)); } +TEST_F(ParquetReaderTest, ReadRowLineage) { + temp_parquet_file_ = "row_lineage.parquet"; + CreateSimpleParquetFile(); + + ICEBERG_UNWRAP_OR_FAIL(auto reader, OpenRowLineageReader(RowLineageSchema(), 100L, 7L)); + + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[1, 100, 7], + [2, 101, 7], + [3, 102, 7]])")); +} + +TEST_F(ParquetReaderTest, ReadPartialLineage) { + temp_parquet_file_ = "partial_lineage.parquet"; + CreateSimpleParquetFile(); + + auto schema = RowLineageSchema(); + ICEBERG_UNWRAP_OR_FAIL(auto reader, OpenRowLineageReader(schema)); + + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, R"([[1, null, null], + [2, null, null], + [3, null, null]])")); + + ICEBERG_UNWRAP_OR_FAIL(auto reader_with_row_id, OpenRowLineageReader(schema, 100L)); + + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader_with_row_id, + R"([[1, 100, null], + [2, 101, null], + [3, 102, null]])")); +} + TEST_F(ParquetReaderTest, ReadNestedUnknownProjection) { temp_parquet_file_ = "nested_unknown.parquet"; auto write_schema = std::make_shared(std::vector{ diff --git a/src/iceberg/test/table_scan_test.cc b/src/iceberg/test/table_scan_test.cc index 87db9b505..1fb02762d 100644 --- a/src/iceberg/test/table_scan_test.cc +++ b/src/iceberg/test/table_scan_test.cc @@ -399,6 +399,50 @@ TEST_P(TableScanTest, PlanFilesWithDataManifests) { "/path/to/data2.parquet")); } +TEST_P(TableScanTest, PlanRowLineage) { + auto version = GetParam(); + if (version < 3) { + GTEST_SKIP() << "row lineage is only assigned in v3 manifests"; + } + + constexpr int64_t kSnapshotId = 1000L; + constexpr int64_t kDataSequenceNumber = 5L; + constexpr int64_t kFileSequenceNumber = 7L; + + std::vector data_entries{MakeEntry( + ManifestStatus::kAdded, kSnapshotId, kDataSequenceNumber, + MakeDataFile("/path/to/data.parquet", PartitionValues(std::vector{}), + unpartitioned_spec_->spec_id(), /*record_count=*/3))}; + data_entries[0].file_sequence_number = kFileSequenceNumber; + auto data_manifest = WriteDataManifest(version, kSnapshotId, std::move(data_entries), + unpartitioned_spec_); + std::string manifest_list_path = + WriteManifestList(version, kSnapshotId, kFileSequenceNumber, {data_manifest}); + + auto timestamp_ms = TimePointMsFromUnixMs(1609459200000L); + auto snapshot = + std::make_shared(Snapshot{.snapshot_id = kSnapshotId, + .parent_snapshot_id = std::nullopt, + .sequence_number = kFileSequenceNumber, + .timestamp_ms = timestamp_ms, + .manifest_list = manifest_list_path, + .schema_id = schema_->schema_id(), + .first_row_id = 0L, + .added_rows = 3L}); + auto metadata = ScanTestBase::MakeTableMetadata( + {snapshot}, kSnapshotId, + {{"main", std::make_shared(SnapshotRef{ + .snapshot_id = kSnapshotId, .retention = SnapshotRef::Branch{}})}}, + unpartitioned_spec_); + + ICEBERG_UNWRAP_OR_FAIL(auto builder, DataTableScanBuilder::Make(metadata, file_io_)); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanFiles()); + ASSERT_EQ(tasks.size(), 1); + EXPECT_EQ(tasks[0]->data_file()->first_row_id, 0L); + EXPECT_EQ(tasks[0]->data_file()->data_sequence_number, kDataSequenceNumber); +} + TEST_P(TableScanTest, PlanFilesWithMultipleManifests) { auto version = GetParam();