diff --git a/mkdocs/docs/configuration.md b/mkdocs/docs/configuration.md index 44b5e395a9..cf83ab9c49 100644 --- a/mkdocs/docs/configuration.md +++ b/mkdocs/docs/configuration.md @@ -250,6 +250,7 @@ Tencent Cloud Object Storage (COS) is S3-compatible and can be used with PyIcebe | ----------- | ------------------------ | --------------------------------------------------------- | | hf.endpoint | | Configure the endpoint for Hugging Face | | hf.token | hf_xxx | The Hugging Face token to access HF Datasets repositories | +| hf.revision | main | Pin reads to a specific branch, tag, or commit. Writes always target the repository's default branch. | diff --git a/pyiceberg/io/__init__.py b/pyiceberg/io/__init__.py index 7dbc651214..1c16f7e871 100644 --- a/pyiceberg/io/__init__.py +++ b/pyiceberg/io/__init__.py @@ -99,6 +99,7 @@ GCS_VERSION_AWARE = "gcs.version-aware" HF_ENDPOINT = "hf.endpoint" HF_TOKEN = "hf.token" +HF_REVISION = "hf.revision" @runtime_checkable diff --git a/pyiceberg/io/fsspec.py b/pyiceberg/io/fsspec.py index 7749268ff5..82f20083a3 100644 --- a/pyiceberg/io/fsspec.py +++ b/pyiceberg/io/fsspec.py @@ -68,6 +68,7 @@ GCS_TOKEN, GCS_VERSION_AWARE, HF_ENDPOINT, + HF_REVISION, HF_TOKEN, S3_ACCESS_KEY_ID, S3_ANONYMOUS, @@ -319,6 +320,7 @@ def _hf(properties: Properties) -> AbstractFileSystem: } _ADLS_SCHEMES = frozenset({"abfs", "abfss", "wasb", "wasbs"}) +_HF_SCHEMES = frozenset({"hf"}) class FsspecInputFile(InputFile): @@ -327,16 +329,19 @@ class FsspecInputFile(InputFile): Args: location (str): A URI to a file location. fs (AbstractFileSystem): An fsspec filesystem instance. + fs_kwargs (Properties): Extra keyword arguments forwarded to the underlying filesystem read calls, + e.g. a pinned `revision` for the HuggingFace Hub filesystem. """ - def __init__(self, location: str, fs: AbstractFileSystem): + def __init__(self, location: str, fs: AbstractFileSystem, fs_kwargs: Properties | None = None): self._fs = fs + self._fs_kwargs = fs_kwargs or {} super().__init__(location=location) @override def __len__(self) -> int: """Return the total length of the file, in bytes.""" - object_info = self._fs.info(self.location) + object_info = self._fs.info(self.location, **self._fs_kwargs) if "Size" in object_info: return object_info["Size"] elif "size" in object_info: @@ -346,6 +351,12 @@ def __len__(self) -> int: @override def exists(self) -> bool: """Check whether the location exists.""" + if self._fs_kwargs: + # fsspec's AbstractFileSystem.lexists() doesn't forward **kwargs to exists()/info(), so + # honoring fs_kwargs (e.g. a pinned HuggingFace Hub revision) requires calling exists() + # directly -- it does forward kwargs to info(), with the same broad exception handling + # lexists() would otherwise provide. + return self._fs.exists(self.location, **self._fs_kwargs) return self._fs.lexists(self.location) @override @@ -362,7 +373,7 @@ def open(self, seekable: bool = True) -> InputStream: FileNotFoundError: If the file does not exist. """ try: - return self._fs.open(self.location, "rb") + return self._fs.open(self.location, "rb", **self._fs_kwargs) except FileNotFoundError as e: # To have a consistent error handling experience, make sure exception contains missing file location. raise e if e.filename else FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), self.location) from e @@ -444,7 +455,7 @@ def new_input(self, location: str) -> FsspecInputFile: """ uri = urlparse(location) fs = self._get_fs_from_uri(uri) - return FsspecInputFile(location=location, fs=fs) + return FsspecInputFile(location=location, fs=fs, fs_kwargs=self._get_fs_kwargs(uri.scheme)) @override def new_output(self, location: str) -> FsspecOutputFile: @@ -484,6 +495,17 @@ def _get_fs_from_uri(self, uri: "ParseResult") -> AbstractFileSystem: return self.get_fs(uri.scheme, uri.hostname) return self.get_fs(uri.scheme) + def _get_fs_kwargs(self, scheme: str) -> Properties: + """Get extra keyword arguments to forward to filesystem read calls for a given scheme. + + Only applies to reads (new_input): a pinned HuggingFace Hub revision may be a tag or + commit, neither of which is a valid write target, so writes/deletes always target the + repository's default (writable) branch. + """ + if scheme in _HF_SCHEMES and (revision := self.properties.get(HF_REVISION)): + return {"revision": revision} + return {} + def get_fs(self, scheme: str, hostname: str | None = None) -> AbstractFileSystem: """Get a filesystem for a specific scheme, cached per thread.""" if not hasattr(self._thread_locals, "get_fs_cached"): diff --git a/tests/io/test_fsspec.py b/tests/io/test_fsspec.py index 8739a5964d..1e50dea49e 100644 --- a/tests/io/test_fsspec.py +++ b/tests/io/test_fsspec.py @@ -1105,3 +1105,122 @@ def auth_header(self) -> str: assert requests_mock.last_request is not None assert requests_mock.last_request.headers["Authorization"] == "Bearer via-manager" assert request.url == new_uri + + +def test_fsspec_hf_session_properties() -> None: + session_properties: Properties = { + "hf.endpoint": "https://huggingface.co", + "hf.token": "hf_xxx", + } + + with mock.patch("huggingface_hub.HfFileSystem") as mock_hf_fs: + hf_fileio = FsspecFileIO(properties=session_properties) + filename = str(uuid.uuid4()) + + hf_fileio.new_input(location=f"hf://datasets/user/repo/{filename}") + + mock_hf_fs.assert_called_with( + endpoint="https://huggingface.co", + token="hf_xxx", + ) + + +def test_fsspec_hf_revision_forwarded_to_reads() -> None: + session_properties: Properties = { + "hf.revision": "a-pinned-revision", + } + location = "hf://datasets/user/repo/file.parquet" + + with mock.patch("huggingface_hub.HfFileSystem") as mock_hf_fs: + mock_fs = mock_hf_fs.return_value + mock_fs.info.return_value = {"size": 123} + mock_fs.exists.return_value = True + + hf_fileio = FsspecFileIO(properties=session_properties) + input_file = hf_fileio.new_input(location=location) + + assert len(input_file) == 123 + mock_fs.info.assert_called_with(location, revision="a-pinned-revision") + + assert input_file.exists() is True + mock_fs.exists.assert_called_with(location, revision="a-pinned-revision") + + input_file.open() + mock_fs.open.assert_called_with(location, "rb", revision="a-pinned-revision") + + +def test_fsspec_hf_revision_not_forwarded_to_writes() -> None: + session_properties: Properties = { + "hf.revision": "a-pinned-revision", + } + location = "hf://datasets/user/repo/file.parquet" + + with mock.patch("huggingface_hub.HfFileSystem") as mock_hf_fs: + mock_fs = mock_hf_fs.return_value + mock_fs.lexists.return_value = False + + hf_fileio = FsspecFileIO(properties=session_properties) + output_file = hf_fileio.new_output(location=location) + + output_file.create() + mock_fs.open.assert_called_with(location, "wb") + + +def test_fsspec_hf_no_revision_by_default() -> None: + location = "hf://datasets/user/repo/file.parquet" + + with mock.patch("huggingface_hub.HfFileSystem") as mock_hf_fs: + mock_fs = mock_hf_fs.return_value + mock_fs.info.return_value = {"size": 123} + + hf_fileio = FsspecFileIO(properties={}) + input_file = hf_fileio.new_input(location=location) + + assert len(input_file) == 123 + mock_fs.info.assert_called_with(location) + + +def test_fsspec_non_hf_scheme_does_not_receive_revision_kwarg() -> None: + session_properties: Properties = { + "hf.revision": "a-pinned-revision", + } + + fileio = FsspecFileIO(properties=session_properties) + input_file = fileio.new_input(location="file:///tmp/foo.parquet") + + assert input_file._fs_kwargs == {} + + +@pytest.mark.skipif(not os.environ.get("HF_TOKEN"), reason="Requires a real Hugging Face Hub token in HF_TOKEN") +def test_fsspec_hf_revision_pins_reads_to_a_fixed_commit() -> None: + """Without hf.revision, reads always follow the repo's current default branch. + + This means an Iceberg data-file location pointing into an `hf://` repo isn't reproducible on + its own: if someone pushes a new commit that changes the same path, every future read of that + "immutable" data file silently returns the new content instead of what existed when the + Iceberg snapshot referencing it was written. hf.revision fixes this by letting a table pin + reads to the exact commit that was current when the file was written. + """ + from huggingface_hub import HfApi + + hf_token = os.environ["HF_TOKEN"] + api = HfApi(token=hf_token) + repo_id = f"pyiceberg-hf-revision-test-{uuid.uuid4().hex[:8]}" + api.create_repo(repo_id=repo_id, repo_type="dataset") # private by default + location = f"hf://datasets/{repo_id}/file.txt" + + try: + first_commit = api.upload_file(path_or_fileobj=b"v1", path_in_repo="file.txt", repo_id=repo_id, repo_type="dataset") + api.upload_file(path_or_fileobj=b"v2", path_in_repo="file.txt", repo_id=repo_id, repo_type="dataset") + + # Without hf.revision, reads follow the moving default branch -- the file now reads "v2", + # even though an Iceberg data file referencing it was "written" back when it was "v1". + default_branch_fileio = FsspecFileIO(properties={"hf.token": hf_token}) + assert default_branch_fileio.new_input(location).open().read() == b"v2" + + # Pinning hf.revision to the first commit makes the read reproducible: it keeps returning + # "v1" regardless of what's since been pushed to the default branch. + pinned_fileio = FsspecFileIO(properties={"hf.token": hf_token, "hf.revision": first_commit.oid}) + assert pinned_fileio.new_input(location).open().read() == b"v1" + finally: + api.delete_repo(repo_id=repo_id, repo_type="dataset")