Repository: huggingface/dataset-viewer

Issue:
Use of hffs dependency is legacy
The `hffs`dependency is legacy and their repo has been archived: https://github.com/huggingface/hffs

We should use `huggingface-hub` >= 0.14 instead.


Discussion:

Resolving patch:
diff --git a/services/api/src/api/routes/rows.py b/services/api/src/api/routes/rows.py
--- a/services/api/src/api/routes/rows.py
+++ b/services/api/src/api/routes/rows.py
@@ -14,7 +14,8 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 from datasets import Features
-from hffs.fs import HfFileSystem
+from huggingface_hub import HfFileSystem
+from huggingface_hub.hf_file_system import safe_quote
 from libcommon.processing_graph import ProcessingGraph
 from libcommon.viewer_utils.asset import (
     glob_rows_in_assets_dir,
@@ -67,19 +68,27 @@ class ParquetDataProcessingError(Exception):
 
 # TODO: how to invalidate the cache when the parquet branch is created or deleted?
 @lru_cache(maxsize=128)
-def get_parquet_fs(dataset: str, hf_token: Optional[str]) -> HfFileSystem:
-    """Get the parquet filesystem for a dataset.
-
-    The parquet files are stored in a separate branch of the dataset repository (see PARQUET_REVISION)
+def get_hf_fs(hf_token: Optional[str]) -> HfFileSystem:
+    """Get the Hugging Face filesystem.
 
     Args:
-        dataset (str): The dataset name.
         hf_token (Optional[str]): The token to access the filesystem.
+    Returns:
+        HfFileSystem: The Hugging Face filesystem.
+    """
+    return HfFileSystem(token=hf_token)
 
+
+def get_hf_parquet_uris(paths: List[str], dataset: str) -> List[str]:
+    """Get the Hugging Face URIs from the Parquet branch of the dataset repository (see PARQUET_REVISION).
+
+    Args:
+        paths (List[str]): List of paths.
+        dataset (str): The dataset name.
     Returns:
-        HfFileSystem: The parquet filesystem.
+        List[str]: List of Parquet URIs.
     """
-    return HfFileSystem(dataset, repo_type="dataset", revision=PARQUET_REVISION, token=hf_token)
+    return [f"hf://datasets/{dataset}@{safe_quote(PARQUET_REVISION)}/{path}" for path in paths]
 
 
 UNSUPPORTED_FEATURES_MAGIC_STRINGS = ["Image(", "Audio(", "'binary'"]
@@ -146,12 +155,14 @@ def __post_init__(
                 if not sources:
                     raise ParquetResponseEmptyError("No parquet files found.")
             with StepProfiler(method="rows.index", step="get the Hub's dataset filesystem"):
-                fs = get_parquet_fs(dataset=self.dataset, hf_token=hf_token)
+                fs = get_hf_fs(hf_token=hf_token)
+            with StepProfiler(method="rows.index", step="get the source URIs"):
+                source_uris = get_hf_parquet_uris(sources, dataset=self.dataset)
             with StepProfiler(method="rows.index", step="get one parquet reader per parquet file"):
                 desc = f"{self.dataset}/{self.config}/{self.split}"
                 try:
                     parquet_files: List[pq.ParquetFile] = thread_map(
-                        partial(pq.ParquetFile, filesystem=fs), sources, desc=desc, unit="pq", disable=True
+                        partial(pq.ParquetFile, filesystem=fs), source_uris, desc=desc, unit="pq", disable=True
                     )
                 except Exception as e:
                     raise FileSystemError(f"Could not read the parquet files: {e}") from e
diff --git a/services/worker/src/worker/job_runners/split/first_rows_from_parquet.py b/services/worker/src/worker/job_runners/split/first_rows_from_parquet.py
--- a/services/worker/src/worker/job_runners/split/first_rows_from_parquet.py
+++ b/services/worker/src/worker/job_runners/split/first_rows_from_parquet.py
@@ -2,14 +2,14 @@
 # Copyright 2022 The HuggingFace Authors.
 
 import logging
-from functools import partial
+from functools import lru_cache, partial
 from http import HTTPStatus
 from typing import Any, List, Literal, Mapping, Optional
 
 import pyarrow as pa
 from datasets import Features
-from fsspec import AbstractFileSystem  # type: ignore
-from hffs.fs import HfFileSystem
+from huggingface_hub import HfFileSystem
+from huggingface_hub.hf_file_system import safe_quote
 from libcommon.constants import (
     PARQUET_REVISION,
     PROCESSING_STEP_SPLIT_FIRST_ROWS_FROM_PARQUET_VERSION,
@@ -135,16 +135,28 @@ def transform_rows(
     ]
 
 
-def get_parquet_fs(dataset: str, hf_token: Optional[str]) -> AbstractFileSystem:
-    """Get the parquet filesystem for a dataset.
-    The parquet files are stored in a separate branch of the dataset repository (see PARQUET_REVISION)
+@lru_cache(maxsize=128)
+def get_hf_fs(hf_token: Optional[str]) -> HfFileSystem:
+    """Get the Hugging Face filesystem.
+
     Args:
-        dataset (str): The dataset name.
         hf_token (Optional[str]): The token to access the filesystem.
     Returns:
-        HfFileSystem: The parquet filesystem.
+        HfFileSystem: The Hugging Face filesystem.
+    """
+    return HfFileSystem(token=hf_token)
+
+
+def get_hf_parquet_uris(paths: List[str], dataset: str) -> List[str]:
+    """Get the Hugging Face URIs from the Parquet branch of the dataset repository (see PARQUET_REVISION).
+
+    Args:
+        paths (List[str]): List of paths.
+        dataset (str): The dataset name.
+    Returns:
+        List[str]: List of Parquet URIs.
     """
-    return HfFileSystem(dataset, repo_type="dataset", revision=PARQUET_REVISION, token=hf_token)
+    return [f"hf://datasets/{dataset}@{safe_quote(PARQUET_REVISION)}/{path}" for path in paths]
 
 
 def compute_first_rows_response(
@@ -179,11 +191,12 @@ def compute_first_rows_response(
 
     logging.debug(f"Found {len(sources)} parquet files for {dataset=}, {config=}, {split=}: {sources}")
 
-    fs = get_parquet_fs(dataset=dataset, hf_token=hf_token)
+    fs = get_hf_fs(hf_token=hf_token)
+    source_uris = get_hf_parquet_uris(sources, dataset=dataset)
     desc = f"{dataset}/{config}/{split}"
     try:
         parquet_files: List[ParquetFile] = thread_map(
-            partial(ParquetFile, filesystem=fs), sources, desc=desc, unit="pq", disable=True
+            partial(ParquetFile, filesystem=fs), source_uris, desc=desc, unit="pq", disable=True
         )
     except Exception as e:
         raise FileSystemError(f"Could not read the parquet files: {e}") from e


Test patch:
diff --git a/services/api/tests/routes/test_rows.py b/services/api/tests/routes/test_rows.py
--- a/services/api/tests/routes/test_rows.py
+++ b/services/api/tests/routes/test_rows.py
@@ -149,6 +149,10 @@ def indexer(app_config: AppConfig, processing_graph: ProcessingGraph) -> Indexer
     )
 
 
+def mock_get_hf_parquet_uris(paths: List[str], dataset: str) -> List[str]:
+    return paths
+
+
 @pytest.fixture
 def rows_index(
     indexer: Indexer,
@@ -156,15 +160,17 @@ def rows_index(
     ds_sharded_fs: AbstractFileSystem,
     dataset_sharded_with_config_parquet: dict[str, Any],
 ) -> Generator[RowsIndex, None, None]:
-    with patch("api.routes.rows.get_parquet_fs", return_value=ds_sharded_fs):
-        yield indexer.get_rows_index("ds_sharded", "plain_text", "train")
+    with patch("api.routes.rows.get_hf_fs", return_value=ds_sharded_fs):
+        with patch("api.routes.rows.get_hf_parquet_uris", side_effect=mock_get_hf_parquet_uris):
+            yield indexer.get_rows_index("ds_sharded", "plain_text", "train")
 
 
 def test_indexer_get_rows_index(
     indexer: Indexer, ds: Dataset, ds_fs: AbstractFileSystem, dataset_with_config_parquet: dict[str, Any]
 ) -> None:
-    with patch("api.routes.rows.get_parquet_fs", return_value=ds_fs):
-        index = indexer.get_rows_index("ds", "plain_text", "train")
+    with patch("api.routes.rows.get_hf_fs", return_value=ds_fs):
+        with patch("api.routes.rows.get_hf_parquet_uris", side_effect=mock_get_hf_parquet_uris):
+            index = indexer.get_rows_index("ds", "plain_text", "train")
     assert index.features == ds.features
     assert index.row_group_offsets.tolist() == [len(ds)]
     assert len(index.row_group_readers) == 1
@@ -180,8 +186,9 @@ def test_indexer_get_rows_index_sharded(
     ds_sharded_fs: AbstractFileSystem,
     dataset_sharded_with_config_parquet: dict[str, Any],
 ) -> None:
-    with patch("api.routes.rows.get_parquet_fs", return_value=ds_sharded_fs):
-        index = indexer.get_rows_index("ds_sharded", "plain_text", "train")
+    with patch("api.routes.rows.get_hf_fs", return_value=ds_sharded_fs):
+        with patch("api.routes.rows.get_hf_parquet_uris", side_effect=mock_get_hf_parquet_uris):
+            index = indexer.get_rows_index("ds_sharded", "plain_text", "train")
     assert index.features == ds_sharded.features
     assert index.row_group_offsets.tolist() == np.cumsum([len(ds)] * 4).tolist()
     assert len(index.row_group_readers) == 4
diff --git a/services/worker/tests/job_runners/split/test_first_rows_from_parquet.py b/services/worker/tests/job_runners/split/test_first_rows_from_parquet.py
--- a/services/worker/tests/job_runners/split/test_first_rows_from_parquet.py
+++ b/services/worker/tests/job_runners/split/test_first_rows_from_parquet.py
@@ -4,7 +4,7 @@
 import os
 from dataclasses import replace
 from http import HTTPStatus
-from typing import Callable
+from typing import Callable, List
 from unittest.mock import Mock, patch
 
 import pytest
@@ -81,6 +81,10 @@ def test_doesnotexist(app_config: AppConfig, get_job_runner: GetJobRunner) -> No
         get_response(kind=job_runner.processing_step.cache_kind, dataset=dataset, config=config, split=split)
 
 
+def mock_get_hf_parquet_uris(paths: List[str], dataset: str) -> List[str]:
+    return paths
+
+
 @pytest.mark.parametrize(
     "rows_max_bytes,columns_max_number,error_code",
     [
@@ -115,63 +119,69 @@ def test_compute(
         http_status=HTTPStatus.OK,
     )
 
-    with patch("worker.job_runners.split.first_rows_from_parquet.get_parquet_fs") as mock_read:
-        initial_location = os.getcwd()
-        os.chdir("tests/job_runners/split")
-        # TODO:  Make localsystem by relative path
-        fs = LocalFileSystem()
-        mock_read.return_value = fs
-        # ^ Mocking file system with local file
-        job_runner = get_job_runner(
-            dataset,
-            config,
-            split,
-            replace(
-                app_config,
-                common=replace(app_config.common, hf_token=None),
-                first_rows=replace(
-                    app_config.first_rows,
-                    max_number=1_000_000,
-                    min_number=10,
-                    max_bytes=rows_max_bytes,
-                    min_cell_bytes=10,
-                    columns_max_number=columns_max_number,
+    with patch("worker.job_runners.split.first_rows_from_parquet.get_hf_fs") as mock_read:
+        with patch(
+            "worker.job_runners.split.first_rows_from_parquet.get_hf_parquet_uris",
+            side_effect=mock_get_hf_parquet_uris,
+        ):
+            initial_location = os.getcwd()
+            os.chdir("tests/job_runners/split")
+            # TODO:  Make localsystem by relative path
+            fs = LocalFileSystem()
+            mock_read.return_value = fs
+            # ^ Mocking file system with local file
+            job_runner = get_job_runner(
+                dataset,
+                config,
+                split,
+                replace(
+                    app_config,
+                    common=replace(app_config.common, hf_token=None),
+                    first_rows=replace(
+                        app_config.first_rows,
+                        max_number=1_000_000,
+                        min_number=10,
+                        max_bytes=rows_max_bytes,
+                        min_cell_bytes=10,
+                        columns_max_number=columns_max_number,
+                    ),
                 ),
-            ),
-            False,
-        )
-
-        job_runner.get_dataset_git_revision = Mock(return_value="1.0.0")  # type: ignore
-        if error_code:
-            with pytest.raises(CustomError) as error_info:
-                job_runner.compute()
-            assert error_info.value.code == error_code
-        else:
-            response = job_runner.compute().content
-            assert get_json_size(response) <= rows_max_bytes
-            assert response
-            assert response["rows"]
-            assert response["features"]
-            assert len(response["rows"]) == 3  # testing file has 3 rows see config/dataset-split.parquet file
-            assert len(response["features"]) == 2  # testing file has 2 columns see config/dataset-split.parquet file
-            assert response["features"][0]["feature_idx"] == 0
-            assert response["features"][0]["name"] == "col1"
-            assert response["features"][0]["type"]["_type"] == "Value"
-            assert response["features"][0]["type"]["dtype"] == "int32"
-            assert response["features"][1]["feature_idx"] == 1
-            assert response["features"][1]["name"] == "col2"
-            assert response["features"][1]["type"]["_type"] == "Value"
-            assert response["features"][1]["type"]["dtype"] == "string"
-            assert response["rows"][0]["row_idx"] == 0
-            assert response["rows"][0]["truncated_cells"] == []
-            assert response["rows"][0]["row"] == {"col1": 1, "col2": "a"}
-            assert response["rows"][1]["row_idx"] == 1
-            assert response["rows"][1]["truncated_cells"] == []
-            assert response["rows"][1]["row"] == {"col1": 2, "col2": "b"}
-            assert response["rows"][2]["row_idx"] == 2
-            assert response["rows"][2]["truncated_cells"] == []
-            assert response["rows"][2]["row"] == {"col1": 3, "col2": "c"}
-        os.chdir(initial_location)
+                False,
+            )
+
+            job_runner.get_dataset_git_revision = Mock(return_value="1.0.0")  # type: ignore
+            if error_code:
+                with pytest.raises(CustomError) as error_info:
+                    job_runner.compute()
+                assert error_info.value.code == error_code
+            else:
+                response = job_runner.compute().content
+                assert get_json_size(response) <= rows_max_bytes
+                assert response
+                assert response["rows"]
+                assert response["features"]
+                assert len(response["rows"]) == 3  # testing file has 3 rows see config/dataset-split.parquet file
+                assert (
+                    len(response["features"]) == 2
+                )  # testing file has 2 columns see config/dataset-split.parquet file
+                assert response["features"][0]["feature_idx"] == 0
+                assert response["features"][0]["name"] == "col1"
+                assert response["features"][0]["type"]["_type"] == "Value"
+                assert response["features"][0]["type"]["dtype"] == "int32"
+                assert response["features"][1]["feature_idx"] == 1
+                assert response["features"][1]["name"] == "col2"
+                assert response["features"][1]["type"]["_type"] == "Value"
+                assert response["features"][1]["type"]["dtype"] == "string"
+                assert response["rows"][0]["row_idx"] == 0
+                assert response["rows"][0]["truncated_cells"] == []
+                assert response["rows"][0]["row"] == {"col1": 1, "col2": "a"}
+                assert response["rows"][1]["row_idx"] == 1
+                assert response["rows"][1]["truncated_cells"] == []
+                assert response["rows"][1]["row"] == {"col1": 2, "col2": "b"}
+                assert response["rows"][2]["row_idx"] == 2
+                assert response["rows"][2]["truncated_cells"] == []
+                assert response["rows"][2]["row"] == {"col1": 3, "col2": "c"}
+            os.chdir(initial_location)
 
 
 @pytest.mark.parametrize(