Đổi schema và đọc snapshot cũ của Apache Iceberg bằng PyIceberg

1python3 -m venv .venv
2source .venv/bin/activate
3python -m pip install 'pyiceberg[sql-sqlite,pyarrow]==0.10.0' 'pyarrow==20.0.0' 'SQLAlchemy==2.0.44'

  1import hashlib
  2import platform
  3from pathlib import Path
  4from tempfile import TemporaryDirectory
  5
  6import pyarrow as pa
  7import pyiceberg
  8import sqlalchemy
  9from pyiceberg.catalog.sql import SqlCatalog
 10from pyiceberg.schema import Schema
 11from pyiceberg.types import LongType, NestedField, StringType
 12
 13
 14def parquet_hashes(warehouse):
 15    return {
 16        path.relative_to(warehouse): hashlib.sha256(path.read_bytes()).hexdigest()
 17        for path in warehouse.rglob("*.parquet")
 18    }
 19
 20
 21def read_rows(table, fields, snapshot_id=None):
 22    data = table.scan(selected_fields=fields, snapshot_id=snapshot_id).to_arrow()
 23    assert tuple(data.column_names) == fields
 24    return sorted(data.to_pylist(), key=lambda row: row["order_id"])
 25
 26
 27def main():
 28    print(f"Python {platform.python_version()} ({platform.machine()}) | "
 29          f"PyIceberg {pyiceberg.__version__} | PyArrow {pa.__version__} | "
 30          f"SQLAlchemy {sqlalchemy.__version__}")
 31    with TemporaryDirectory(prefix="iceberg-orders-") as temp:
 32        root = Path(temp).resolve()
 33        warehouse = root / "warehouse"
 34        warehouse.mkdir()
 35        catalog = SqlCatalog(
 36            "demo", uri=f"sqlite:///{root / 'catalog.db'}", warehouse=warehouse.as_uri()
 37        )
 38        catalog.create_namespace("demo")
 39        name = "demo.orders"
 40        table = catalog.create_table(
 41            name,
 42            schema=Schema(
 43                NestedField(1, "order_id", LongType(), required=True),
 44                NestedField(2, "customer", StringType(), required=False),
 45            ),
 46            properties={"format-version": "2"},
 47        )
 48        original = [
 49            {"order_id": 101, "customer": "An"},
 50            {"order_id": 102, "customer": "Binh"},
 51        ]
 52        table.append(pa.Table.from_pylist(original, schema=table.schema().as_arrow()))
 53        first = table.current_snapshot()
 54        assert first is not None
 55        before_schema = table.schema()
 56        before_files = parquet_hashes(warehouse)
 57        assert len(before_files) == 1
 58        before_metadata = table.metadata_location
 59        print(f"S1 id={first.snapshot_id}, schema_id={first.schema_id}")
 60
 61        with table.update_schema() as update:
 62            update.rename_column("customer", "buyer")
 63            update.add_column("channel", StringType(), required=False)
 64
 65        table = catalog.load_table(name)
 66        after_schema = table.schema()
 67        old_id = before_schema.find_field("customer").field_id
 68        new_id = after_schema.find_field("buyer").field_id
 69        channel = after_schema.find_field("channel")
 70        assert new_id == old_id and channel.field_id != new_id and not channel.required
 71        assert after_schema.schema_id != before_schema.schema_id
 72        assert table.metadata_location != before_metadata
 73        assert [s.snapshot_id for s in table.snapshots()] == [first.snapshot_id]
 74        assert table.current_snapshot().snapshot_id == first.snapshot_id
 75        assert parquet_hashes(warehouse) == before_files
 76        print(f"Schema {before_schema.schema_id} -> {after_schema.schema_id}; "
 77              f"customer ID {old_id} -> buyer ID {new_id}; channel ID {channel.field_id}")
 78        print("After schema commit: 1 snapshot, current=S1, Parquet SHA-256 unchanged")
 79
 80        fields = ("order_id", "buyer", "channel")
 81        expected_old_rows = [
 82            {"order_id": 101, "buyer": "An", "channel": None},
 83            {"order_id": 102, "buyer": "Binh", "channel": None},
 84        ]
 85        after_ddl = read_rows(table, fields)
 86        assert after_ddl == expected_old_rows
 87        print("Current after schema commit:", after_ddl)
 88
 89        appended = {"order_id": 103, "buyer": "Chi", "channel": "web"}
 90        table.append(pa.Table.from_pylist([appended], schema=table.schema().as_arrow()))
 91        table = catalog.load_table(name)
 92        second = table.current_snapshot()
 93        assert second is not None and second.parent_snapshot_id == first.snapshot_id
 94        assert second.snapshot_id != first.snapshot_id and second.schema_id == after_schema.schema_id
 95        assert len(table.snapshots()) == 2
 96        print(f"S2 id={second.snapshot_id}, schema_id={second.schema_id}")
 97
 98        current = read_rows(table, fields)
 99        historical = read_rows(table, ("order_id", "customer"), snapshot_id=first.snapshot_id)
100        assert current == expected_old_rows + [appended]
101        assert historical == original
102        print("Current S2:", current)
103        print("Time travel S1:", historical)
104
105        try:
106            read_rows(table, ("order_id", "buyer"), snapshot_id=first.snapshot_id)
107        except ValueError as error:
108            assert "buyer" in str(error)
109            print("Time travel S1 projection of buyer:", error)
110        else:
111            raise AssertionError("buyer should not be part of the S1 schema")
112        catalog.engine.dispose()
113
114
115if __name__ == "__main__":
116    main()