Add row-lineage support with Spark tests

pyiceberg/manifest.py

@@ -548,6 +548,10 @@ def sort_order_id(self) -> int | None:
     def first_row_id(self) -> int | None:
         return self._data[16]
 
+    @first_row_id.setter
+    def first_row_id(self, value: int | None) -> None:
+        self._data[16] = value
+
     @property
     def referenced_data_file(self) -> str | None:
         return self._data[17]
@@ -916,12 +920,21 @@ def fetch_manifest_entry(
             read_enums={0: ManifestEntryStatus, 101: FileFormat, 134: DataFileContent},
         ) as reader:
             result = []
+            next_row_id = self.first_row_id if self.content == ManifestContent.DATA else None
 
             for entry in reader:
                 if discard_deleted and entry.status == ManifestEntryStatus.DELETED:
                     continue
                 _inherit_from_manifest(entry, self)
 
+                # Inherit first row IDs, which are only valid when a data manifest has a first_row_id
+                data_file = entry.data_file
+                if next_row_id is None:
+                    data_file.first_row_id = None
+                elif entry.status != ManifestEntryStatus.DELETED and data_file.first_row_id is None:
+                    data_file.first_row_id = next_row_id
+                    next_row_id += data_file.record_count
+
                 if entry_filter is None or entry_filter(entry):
                     result.append(entry)
 

pyiceberg/table/metadata.py

@@ -68,6 +68,7 @@
 
 INITIAL_SEQUENCE_NUMBER = 0
 INITIAL_SPEC_ID = 0
+INITIAL_ROW_ID = 0
 DEFAULT_SCHEMA_ID = 0
 
 SUPPORTED_TABLE_FORMAT_VERSION = 2

pyiceberg/table/update/__init__.py

@@ -29,7 +29,7 @@
 from pyiceberg.exceptions import CommitFailedException
 from pyiceberg.partitioning import PARTITION_FIELD_ID_START, PartitionSpec
 from pyiceberg.schema import Schema
-from pyiceberg.table.metadata import SUPPORTED_TABLE_FORMAT_VERSION, TableMetadata, TableMetadataUtil
+from pyiceberg.table.metadata import INITIAL_ROW_ID, SUPPORTED_TABLE_FORMAT_VERSION, TableMetadata, TableMetadataUtil
 from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import (
     MetadataLogEntry,
@@ -458,18 +458,18 @@ def _(update: AddSnapshotUpdate, base_metadata: TableMetadata, context: _TableMe
         )
 
     context.add_update(update)
-    return base_metadata.model_copy(
-        update={
-            "last_updated_ms": update.snapshot.timestamp_ms,
-            "last_sequence_number": update.snapshot.sequence_number,
-            "snapshots": base_metadata.snapshots + [update.snapshot],
-            "next_row_id": base_metadata.next_row_id + update.snapshot.added_rows
-            if base_metadata.format_version >= 3
-            and base_metadata.next_row_id is not None
-            and update.snapshot.added_rows is not None
-            else None,
-        }
-    )
+    metadata_updates: dict[str, Any] = {
+        "last_updated_ms": update.snapshot.timestamp_ms,
+        "last_sequence_number": update.snapshot.sequence_number,
+        "snapshots": base_metadata.snapshots + [update.snapshot],
+    }
+
+    if base_metadata.format_version >= 3:
+        if update.snapshot.added_rows is None:
+            raise ValueError("Cannot add snapshot without added rows")
+        metadata_updates["next_row_id"] = (base_metadata.next_row_id or INITIAL_ROW_ID) + update.snapshot.added_rows
+
+    return base_metadata.model_copy(update=metadata_updates)
 
 
 @_apply_table_update.register(SetSnapshotRefUpdate)

tests/integration/test_reads.py

@@ -1254,6 +1254,34 @@ def test_initial_default(catalog: Catalog, spark: SparkSession) -> None:
     assert len(result_table) == 10
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [lf("session_catalog")])
+def test_read_first_row_ids_written_by_spark(catalog: Catalog, spark: SparkSession) -> None:
+    identifier = "default.test_read_first_row_ids_written_by_spark"
+    spark.sql(f"DROP TABLE IF EXISTS {identifier}")
+    spark.sql(f"CREATE TABLE {identifier} (id int) USING ICEBERG TBLPROPERTIES ('format-version'='3')")
+    spark.sql(f"INSERT INTO {identifier} VALUES (1), (2), (3)")
+    spark.sql(f"INSERT INTO {identifier} VALUES (4), (5)")
+    # Rewriting manifests stores explicit first row IDs on the existing data files
+    spark.sql(f"CALL rest.system.rewrite_manifests('{identifier}')")
+    spark.sql(f"INSERT INTO {identifier} VALUES (6), (7)")
+
+    expected = {
+        row._file: row.first_row_id
+        for row in spark.sql(f"SELECT _file, MIN(_row_id) AS first_row_id FROM {identifier} GROUP BY _file").collect()
+    }
+
+    table = catalog.load_table(identifier)
+    snapshot = table.current_snapshot()
+    assert snapshot is not None
+    actual = {
+        entry.data_file.file_path: entry.data_file.first_row_id
+        for manifest in snapshot.manifests(table.io)
+        for entry in manifest.fetch_manifest_entry(table.io)
+    }
+    assert actual == expected
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")])
 def test_filter_after_arrow_scan(catalog: Catalog) -> None:

tests/table/test_init.py

@@ -1952,6 +1952,40 @@ def test_add_snapshot_update_updates_next_row_id(table_v3: Table) -> None:
     assert new_metadata.next_row_id == 11
 
 
+def test_add_snapshot_update_fails_without_added_rows(table_v3: Table) -> None:
+    new_snapshot = Snapshot(
+        snapshot_id=25,
+        parent_snapshot_id=19,
+        sequence_number=200,
+        timestamp_ms=1602638593590,
+        manifest_list="s3:/a/b/c.avro",
+        summary=Summary(Operation.APPEND),
+        schema_id=3,
+        first_row_id=2,
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="Cannot add snapshot without added rows",
+    ):
+        update_table_metadata(table_v3.metadata, (AddSnapshotUpdate(snapshot=new_snapshot),))
+
+
+def test_add_snapshot_update_keeps_next_row_id_unset_below_v3(table_v2: Table) -> None:
+    new_snapshot = Snapshot(
+        snapshot_id=25,
+        parent_snapshot_id=3055729675574597004,
+        sequence_number=200,
+        timestamp_ms=1602638593590,
+        manifest_list="s3:/a/b/c.avro",
+        summary=Summary(Operation.APPEND),
+        schema_id=3,
+    )
+
+    new_metadata = update_table_metadata(table_v2.metadata, (AddSnapshotUpdate(snapshot=new_snapshot),))
+    assert "next-row-id" not in json.loads(new_metadata.model_dump_json())
+
+
 def model_roundtrips(model: BaseModel) -> bool:
     """Helper assertion that tests if a pydantic model roundtrips
     successfully.

tests/utils/test_manifest.py

@@ -228,6 +228,67 @@ def test_fetch_manifest_entry_with_filter(generated_manifest_entry_file: str) ->
     assert len(no_match) == 0
 
 
+def test_fetch_manifest_entry_inherits_first_row_id(tmp_path: Path) -> None:
+    """Data files inherit first row IDs from data manifests in file order, and have them cleared otherwise."""
+    io = PyArrowFileIO()
+    manifest_path = str(tmp_path / "manifest.avro")
+
+    def entry(status: ManifestEntryStatus, record_count: int, first_row_id: int | None = None) -> ManifestEntry:
+        return ManifestEntry.from_args(
+            _table_format_version=3,
+            status=status,
+            snapshot_id=25,
+            sequence_number=1,
+            file_sequence_number=1,
+            data_file=DataFile.from_args(
+                _table_format_version=3,
+                content=DataFileContent.DATA,
+                file_path=f"s3://bucket/data-{record_count}.parquet",
+                file_format=FileFormat.PARQUET,
+                partition=Record(),
+                record_count=record_count,
+                file_size_in_bytes=1024,
+                first_row_id=first_row_id,
+            ),
+        )
+
+    with AvroOutputFile[ManifestEntry](
+        output_file=io.new_output(manifest_path),
+        file_schema=MANIFEST_ENTRY_SCHEMAS[3],
+        record_schema=MANIFEST_ENTRY_SCHEMAS[3],
+        schema_name="manifest_entry",
+        metadata={"format-version": "3"},
+    ) as writer:
+        writer.write_block(
+            [
+                entry(ManifestEntryStatus.ADDED, 10),
+                entry(ManifestEntryStatus.EXISTING, 5, first_row_id=500),
+                entry(ManifestEntryStatus.DELETED, 7),
+                entry(ManifestEntryStatus.ADDED, 3),
+            ]
+        )
+
+    def first_row_ids(
+        manifest_first_row_id: int | None, discard_deleted: bool, content: ManifestContent = ManifestContent.DATA
+    ) -> list[int | None]:
+        manifest = ManifestFile.from_args(
+            manifest_path=manifest_path,
+            manifest_length=0,
+            partition_spec_id=0,
+            content=content,
+            added_snapshot_id=25,
+            sequence_number=1,
+            min_sequence_number=1,
+            first_row_id=manifest_first_row_id,
+        )
+        return [e.data_file.first_row_id for e in manifest.fetch_manifest_entry(io, discard_deleted=discard_deleted)]
+
+    assert first_row_ids(1000, discard_deleted=False) == [1000, 500, None, 1010]
+    assert first_row_ids(1000, discard_deleted=True) == [1000, 500, 1010]
+    assert first_row_ids(None, discard_deleted=False) == [None, None, None, None]
+    assert first_row_ids(1000, discard_deleted=False, content=ManifestContent.DELETES) == [None, None, None, None]
+
+
 def test_read_manifest_entry_v3_fields(tmp_path: Path) -> None:
     io = PyArrowFileIO()
 
@@ -257,6 +318,7 @@ def write_and_read(file_name: str, data_file: DataFile) -> DataFile:
             added_snapshot_id=25,
             sequence_number=1,
             min_sequence_number=1,
+            first_row_id=0,
         )
         return manifest.fetch_manifest_entry(io)[0].data_file