Skip to content

Commit 551c5f4

Browse files
committed
Support row lineage assignment on commit
V3 snapshots set first-row-id but not added-rows, so REST catalogs reject every V3 commit and next-row-id is lost after the first one. Set added-rows to the row IDs that the V3 manifest list writer assigns, as Java does, and remove the unused _calculate_added_rows. Closes #3621
1 parent c152089 commit 551c5f4

3 files changed

Lines changed: 84 additions & 14 deletions

File tree

‎pyiceberg/table/update/snapshot.py‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
ManifestEntry,
5252
ManifestEntryStatus,
5353
ManifestFile,
54+
ManifestListWriterV3,
5455
ManifestWriter,
5556
write_manifest,
5657
write_manifest_list,
@@ -209,19 +210,6 @@ def delete_data_file(self, data_file: DataFile) -> _SnapshotProducer[U]:
209210
self._deleted_data_files.add(data_file)
210211
return self
211212

212-
def _calculate_added_rows(self, manifests: list[ManifestFile]) -> int:
213-
"""Calculate the number of added rows from a list of manifest files."""
214-
added_rows = 0
215-
for manifest in manifests:
216-
if manifest.added_snapshot_id is None or manifest.added_snapshot_id == self._snapshot_id:
217-
if manifest.added_rows_count is None:
218-
raise ValueError(
219-
"Cannot determine number of added rows in snapshot because "
220-
f"the entry for manifest {manifest.manifest_path} is missing the field `added-rows-count`"
221-
)
222-
added_rows += manifest.added_rows_count
223-
return added_rows
224-
225213
@abstractmethod
226214
def _deleted_entries(self) -> list[ManifestEntry]: ...
227215

@@ -335,6 +323,7 @@ def _commit(self) -> UpdatesAndRequirements:
335323
self._written_manifest_lists.append(manifest_list_file_path)
336324

337325
first_row_id: int | None = None
326+
added_rows: int | None = None
338327

339328
if self._transaction.table_metadata.format_version >= 3:
340329
first_row_id = self._transaction.table_metadata.next_row_id
@@ -350,6 +339,10 @@ def _commit(self) -> UpdatesAndRequirements:
350339
) as writer:
351340
writer.add_manifests(new_manifests)
352341

342+
if isinstance(writer, ManifestListWriterV3) and first_row_id is not None:
343+
# row IDs assigned to manifests without a first_row_id, including existing rows of upgraded tables
344+
added_rows = writer.next_row_id - first_row_id
345+
353346
snapshot = Snapshot(
354347
snapshot_id=self._snapshot_id,
355348
parent_snapshot_id=self._parent_snapshot_id,
@@ -358,6 +351,7 @@ def _commit(self) -> UpdatesAndRequirements:
358351
summary=summary,
359352
schema_id=self._transaction.table_metadata.current_schema_id,
360353
first_row_id=first_row_id,
354+
added_rows=added_rows,
361355
)
362356

363357
add_snapshot_update = AddSnapshotUpdate(snapshot=snapshot)

‎tests/integration/test_writes/test_writes.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2532,7 +2532,6 @@ def test_stage_only_overwrite_files(
25322532
assert parent_snapshot_id == [None, first_snapshot, second_snapshot, second_snapshot, second_snapshot]
25332533

25342534

2535-
@pytest.mark.skip("V3 writer support is not enabled.")
25362535
@pytest.mark.integration
25372536
def test_v3_write_and_read_row_lineage(spark: SparkSession, session_catalog: Catalog) -> None:
25382537
"""Test writing to a v3 table and reading with Spark."""
@@ -2569,6 +2568,14 @@ def test_v3_write_and_read_row_lineage(spark: SparkSession, session_catalog: Cat
25692568
"Expected next_row_id to be incremented by the number of added rows"
25702569
)
25712570

2571+
tbl.append(test_data)
2572+
assert tbl.metadata.next_row_id == initial_next_row_id + 2 * len(test_data)
2573+
2574+
rows = spark.sql(f"SELECT _row_id, _last_updated_sequence_number FROM {identifier} ORDER BY _row_id").collect()
2575+
assert [(row._row_id, row._last_updated_sequence_number) for row in rows] == [
2576+
(initial_next_row_id + i, 1 + i // len(test_data)) for i in range(2 * len(test_data))
2577+
]
2578+
25722579

25732580
# RecordBatchReader streaming append/overwrite — see https://github.com/apache/iceberg-python/issues/2152
25742581
#

‎tests/table/test_snapshots.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -858,3 +858,72 @@ def test_manifest_list_has_no_first_row_id_for_v2(table_v2: Table) -> None:
858858

859859
assert writer.call_args.kwargs["format_version"] == 2
860860
assert writer.call_args.kwargs["first_row_id"] is None
861+
862+
863+
@pytest.fixture
864+
def empty_table_v3(tmp_path: Path, arrow_table_simple: pa.Table) -> Table:
865+
"""An empty V3 table; writing V3 metadata files is not supported, so commits stay staged in a transaction."""
866+
from pyiceberg.catalog.noop import NoopCatalog
867+
from pyiceberg.io import load_file_io
868+
from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC
869+
from pyiceberg.table.metadata import new_table_metadata
870+
from pyiceberg.table.sorting import UNSORTED_SORT_ORDER
871+
872+
metadata = new_table_metadata(
873+
schema=Catalog._convert_schema_if_needed(arrow_table_simple.schema, format_version=3),
874+
partition_spec=UNPARTITIONED_PARTITION_SPEC,
875+
sort_order=UNSORTED_SORT_ORDER,
876+
location=str(tmp_path),
877+
properties={"format-version": "3"},
878+
).model_copy(update={"next_row_id": 0})
879+
return Table(
880+
identifier=("default", "lineage"),
881+
metadata=metadata,
882+
metadata_location=f"{tmp_path}/metadata/00000.metadata.json",
883+
io=load_file_io(),
884+
catalog=NoopCatalog("NoopCatalog"),
885+
)
886+
887+
888+
def test_v3_append_assigns_row_lineage(empty_table_v3: Table, arrow_table_simple: pa.Table) -> None:
889+
txn = empty_table_v3.transaction()
890+
891+
txn.append(arrow_table_simple)
892+
first = txn.table_metadata.current_snapshot()
893+
assert first is not None
894+
assert (first.first_row_id, first.added_rows) == (0, 3)
895+
assert txn.table_metadata.next_row_id == 3
896+
897+
txn.append(arrow_table_simple)
898+
second = txn.table_metadata.current_snapshot()
899+
assert second is not None
900+
assert (second.first_row_id, second.added_rows) == (3, 3)
901+
assert txn.table_metadata.next_row_id == 6
902+
903+
904+
def test_v3_add_files_assigns_row_lineage(empty_table_v3: Table, arrow_table_simple: pa.Table, tmp_path: Path) -> None:
905+
import pyarrow.parquet as pq
906+
907+
file_path = str(tmp_path / "data.parquet")
908+
pq.write_table(arrow_table_simple, file_path)
909+
txn = empty_table_v3.transaction()
910+
911+
txn.add_files([file_path])
912+
913+
snapshot = txn.table_metadata.current_snapshot()
914+
assert snapshot is not None
915+
assert (snapshot.first_row_id, snapshot.added_rows) == (0, 3)
916+
assert txn.table_metadata.next_row_id == 3
917+
918+
919+
def test_v3_delete_assigns_row_ids_to_rewritten_rows_only(empty_table_v3: Table, arrow_table_simple: pa.Table) -> None:
920+
txn = empty_table_v3.transaction()
921+
txn.append(arrow_table_simple)
922+
923+
txn.delete("bar == 1")
924+
925+
snapshot = txn.table_metadata.current_snapshot()
926+
assert snapshot is not None
927+
# the rewritten data file holds the 2 surviving rows; the old file's DELETED entry uses no row IDs
928+
assert (snapshot.first_row_id, snapshot.added_rows) == (3, 2)
929+
assert txn.table_metadata.next_row_id == 5

0 commit comments

Comments
 (0)