diff --git a/pyiceberg/io/__init__.py b/pyiceberg/io/__init__.py index c44e105e62..61fc934e8c 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.""" @@ -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" diff --git a/pyiceberg/io/fsspec.py b/pyiceberg/io/fsspec.py index 09bbe6f1d6..9adefb5539 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,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): diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..83a700fda2 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,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__) diff --git a/tests/io/test_fsspec.py b/tests/io/test_fsspec.py index 45835a08eb..50ff7c29d1 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 @@ -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") diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..d3495b959c 100644 --- a/tests/io/test_pyarrow.py +++ b/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"""