feat: hf.revision property for HuggingFace Hub locations

mkdocs/docs/configuration.md

@@ -272,6 +272,7 @@ Tencent Cloud Object Storage (COS) is S3-compatible and can be used with PyIcebe
 | ----------- | ------------------------ | --------------------------------------------------------- |
 | hf.endpoint | <https://huggingface.co> | Configure the endpoint for Hugging Face                   |
 | hf.token    | hf_xxx                   | The Hugging Face token to access HF Datasets repositories |
+| hf.revision | main                     | The branch, tag or commit to use for locations that don't pin one themselves as `hf://datasets/user/repo@revision/path`. Defaults to the repository's default branch. |
 
 <!-- markdown-link-check-enable-->
 

pyiceberg/io/__init__.py

@@ -114,6 +114,7 @@ def _is_local_path(path: str) -> bool:
 GCS_VERSION_AWARE = "gcs.version-aware"
 HF_ENDPOINT = "hf.endpoint"
 HF_TOKEN = "hf.token"
+HF_REVISION = "hf.revision"
 
 
 @runtime_checkable

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,
@@ -332,6 +333,25 @@ def _hf(properties: Properties) -> AbstractFileSystem:
 }
 
 _ADLS_SCHEMES = frozenset({"abfs", "abfss", "wasb", "wasbs"})
+_HF_SCHEMES = frozenset({"hf"})
+
+
+def _has_revision_in_path(uri: "ParseResult") -> bool:
+    """Check whether a HuggingFace Hub location pins a revision itself."""
+    # The revision is attached to the repo id, which spans at most the netloc and the two path
+    # segments following it: hf://<repo_type>/<owner>/<repo>@<revision>/<path_in_repo>.
+    return any("@" in segment for segment in (uri.netloc, *uri.path.lstrip("/").split("/")[:2]))
+
+
+def _exists(fs: AbstractFileSystem, location: str, fs_kwargs: Properties) -> bool:
+    """Check whether a location exists, honoring extra filesystem keyword arguments."""
+    if 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 fs.exists(location, **fs_kwargs)
+    return fs.lexists(location)
 
 
 class FsspecInputFile(InputFile):
@@ -340,16 +360,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 calls,
+            e.g. a `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:
@@ -359,7 +382,7 @@ def __len__(self) -> int:
     @override
     def exists(self) -> bool:
         """Check whether the location exists."""
-        return self._fs.lexists(self.location)
+        return _exists(self._fs, self.location, self._fs_kwargs)
 
     @override
     def open(self, seekable: bool = True) -> InputStream:
@@ -375,7 +398,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
@@ -387,16 +410,19 @@ class FsspecOutputFile(OutputFile):
     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 calls,
+            e.g. a `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:
@@ -406,7 +432,7 @@ def __len__(self) -> int:
     @override
     def exists(self) -> bool:
         """Check whether the location exists."""
-        return self._fs.lexists(self.location)
+        return _exists(self._fs, self.location, self._fs_kwargs)
 
     @override
     def create(self, overwrite: bool = False) -> OutputStream:
@@ -429,12 +455,12 @@ def create(self, overwrite: bool = False) -> OutputStream:
         """
         if not overwrite and self.exists():
             raise FileExistsError(f"Cannot create file, file already exists: {self.location}")
-        return self._fs.open(self.location, "wb")
+        return self._fs.open(self.location, "wb", **self._fs_kwargs)
 
     @override
     def to_input_file(self) -> FsspecInputFile:
         """Return a new FsspecInputFile for the location at `self.location`."""
-        return FsspecInputFile(location=self.location, fs=self._fs)
+        return FsspecInputFile(location=self.location, fs=self._fs, fs_kwargs=self._fs_kwargs)
 
 
 class FsspecFileIO(FileIO):
@@ -457,7 +483,7 @@ def new_input(self, location: str) -> FsspecInputFile:
         """
         uri = urlparse(location)
         fs = self._get_fs_from_uri(uri, location)
-        return FsspecInputFile(location=location, fs=fs)
+        return FsspecInputFile(location=location, fs=fs, fs_kwargs=self._get_fs_kwargs(uri))
 
     @override
     def new_output(self, location: str) -> FsspecOutputFile:
@@ -471,7 +497,7 @@ def new_output(self, location: str) -> FsspecOutputFile:
         """
         uri = urlparse(location)
         fs = self._get_fs_from_uri(uri, location)
-        return FsspecOutputFile(location=location, fs=fs)
+        return FsspecOutputFile(location=location, fs=fs, fs_kwargs=self._get_fs_kwargs(uri))
 
     @override
     def delete(self, location: str | InputFile | OutputFile) -> None:
@@ -489,7 +515,7 @@ def delete(self, location: str | InputFile | OutputFile) -> None:
 
         uri = urlparse(str_location)
         fs = self._get_fs_from_uri(uri, str_location)
-        fs.rm(str_location)
+        fs.rm(str_location, **self._get_fs_kwargs(uri))
 
     def _get_fs_from_uri(self, uri: "ParseResult", location: str = "") -> AbstractFileSystem:
         """Get a filesystem from a parsed URI, using hostname for ADLS account resolution."""
@@ -499,6 +525,17 @@ def _get_fs_from_uri(self, uri: "ParseResult", location: str = "") -> AbstractFi
             return self.get_fs(uri.scheme, uri.hostname)
         return self.get_fs(uri.scheme)
 
+    def _get_fs_kwargs(self, uri: "ParseResult") -> Properties:
+        """Get extra keyword arguments to forward to the filesystem calls for a given location.
+
+        `hf.revision` is the default revision for locations that don't pin one themselves, matching
+        iceberg-rust. A revision embedded in the location (`hf://datasets/user/repo@revision/path`)
+        takes precedence, since huggingface_hub rejects two conflicting revisions.
+        """
+        if uri.scheme in _HF_SCHEMES and not _has_revision_in_path(uri) 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"):

tests/io/test_fsspec.py

@@ -1152,3 +1152,145 @@ 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_read_calls() -> 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_forwarded_to_write_calls() -> 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.exists.return_value = False
+
+        hf_fileio = FsspecFileIO(properties=session_properties)
+        output_file = hf_fileio.new_output(location=location)
+
+        output_file.create()
+        mock_fs.exists.assert_called_with(location, revision="a-pinned-revision")
+        mock_fs.open.assert_called_with(location, "wb", revision="a-pinned-revision")
+
+        hf_fileio.delete(location)
+        mock_fs.rm.assert_called_with(location, revision="a-pinned-revision")
+
+
+def test_fsspec_hf_revision_in_location_takes_precedence() -> None:
+    session_properties: Properties = {
+        "hf.revision": "a-pinned-revision",
+    }
+    location = "hf://datasets/user/repo@another-revision/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=session_properties)
+        input_file = hf_fileio.new_input(location=location)
+
+        # huggingface_hub raises on a revision kwarg conflicting with the one in the path
+        assert len(input_file) == 123
+        mock_fs.info.assert_called_with(location)
+
+
+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.integration
+@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")