final version of position_deletes

Author: amitgilad3

pyiceberg/io/pyarrow.py

@@ -889,17 +889,9 @@ def _construct_fragment(fs: FileSystem, data_file: DataFile, file_format_kwargs:
     return _get_file_format(data_file.file_format, **file_format_kwargs).make_fragment(path, fs)
 
 
-def _read_delete_file(fs: FileSystem, data_file: DataFile) -> pa.Table:
-    positinal_delete_schema = pa.schema(
-        [
-            pa.field("file_path", pa.string(), nullable=False),
-            pa.field("pos", pa.int64(), nullable=False),
-            pa.field("row", pa.int64(), nullable=True),
-        ]
-    )
-
+def _read_delete_file(fs: FileSystem, data_file: DataFile, schema: "pa.Schema") -> pa.Table:
     delete_fragment = _construct_fragment(fs, data_file, file_format_kwargs={"pre_buffer": True, "buffer_size": ONE_MEGABYTE})
-    table = ds.Scanner.from_fragment(fragment=delete_fragment, schema=positinal_delete_schema).to_table()
+    table = ds.Scanner.from_fragment(fragment=delete_fragment, schema=schema).to_table()
     return table
 
 

pyiceberg/table/inspect.py

@@ -20,7 +20,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Tuple
 
 from pyiceberg.conversions import from_bytes
-from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, PartitionFieldSummary
+from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestFile, PartitionFieldSummary
 from pyiceberg.partitioning import PartitionSpec
 from pyiceberg.table.snapshots import Snapshot, ancestors_of
 from pyiceberg.types import PrimitiveType
@@ -384,14 +384,35 @@ def _get_all_manifests_schema(self) -> "pa.Schema":
         all_manifests_schema = all_manifests_schema.append(pa.field("reference_snapshot_id", pa.int64(), nullable=False))
         return all_manifests_schema
 
+    def _get_positional_file_schema(self) -> "pa.Schema":
+        import pyarrow as pa
+
+        from pyiceberg.io.pyarrow import schema_to_pyarrow
+
+        pa_row_struct = schema_to_pyarrow(self.tbl.schema().as_struct())
+        positinal_delete_schema = pa.schema(
+            [
+                pa.field("file_path", pa.string(), nullable=False),
+                pa.field("pos", pa.int64(), nullable=False),
+                pa.field("row", pa_row_struct, nullable=True),
+            ]
+        )
+        return positinal_delete_schema
+
     def _get_positional_deletes_schema(self) -> "pa.Schema":
         import pyarrow as pa
 
+        from pyiceberg.io.pyarrow import schema_to_pyarrow
+
+        partition_record = self.tbl.metadata.specs_struct()
+        pa_partition_struct = schema_to_pyarrow(partition_record)
+        pa_row_struct = schema_to_pyarrow(self.tbl.schema().as_struct())
         positinal_delete_schema = pa.schema(
             [
                 pa.field("file_path", pa.string(), nullable=False),
                 pa.field("pos", pa.int64(), nullable=False),
-                pa.field("row", pa.int64(), nullable=True),
+                pa.field("row", pa_row_struct, nullable=True),
+                pa.field("partition", pa_partition_struct, nullable=False),
                 pa.field("spec_id", pa.int64(), nullable=True),
                 pa.field("delete_file_path", pa.string(), nullable=False),
             ]
@@ -467,23 +488,30 @@ def _partition_summaries_to_rows(
             schema=self._get_all_manifests_schema() if is_all_manifests_table else self._get_manifests_schema(),
         )
 
-    # def _generate_positional_delete_table(self, manifest_list: ManifestFile) -> "pa.Table":
-    #     import pyarrow as pa
-    #     all_deletes = []
-    #     if manifest_list.content == ManifestContent.DELETES:
-    #         for manifest_entry in manifest_list.fetch_manifest_entry(self.tbl.io):
-    #             if manifest_entry.data_file.content == DataFileContent.POSITION_DELETES:
-    #                 from pyiceberg.io.pyarrow import _read_delete_file
-    #                 from pyiceberg.io.pyarrow import _fs_from_file_path
-    #                 positional_delete = _read_delete_file(
-    #                     _fs_from_file_path(self.tbl.io, manifest_entry.data_file.file_path),
-    #                     manifest_entry.data_file)
-    #
-    #                 positional_delete = positional_delete.append_column("spec_id", pa.array(
-    #                     [manifest_list.partition_spec_id] * len(positional_delete))).append_column("delete_file_path",pa.array([manifest_entry.data_file.file_path] * len(positional_delete)))
-    #
-    #                 all_deletes.append(positional_delete)
-    #         return pa.concat_tables(all_deletes)
+    def _generate_positional_delete_table(self, manifest: ManifestFile, position_deletes_schema: "pa.Schema") -> "pa.Table":
+        import pyarrow as pa
+
+        positional_deletes: List["pa.Table"] = []
+        if manifest.content == ManifestContent.DELETES:
+            for entry in manifest.fetch_manifest_entry(self.tbl.io):
+                if entry.data_file.content == DataFileContent.POSITION_DELETES:
+                    from pyiceberg.io.pyarrow import _fs_from_file_path, _read_delete_file
+
+                    positional_delete_file = _read_delete_file(
+                        _fs_from_file_path(self.tbl.io, entry.data_file.file_path),
+                        entry.data_file,
+                        self._get_positional_file_schema(),
+                    ).to_pylist()
+                    for record in positional_delete_file:
+                        record["partition"] = entry.data_file.partition.__dict__
+                        record["spec_id"] = manifest.partition_spec_id
+                        record["delete_file_path"] = entry.data_file.file_path
+
+                    positional_deletes.append(pa.Table.from_pylist(positional_delete_file, position_deletes_schema))
+
+        if not positional_deletes:
+            return pa.Table.from_pylist([], position_deletes_schema)
+        return pa.concat_tables(positional_deletes)
 
     def manifests(self) -> "pa.Table":
         return self._generate_manifests_table(self.tbl.current_snapshot())
@@ -693,38 +721,15 @@ def all_manifests(self) -> "pa.Table":
     def position_deletes(self) -> "pa.Table":
         import pyarrow as pa
 
-        snapshots = self.tbl.snapshots()
-        if not snapshots:
-            return pa.Table.from_pylist([], schema=self._get_positional_deletes_schema())
+        position_deletes_schema = self._get_positional_deletes_schema()
         current_snapshot = self.tbl.current_snapshot()
 
-        #
-        # executor = ExecutorFactory.get_or_create()
-        # positonal_deletes: Iterator["pa.Table"] = executor.map(
-        #     lambda manifest_list: self._generate_positional_delete_table(manifest_list),current_snapshot.manifests(self.tbl.io)
-        # )
-        # all_deletes = []
-        positional_deletes = []
-
-        for manifest_list in current_snapshot.manifests(self.tbl.io):
-            import pyarrow as pa
-            if manifest_list.content == ManifestContent.DELETES:
-                defaultSpecId = self.tbl.spec().spec_id
-                for manifest_entry in manifest_list.fetch_manifest_entry(self.tbl.io):
-
-                    if manifest_entry.data_file.content == DataFileContent.POSITION_DELETES:
-                        from pyiceberg.io.pyarrow import _read_delete_file
-                        from pyiceberg.io.pyarrow import _fs_from_file_path
-                        positional_delete = _read_delete_file(
-                            _fs_from_file_path(self.tbl.io, manifest_entry.data_file.file_path),
-                            manifest_entry.data_file)
-
-                        positional_delete = positional_delete.append_column("spec_id", pa.array(
-                            [manifest_list.partition_spec_id] * len(positional_delete))).append_column("partition", pa.array(
-                            [self.] * len(positional_delete))).append_column(
-                            "delete_file_path", pa.array([manifest_entry.data_file.file_path] * len(positional_delete)))
-
-                        positional_deletes.append(positional_delete)
-                # return pa.concat_tables(all_deletes)
+        if not current_snapshot:
+            return pa.Table.from_pylist([], schema=position_deletes_schema)
 
+        executor = ExecutorFactory.get_or_create()
+        positional_deletes: Iterator["pa.Table"] = executor.map(
+            lambda manifest: self._generate_positional_delete_table(manifest, position_deletes_schema),
+            current_snapshot.manifests(self.tbl.io),
+        )
         return pa.concat_tables(positional_deletes)

tests/integration/test_inspect_table.py

@@ -941,7 +941,7 @@ def test_inspect_all_manifests(spark: SparkSession, session_catalog: Catalog, fo
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize("format_version", [2])
+@pytest.mark.parametrize("format_version", [1, 2])
 def test_inspect_positional_deletes(spark: SparkSession, session_catalog: Catalog, format_version: int) -> None:
     from pandas.testing import assert_frame_equal
 
@@ -964,8 +964,6 @@ def test_inspect_positional_deletes(spark: SparkSession, session_catalog: Catalo
     )
     tbl = session_catalog.load_table(identifier)
 
-    # check all_manifests when there are no snapshots
-
     spark.sql(f"INSERT INTO {identifier} VALUES (1, 'a')")
 
     spark.sql(f"INSERT INTO {identifier} VALUES (2, 'b')")
@@ -977,21 +975,10 @@ def test_inspect_positional_deletes(spark: SparkSession, session_catalog: Catalo
     tbl.refresh()
     df = tbl.inspect.position_deletes()
 
-    assert df.column_names == [
-        "file_path",
-        "pos",
-        "row",
-        "spec_id",
-        "delete_file_path"
-    ]
+    assert df.column_names == ["file_path", "pos", "row", "partition", "spec_id", "delete_file_path"]
 
-    int_cols = [
-        "pos"
-    ]
-    string_cols = [
-        "file_path",
-        "delete_file_path"
-    ]
+    int_cols = ["pos"]
+    string_cols = ["file_path", "delete_file_path"]
 
     for column in int_cols:
         for value in df[column]:
@@ -1001,7 +988,6 @@ def test_inspect_positional_deletes(spark: SparkSession, session_catalog: Catalo
         for value in df[column]:
             assert isinstance(value.as_py(), str)
 
-    new_df = spark.sql(f"select * from {identifier}.position_deletes").toPandas()
     lhs = spark.table(f"{identifier}.position_deletes").toPandas()
     rhs = df.to_pandas()
     assert_frame_equal(lhs, rhs, check_dtype=False)