Partial deletes can bring back rows removed by earlier delete files

import struct
import tempfile
import uuid

import pyarrow as pa
import pyarrow.parquet as pq
from pyiceberg.catalog.sql import SqlCatalog
from pyiceberg.expressions import EqualTo
from pyiceberg.manifest import (
    DataFile,
    DataFileContent,
    FileFormat,
    ManifestContent,
    ManifestEntry,
    ManifestEntryStatus,
    ManifestWriterV2,
    write_manifest_list,
)
from pyiceberg.schema import Schema
from pyiceberg.table.refs import SnapshotRefType
from pyiceberg.table.snapshots import (
    Operation,
    Snapshot,
    SnapshotSummaryCollector,
    Summary,
    update_snapshot_summaries,
)
from pyiceberg.table.update import AddSnapshotUpdate, SetSnapshotRefUpdate
from pyiceberg.typedef import Record
from pyiceberg.types import LongType, NestedField, StringType

# PyIceberg doesn't write delete files, so this commits one with its low-level writers


class DeleteManifestWriter(ManifestWriterV2):
    def content(self):
        return ManifestContent.DELETES

    @property
    def _meta(self):
        return {**super()._meta, "content": "deletes"}


def field(name, type_, field_id):
    return pa.field(name, type_, nullable=False, metadata={"PARQUET:field_id": str(field_id)})


def add_delete_file(table, content, rows, bounds, equality_ids=None):
    """Commit one delete file with `rows` (a pyarrow table) and these column bounds."""
    path = f"{table.location()}/data/delete-{uuid.uuid4()}.parquet"
    pq.write_table(rows, path.removeprefix("file://"))
    delete_file = DataFile.from_args(
        content=content,
        file_path=path,
        file_format=FileFormat.PARQUET,
        partition=Record(),
        record_count=rows.num_rows,
        file_size_in_bytes=len(table.io.new_input(path)),
        lower_bounds=bounds,
        upper_bounds=bounds,
        equality_ids=equality_ids,
    )
    delete_file.spec_id = table.spec().spec_id
    parent = table.current_snapshot()
    snapshot_id = uuid.uuid4().int >> 65
    sequence_number = table.metadata.last_sequence_number + 1
    manifest = table.io.new_output(f"{table.location()}/metadata/{uuid.uuid4()}-m0.avro")
    with DeleteManifestWriter(
        spec=table.spec(),
        schema=table.schema(),
        output_file=manifest,
        snapshot_id=snapshot_id,
        avro_compression="deflate",
    ) as writer:
        writer.add(
            ManifestEntry.from_args(
                status=ManifestEntryStatus.ADDED, snapshot_id=snapshot_id, data_file=delete_file
            )
        )
    manifest_list = f"{table.location()}/metadata/snap-{snapshot_id}.avro"
    with write_manifest_list(
        format_version=2,
        output_file=table.io.new_output(manifest_list),
        snapshot_id=snapshot_id,
        parent_snapshot_id=parent.snapshot_id,
        sequence_number=sequence_number,
        avro_compression="deflate",
    ) as list_writer:
        list_writer.add_manifests([writer.to_manifest_file(), *parent.manifests(table.io)])
    collector = SnapshotSummaryCollector()
    collector.add_file(delete_file, schema=table.schema())
    summary = update_snapshot_summaries(
        Summary(Operation.DELETE, **collector.build()), previous_summary=parent.summary
    )
    snapshot = Snapshot(
        snapshot_id=snapshot_id,
        parent_snapshot_id=parent.snapshot_id,
        sequence_number=sequence_number,
        manifest_list=manifest_list,
        summary=summary,
        schema_id=table.schema().schema_id,
    )
    ref = SetSnapshotRefUpdate(
        ref_name="main", type=SnapshotRefType.BRANCH, snapshot_id=snapshot_id
    )
    table.catalog.commit_table(table, (), (AddSnapshotUpdate(snapshot=snapshot), ref))
    return table.catalog.load_table(table.name())


def long(x):
    return struct.pack("<q", x)


tmp = tempfile.mkdtemp()
catalog = SqlCatalog("lake", uri=f"sqlite:///{tmp}/c.db", warehouse=f"file://{tmp}")
catalog.create_namespace("ns")
schema = Schema(NestedField(1, "k", LongType()), NestedField(2, "v", StringType()))
data = pa.table({"k": [1, 2, 3], "v": ["a", "b", "c"]})

# 1. A position delete that stores the deleted row, as the spec allows
table = catalog.create_table(("ns", "pos"), schema=schema)
table.append(data)
data_path = next(iter(table.scan().plan_files())).file.file_path
row = pa.struct([field("k", pa.int64(), 1), field("v", pa.string(), 2)])
deletes = pa.Table.from_pylist(
    [{"file_path": data_path, "pos": 1, "row": {"k": 2, "v": "b"}}],
    schema=pa.schema(
        [
            field("file_path", pa.string(), 2147483546),
            field("pos", pa.int64(), 2147483545),
            field("row", row, 2147483544),
        ]
    ),
)
bounds = {2147483546: data_path.encode(), 2147483545: long(1), 1: long(2), 2: b"b"}
table = add_delete_file(table, DataFileContent.POSITION_DELETES, deletes, bounds)
print(table.scan().to_arrow().sort_by("k").to_pylist())
# [{'k': 1, 'v': 'a'}, {'k': 3, 'v': 'c'}]

# The filtered plan leaves out the delete file, as its bounds exclude k = 1
print([len(t.delete_files) for t in table.scan(row_filter=EqualTo("k", 1)).plan_files()])
# [0]
print([len(t.delete_files) for t in table.scan().plan_files()])
# [1]

table.delete(EqualTo("k", 1))
table = catalog.load_table(("ns", "pos"))
print(table.scan().to_arrow().sort_by("k").to_pylist())
# [{'k': 2, 'v': 'b'}, {'k': 3, 'v': 'c'}]: expected [{'k': 3, 'v': 'c'}]

# 2. An equality delete of k = 2
table = catalog.create_table(("ns", "eq"), schema=schema)
table.append(data)
deletes = pa.Table.from_pylist([{"k": 2}], schema=pa.schema([field("k", pa.int64(), 1)]))
table = add_delete_file(
    table, DataFileContent.EQUALITY_DELETES, deletes, {1: long(2)}, equality_ids=[1]
)
# PyIceberg refuses to read equality deletes, but the filtered plan leaves it out
table.delete(EqualTo("k", 1))
table = catalog.load_table(("ns", "eq"))
for manifest in table.current_snapshot().manifests(table.io):
    for entry in manifest.fetch_manifest_entry(table.io):
        file = entry.data_file
        rows = pq.read_table(file.file_path.removeprefix("file://")).to_pylist()
        print(entry.sequence_number, file.content.name, rows)
# 3 DATA [{'k': 2, 'v': 'b'}, {'k': 3, 'v': 'c'}]
# 2 EQUALITY_DELETES [{'k': 2}]
# The new data file holds (2, 'b'). The equality delete has a lower sequence number,
# so it no longer applies to that file.