Add FileIO.list_prefix

pyiceberg/io/__init__.py

@@ -30,6 +30,9 @@
 import os
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from dataclasses import dataclass
+from datetime import datetime
 from io import SEEK_SET
 from types import TracebackType
 from typing import (
@@ -269,6 +272,15 @@ def create(self, overwrite: bool = False) -> OutputStream:
         """
 
 
+@dataclass(frozen=True)
+class FileEntry:
+    """Metadata of a single file."""
+
+    location: str
+    size: int
+    last_modified: datetime | None = None
+
+
 class FileIO(ABC):
     """A base class for FileIO implementations."""
 
@@ -307,6 +319,25 @@ def delete(self, location: str | InputFile | OutputFile) -> None:
         """
 
 
+class SupportsPrefixOperations(ABC):
+    """An extension for FileIO implementations that support prefix based operations."""
+
+    @abstractmethod
+    def list_prefix(self, location: str) -> Iterator[FileEntry]:
+        """Recursively list every file under the given location.
+
+        Listing is paged and expensive on object stores, so prefer a storage specific inventory
+        for anything beyond low-volume maintenance. Hierarchical filesystems may require the
+        prefix to be a directory, while object stores allow for arbitrary prefixes.
+
+        Args:
+            location (str): A URI or path to recursively list.
+
+        Returns:
+            Iterator[FileEntry]: The metadata of every file under the location.
+        """
+
+
 LOCATION = "location"
 WAREHOUSE = "warehouse"
 

pyiceberg/io/fsspec.py

@@ -22,8 +22,9 @@
 import logging
 import os
 import threading
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from copy import copy
+from datetime import datetime, timezone
 from functools import lru_cache
 from typing import (
     TYPE_CHECKING,
@@ -86,11 +87,13 @@
     S3_SIGNER_ENDPOINT_DEFAULT,
     S3_SIGNER_URI,
     S3_SSE_KMS_KEY_ID,
+    FileEntry,
     FileIO,
     InputFile,
     InputStream,
     OutputFile,
     OutputStream,
+    SupportsPrefixOperations,
     _is_local_path,
 )
 from pyiceberg.typedef import Properties
@@ -437,7 +440,7 @@ def to_input_file(self) -> FsspecInputFile:
         return FsspecInputFile(location=self.location, fs=self._fs)
 
 
-class FsspecFileIO(FileIO):
+class FsspecFileIO(FileIO, SupportsPrefixOperations):
     """A FileIO implementation that uses fsspec."""
 
     def __init__(self, properties: Properties):
@@ -491,6 +494,35 @@ def delete(self, location: str | InputFile | OutputFile) -> None:
         fs = self._get_fs_from_uri(uri, str_location)
         fs.rm(str_location)
 
+    @override
+    def list_prefix(self, location: str) -> Iterator[FileEntry]:
+        """Recursively list every file under the given location.
+
+        Args:
+            location (str): A URI or a path to recursively list.
+
+        Returns:
+            Iterator[FileEntry]: The metadata of every file under the location.
+        """
+        uri = urlparse(location)
+        fs = self._get_fs_from_uri(uri, location)
+        # fsspec strips the scheme from the listed paths, so it is put back to match table metadata
+        scheme = "" if _is_local_path(location) else uri.scheme
+
+        for path, info in fs.find(location, detail=True).items():
+            mtime = info.get("mtime") or info.get("LastModified") or info.get("last_modified")
+            last_modified = datetime.fromtimestamp(mtime, tz=timezone.utc) if isinstance(mtime, (int, float)) else mtime
+
+            if not scheme:
+                file_location = path
+            elif scheme in _ADLS_SCHEMES:
+                # adlfs also drops the account from the authority
+                file_location = f"{scheme}://{uri.netloc}/{path.partition('/')[2]}"
+            else:
+                file_location = f"{scheme}://{path}"
+
+            yield FileEntry(location=file_location, size=info["size"], last_modified=last_modified)
+
     def _get_fs_from_uri(self, uri: "ParseResult", location: str = "") -> AbstractFileSystem:
         """Get a filesystem from a parsed URI, using hostname for ADLS account resolution."""
         if _is_local_path(location):

pyiceberg/io/pyarrow.py

@@ -61,6 +61,7 @@
 from pyarrow._s3fs import S3RetryStrategy
 from pyarrow.fs import (
     FileInfo,
+    FileSelector,
     FileSystem,
     FileType,
 )
@@ -116,11 +117,13 @@
     S3_ROLE_SESSION_NAME,
     S3_SECRET_ACCESS_KEY,
     S3_SESSION_TOKEN,
+    FileEntry,
     FileIO,
     InputFile,
     InputStream,
     OutputFile,
     OutputStream,
+    SupportsPrefixOperations,
     _is_local_path,
 )
 from pyiceberg.io.fileformat import DataFileStatistics as DataFileStatistics
@@ -393,7 +396,7 @@ def to_input_file(self) -> PyArrowFile:
         return self
 
 
-class PyArrowFileIO(FileIO):
+class PyArrowFileIO(FileIO, SupportsPrefixOperations):
     fs_by_scheme: Callable[[str, str | None], FileSystem]
 
     def __init__(self, properties: Properties = EMPTY_DICT):
@@ -694,6 +697,33 @@ def delete(self, location: str | InputFile | OutputFile) -> None:
                 raise PermissionError(f"Cannot delete file, access denied: {location}") from e
             raise  # pragma: no cover - If some other kind of OSError, raise the raw error
 
+    @override
+    def list_prefix(self, location: str) -> Iterator[FileEntry]:
+        """Recursively list every file under the given location.
+
+        Args:
+            location (str): A URI or a path to recursively list.
+
+        Returns:
+            Iterator[FileEntry]: The metadata of every file under the location.
+        """
+        scheme, netloc, path = self.parse_location(location, self.properties)
+        fs = self.fs_by_scheme(scheme, netloc)
+        selector = FileSelector(path, recursive=True, allow_not_found=True)
+
+        # PyArrow strips the scheme from the listed paths, so it is put back to match table metadata
+        original_scheme = "" if _is_local_path(location) else urlparse(location).scheme
+        if original_scheme in ("hdfs", "viewfs"):
+            uri_prefix = f"{original_scheme}://{netloc}"
+        elif original_scheme:
+            uri_prefix = f"{original_scheme}://"
+        else:
+            uri_prefix = ""
+
+        for info in fs.get_file_info(selector):
+            if info.type == FileType.File:
+                yield FileEntry(location=f"{uri_prefix}{info.path}", size=info.size, last_modified=info.mtime)
+
     def __getstate__(self) -> dict[str, Any]:
         """Create a dictionary of the PyArrowFileIO fields used when pickling."""
         fileio_copy = copy(self.__dict__)

tests/io/test_fsspec.py

@@ -17,9 +17,11 @@
 
 import os
 import pickle
+import sys
 import tempfile
 import threading
 import uuid
+from pathlib import Path
 from unittest import mock
 
 import pytest
@@ -30,7 +32,7 @@
 
 from pyiceberg.catalog.rest.auth import AUTH_MANAGER
 from pyiceberg.exceptions import SignError
-from pyiceberg.io import fsspec
+from pyiceberg.io import SupportsPrefixOperations, fsspec
 from pyiceberg.io.fsspec import FsspecFileIO, S3V4RestSigner
 from pyiceberg.io.pyarrow import PyArrowFileIO
 from pyiceberg.typedef import Properties
@@ -57,6 +59,31 @@ def test_fsspec_local_fs_can_create_path_without_parent_dir(fsspec_fileio: Fsspe
             pytest.fail("Failed to write to file without parent directory")
 
 
+def test_fsspec_list_prefix(fsspec_fileio: FsspecFileIO, tmp_path: Path) -> None:
+    """Test recursively listing a directory using FsspecFileIO.list_prefix(...)"""
+    assert isinstance(fsspec_fileio, SupportsPrefixOperations)
+
+    (tmp_path / "nested").mkdir()
+    (tmp_path / "a.txt").write_bytes(b"foo")
+    (tmp_path / "nested" / "b.txt").write_bytes(b"barr")
+
+    entries = sorted(fsspec_fileio.list_prefix(str(tmp_path)), key=lambda entry: entry.location)
+
+    assert [Path(entry.location) for entry in entries] == [tmp_path / "a.txt", tmp_path / "nested" / "b.txt"]
+    assert [entry.size for entry in entries] == [3, 4]
+    assert all(entry.last_modified is not None for entry in entries)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="A file:// URI cannot carry a Windows drive letter")
+def test_fsspec_list_prefix_retains_scheme(fsspec_fileio: FsspecFileIO, tmp_path: Path) -> None:
+    """Test that a location with a scheme is listed as URIs with that same scheme"""
+    (tmp_path / "a.txt").write_bytes(b"foo")
+
+    entries = list(fsspec_fileio.list_prefix(f"file://{tmp_path}"))
+
+    assert [entry.location for entry in entries] == [f"file://{tmp_path}/a.txt"]
+
+
 def test_fsspec_get_fs_instance_per_thread_caching(fsspec_fileio: FsspecFileIO) -> None:
     """Test that filesystem instances are cached per-thread by `FsspecFileIO.get_fs`"""
     fs_instances: list[AbstractFileSystem] = []
@@ -633,6 +660,20 @@ def test_writing_avro_file_adls(generated_manifest_entry_file: str, adls_fsspec_
     adls_fsspec_fileio.delete(f"abfss://tests/{filename}")
 
 
+@pytest.mark.adls
+def test_fsspec_list_prefix_retains_account_adls(adls_fsspec_fileio: FsspecFileIO, request: pytest.FixtureRequest) -> None:
+    """Test that listing an account-qualified ADLS location keeps the account in every listed URI"""
+    account_name = request.config.getoption("--adls.account-name")
+    prefix = f"abfss://tests@{account_name}.dfs.core.windows.net/{uuid.uuid4()}"
+    with adls_fsspec_fileio.new_output(f"{prefix}/nested/a.txt").create() as f:
+        f.write(b"foo")
+
+    entries = list(adls_fsspec_fileio.list_prefix(prefix))
+
+    assert [entry.location for entry in entries] == [f"{prefix}/nested/a.txt"]
+    adls_fsspec_fileio.delete(f"{prefix}/nested/a.txt")
+
+
 @pytest.mark.adls
 def test_fsspec_pickle_round_trip_aldfs(adls_fsspec_fileio: FsspecFileIO) -> None:
     _test_fsspec_pickle_round_trip(adls_fsspec_fileio, "abfss://tests/foo.txt")

tests/io/test_pyarrow.py

@@ -63,7 +63,7 @@
     Or,
 )
 from pyiceberg.expressions.literals import literal
-from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, InputStream, OutputStream, load_file_io
+from pyiceberg.io import S3_RETRY_STRATEGY_IMPL, InputStream, OutputStream, SupportsPrefixOperations, load_file_io
 from pyiceberg.io.pyarrow import (
     ICEBERG_SCHEMA,
     PYARROW_PARQUET_FIELD_ID_KEY,
@@ -147,6 +147,32 @@ def test_pyarrow_local_fs_can_create_path_without_parent_dir() -> None:
             pytest.fail("Failed to write to file without parent directory")
 
 
+def test_pyarrow_list_prefix(tmp_path: Path) -> None:
+    """Test recursively listing a directory using PyArrowFileIO.list_prefix(...)"""
+    file_io = PyArrowFileIO()
+    assert isinstance(file_io, SupportsPrefixOperations)
+
+    (tmp_path / "nested").mkdir()
+    (tmp_path / "a.txt").write_bytes(b"foo")
+    (tmp_path / "nested" / "b.txt").write_bytes(b"barr")
+
+    entries = sorted(file_io.list_prefix(str(tmp_path)), key=lambda entry: entry.location)
+
+    assert [Path(entry.location) for entry in entries] == [tmp_path / "a.txt", tmp_path / "nested" / "b.txt"]
+    assert [entry.size for entry in entries] == [3, 4]
+    assert all(entry.last_modified is not None for entry in entries)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="A file:// URI cannot carry a Windows drive letter")
+def test_pyarrow_list_prefix_retains_scheme(tmp_path: Path) -> None:
+    """Test that a location with a scheme is listed as URIs with that same scheme"""
+    (tmp_path / "a.txt").write_bytes(b"foo")
+
+    entries = list(PyArrowFileIO().list_prefix(f"file://{tmp_path}"))
+
+    assert [entry.location for entry in entries] == [f"file://{tmp_path}/a.txt"]
+
+
 def test_pyarrow_input_file() -> None:
     """Test reading a file using PyArrowFile"""