diff --git a/pyiceberg/io/__init__.py b/pyiceberg/io/__init__.py index c44e105e62..0f402acee6 100644 --- a/pyiceberg/io/__init__.py +++ b/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.""" @@ -306,6 +318,20 @@ def delete(self, location: str | InputFile | OutputFile) -> None: FileNotFoundError: When the file at the provided location does not exist. """ + def list_prefix(self, location: str) -> Iterator[FileEntry]: + """Recursively list every file under the given location. + + Args: + location (str): A URI or path to recursively list. + + Returns: + Iterator[FileEntry]: The metadata of every file under the location. + + Raises: + NotImplementedError: If the FileIO implementation does not support listing. + """ + raise NotImplementedError(f"{type(self).__name__} does not support list_prefix") + LOCATION = "location" WAREHOUSE = "warehouse" diff --git a/pyiceberg/io/fsspec.py b/pyiceberg/io/fsspec.py index 09bbe6f1d6..1014330567 100644 --- a/pyiceberg/io/fsspec.py +++ b/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,6 +87,7 @@ S3_SIGNER_ENDPOINT_DEFAULT, S3_SIGNER_URI, S3_SSE_KMS_KEY_ID, + FileEntry, FileIO, InputFile, InputStream, @@ -491,6 +493,46 @@ 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 lists paths without a scheme, and adlfs also drops the account from the authority, so + # each path is turned back into a URI that matches the locations recorded in table metadata. + # On Windows a drive letter parses as a URI scheme, so local paths are reported as-is. + 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 | None + if isinstance(mtime, datetime): + last_modified = mtime + elif isinstance(mtime, (int, float)): + last_modified = datetime.fromtimestamp(mtime, tz=timezone.utc) + else: + last_modified = None + + if not scheme: + file_location = path + elif scheme in _ADLS_SCHEMES: + file_location = f"{scheme}://{uri.netloc}/{path.partition('/')[2]}" + else: + file_location = f"{scheme}://{path}" + + yield FileEntry( + location=file_location, + size=int(info.get("size") or 0), + 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): diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..9bbe81fd4c 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -61,6 +61,7 @@ from pyarrow._s3fs import S3RetryStrategy from pyarrow.fs import ( FileInfo, + FileSelector, FileSystem, FileType, ) @@ -116,6 +117,7 @@ S3_ROLE_SESSION_NAME, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, + FileEntry, FileIO, InputFile, InputStream, @@ -694,6 +696,38 @@ 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 reports paths without a scheme, and for object stores the bucket is part of + # the path, so the prefix that reconstructs the original URI differs per scheme. + 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 or 0, + 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__) diff --git a/tests/io/test_fsspec.py b/tests/io/test_fsspec.py index 45835a08eb..2ed80808f7 100644 --- a/tests/io/test_fsspec.py +++ b/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 @@ -57,6 +59,29 @@ 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(...)""" + (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 +658,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") diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..b33723ab5a 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -147,6 +147,29 @@ 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(...)""" + (tmp_path / "nested").mkdir() + (tmp_path / "a.txt").write_bytes(b"foo") + (tmp_path / "nested" / "b.txt").write_bytes(b"barr") + + entries = sorted(PyArrowFileIO().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"""