move under optimize namespace

Author: jayceslesar

pyiceberg/table/__init__.py

@@ -16,15 +16,12 @@
 # under the License.
 from __future__ import annotations
 
-import contextlib
 import itertools
-import logging
 import os
 import uuid
 import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from datetime import timedelta
 from functools import cached_property
 from itertools import chain
 from types import TracebackType
@@ -90,6 +87,7 @@
 from pyiceberg.table.name_mapping import (
     NameMapping,
 )
+from pyiceberg.table.optimize import OptimizeTable
 from pyiceberg.table.refs import SnapshotRef
 from pyiceberg.table.snapshots import (
     Snapshot,
@@ -153,8 +151,6 @@
 
     from pyiceberg.catalog import Catalog
 
-logger = logging.getLogger(__name__)
-
 ALWAYS_TRUE = AlwaysTrue()
 DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE = "downcast-ns-timestamp-to-us-on-write"
 
@@ -912,6 +908,15 @@ def inspect(self) -> InspectTable:
         """
         return InspectTable(self)
 
+    @property
+    def optimize(self) -> OptimizeTable:
+        """Return the OptimizeTable object to optimize.
+
+        Returns:
+            OptimizeTable object based on this Table.
+        """
+        return OptimizeTable(self)
+
     def refresh(self) -> Table:
         """Refresh the current table metadata.
 
@@ -1376,28 +1381,6 @@ def to_polars(self) -> pl.LazyFrame:
 
         return pl.scan_iceberg(self)
 
-    def delete_orphaned_files(self, older_than: timedelta = timedelta(days=3), dry_run: bool = False) -> None:
-        """Delete orphaned files in the table."""
-        location = self.location()
-        orphaned_files = self.inspect.orphaned_files(location, older_than)
-        logger.info(f"Found {len(orphaned_files)} orphaned files at {location}!")
-
-        def _delete(file: str) -> None:
-            # don't error if the file doesn't exist
-            # still catch ctrl-c, etc.
-            with contextlib.suppress(Exception):
-                self.io.delete(file)
-
-        if orphaned_files:
-            if dry_run:
-                logger.info(f"(Dry Run) Deleted {len(orphaned_files)} orphaned files at {location}!")
-            else:
-                executor = ExecutorFactory.get_or_create()
-                deletes = executor.map(_delete, orphaned_files)
-                # exhaust
-                list(deletes)
-                logger.info(f"Deleted {len(orphaned_files)} orphaned files at {location}!")
-
 
 class StaticTable(Table):
     """Load a table directly from a metadata file (i.e., without using a catalog)."""

pyiceberg/table/inspect.py

@@ -16,12 +16,11 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timezone
 from functools import reduce
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Tuple, Union, cast
 
 from pyiceberg.conversions import from_bytes
-from pyiceberg.io import _parse_location
 from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, PartitionFieldSummary
 from pyiceberg.partitioning import PartitionSpec
 from pyiceberg.table.snapshots import Snapshot, ancestors_of
@@ -687,40 +686,3 @@ def all_known_files(self) -> dict[str, set[str]]:
         _all_known_files["datafiles"] = reduce(set.union, files_by_snapshots, set())
 
         return _all_known_files
-
-    def orphaned_files(self, location: str, older_than: timedelta = timedelta(days=3)) -> Set[str]:
-        """Get all the orphaned files in the table.
-
-        Args:
-            location: The location to check for orphaned files.
-            older_than: The time period to check for orphaned files. Defaults to 3 days.
-
-        Returns:
-            A set of orphaned file paths.
-
-        """
-        try:
-            import pyarrow as pa  # noqa: F401
-        except ModuleNotFoundError as e:
-            raise ModuleNotFoundError("For deleting orphaned files PyArrow needs to be installed") from e
-
-        from pyarrow.fs import FileSelector, FileType
-
-        from pyiceberg.io.pyarrow import _fs_from_file_path
-
-        all_known_files = self.all_known_files()
-        flat_known_files: set[str] = reduce(set.union, all_known_files.values(), set())
-
-        fs = _fs_from_file_path(self.tbl.io, location)
-
-        _, _, path = _parse_location(location)
-        selector = FileSelector(path, recursive=True)
-        # filter to just files as it may return directories, and filter on time
-        as_of = datetime.now(timezone.utc) - older_than
-        all_files = [
-            f.path for f in fs.get_file_info(selector) if f.type == FileType.File and (as_of is None or (f.mtime < as_of))
-        ]
-
-        orphaned_files = set(all_files).difference(flat_known_files)
-
-        return orphaned_files

pyiceberg/table/optimize.py

@@ -0,0 +1,107 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import contextlib
+import logging
+from datetime import datetime, timedelta, timezone
+from functools import reduce
+from typing import TYPE_CHECKING, Set
+
+from pyiceberg.io import _parse_location
+from pyiceberg.utils.concurrent import ExecutorFactory
+
+logger = logging.getLogger(__name__)
+
+
+if TYPE_CHECKING:
+    from pyiceberg.table import Table
+
+
+class OptimizeTable:
+    tbl: Table
+
+    def __init__(self, tbl: Table) -> None:
+        self.tbl = tbl
+
+        try:
+            import pyarrow as pa  # noqa
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError("For metadata operations PyArrow needs to be installed") from e
+
+    def orphaned_files(self, location: str, older_than: timedelta = timedelta(days=3)) -> Set[str]:
+        """Get all files which are not referenced in any metadata files of an Iceberg table and can thus be considered "orphaned".
+
+        Args:
+            location: The location to check for orphaned files.
+            older_than: The time period to check for orphaned files. Defaults to 3 days.
+
+        Returns:
+            A set of orphaned file paths.
+        """
+        try:
+            import pyarrow as pa  # noqa: F401
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError("For deleting orphaned files PyArrow needs to be installed") from e
+
+        from pyarrow.fs import FileSelector, FileType
+
+        from pyiceberg.io.pyarrow import _fs_from_file_path
+
+        all_known_files = self.tbl.inspect.all_known_files()
+        flat_known_files: set[str] = reduce(set.union, all_known_files.values(), set())
+
+        fs = _fs_from_file_path(self.tbl.io, location)
+
+        _, _, path = _parse_location(location)
+        selector = FileSelector(path, recursive=True)
+        # filter to just files as it may return directories, and filter on time
+        as_of = datetime.now(timezone.utc) - older_than
+        all_files = [
+            f.path for f in fs.get_file_info(selector) if f.type == FileType.File and (as_of is None or (f.mtime < as_of))
+        ]
+
+        orphaned_files = set(all_files).difference(flat_known_files)
+
+        return orphaned_files
+
+    def remove_orphaned_files(self, older_than: timedelta = timedelta(days=3), dry_run: bool = False) -> None:
+        """Remove files which are not referenced in any metadata files of an Iceberg table and can thus be considered "orphaned".
+
+        Args:
+            older_than: The time period to check for orphaned files. Defaults to 3 days.
+            dry_run: If True, only log the files that would be deleted. Defaults to False.
+        """
+        location = self.tbl.location()
+        orphaned_files = self.orphaned_files(location, older_than)
+        logger.info(f"Found {len(orphaned_files)} orphaned files at {location}!")
+
+        def _delete(file: str) -> None:
+            # don't error if the file doesn't exist
+            # still catch ctrl-c, etc.
+            with contextlib.suppress(Exception):
+                self.tbl.io.delete(file)
+
+        if orphaned_files:
+            if dry_run:
+                logger.info(f"(Dry Run) Deleted {len(orphaned_files)} orphaned files at {location}!")
+            else:
+                executor = ExecutorFactory.get_or_create()
+                deletes = executor.map(_delete, orphaned_files)
+                # exhaust
+                list(deletes)
+                logger.info(f"Deleted {len(orphaned_files)} orphaned files at {location}!")

tests/table/test_delete_orphans.py tests/table/test_remove_orphans.pytests/table/test_delete_orphans.py renamed to tests/table/test_remove_orphans.py

@@ -17,7 +17,7 @@
 import os
 from datetime import datetime, timedelta
 from pathlib import Path, PosixPath
-from unittest.mock import PropertyMock, patch
+from unittest.mock import patch
 
 import pyarrow as pa
 import pytest
@@ -35,8 +35,8 @@ def catalog(tmp_path: PosixPath) -> InMemoryCatalog:
     return catalog
 
 
-def test_delete_orphaned_files(catalog: Catalog) -> None:
-    identifier = "default.test_delete_orphaned_files"
+def test_remove_orphaned_files(catalog: Catalog) -> None:
+    identifier = "default.test_remove_orphaned_files"
 
     schema = Schema(
         NestedField(1, "city", StringType(), required=True),
@@ -69,17 +69,17 @@ def test_delete_orphaned_files(catalog: Catalog) -> None:
     assert orphaned_file.exists()
 
     # assert no files deleted if dry run...
-    tbl.delete_orphaned_files(dry_run=True)
+    tbl.optimize.remove_orphaned_files(dry_run=True)
     assert orphaned_file.exists()
 
     # should not delete because it was just created...
-    tbl.delete_orphaned_files()
+    tbl.optimize.remove_orphaned_files()
     assert orphaned_file.exists()
 
     # modify creation date to be older than 3 days
     five_days_ago = (datetime.now() - timedelta(days=5)).timestamp()
     os.utime(orphaned_file, (five_days_ago, five_days_ago))
-    tbl.delete_orphaned_files()
+    tbl.optimize.remove_orphaned_files()
     assert not orphaned_file.exists()
 
     # assert that all known files still exist...
@@ -89,8 +89,8 @@ def test_delete_orphaned_files(catalog: Catalog) -> None:
             assert Path(file).exists()
 
 
-def test_delete_orphaned_files_with_invalid_file_doesnt_error(catalog: Catalog) -> None:
-    identifier = "default.test_delete_orphaned_files"
+def test_remove_orphaned_files_with_invalid_file_doesnt_error(catalog: Catalog) -> None:
+    identifier = "default.test_remove_orphaned_files"
 
     schema = Schema(
         NestedField(1, "city", StringType(), required=True),
@@ -118,8 +118,7 @@ def test_delete_orphaned_files_with_invalid_file_doesnt_error(catalog: Catalog)
     tbl.append(df)
 
     file_that_does_not_exist = "foo/bar.baz"
-    with patch.object(type(tbl), "inspect", new_callable=PropertyMock) as mock_inspect:
-        mock_inspect.return_value.orphaned_files = lambda location, older_than: {file_that_does_not_exist}
+    with patch.object(type(tbl.optimize), "orphaned_files", return_value={file_that_does_not_exist}):
         with patch.object(tbl.io, "delete", wraps=tbl.io.delete) as mock_delete:
-            tbl.delete_orphaned_files()
+            tbl.optimize.remove_orphaned_files(timedelta(days=3))
             mock_delete.assert_called_with(file_that_does_not_exist)