Repository: iterative/datachain

Issue:
Support for `hf://` filesystem
> * I was wondering if it works to write the transformed dataset using fsspec and `hf://` paths as well ?

It seems like we either need to use the hf `load_dataset()` api, which makes it easy to apply all the high-level hf logic to the dataset, or use `hf://`, which allows low-level access to the files but loses all of the hf loading logic.

Again, I think we can merge and in next steps add read/write functionality with the `hf://` filesystem to enable syntax like `from_storage("hf://...").parse_hf()` or `from_storage("hf://...").parse_tabular()` to parse files like parquet directly from hf hub and eventually to export to hf hub.


_Originally posted by @dberenbaum in https://github.com/iterative/datachain/issues/311#issuecomment-2313106909_

Discussion:
dberenbaum: Not sure if there is a good place or way to add tests for the `HfClient`.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/iterative/datachain/pull/375?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) Report
Attention: Patch coverage is `80.85106%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 86.34%. Comparing base [(`13f6982`)](https://app.codecov.io/gh/iterative/datachain/commit/13f69827096213a35f2f78f0257292d67c174934?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) to head [(`6cb17e9`)](https://app.codecov.io/gh/iterative/datachain/commit/6cb17e95319f08dc01edf4b92eefccbc70ff250f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).

| [Files with missing lines](https://app.codecov.io/gh/iterative/datachain/pull/375?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) | Patch % | Lines |
|---|---|---|
| [src/datachain/client/hf.py](https://app.codecov.io/gh/iterative/datachain/pull/375?src=pr&el=tree&filepath=src%2Fdatachain%2Fclient%2Fhf.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#diff-c3JjL2RhdGFjaGFpbi9jbGllbnQvaGYucHk=) | 70.83% | [7 Missing :warning: ](https://app.codecov.io/gh/iterative/datachain/pull/375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) |
| [src/datachain/client/fsspec.py](https://app.codecov.io/gh/iterative/datachain/pull/375?src=pr&el=tree&filepath=src%2Fdatachain%2Fclient%2Ffsspec.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#diff-c3JjL2RhdGFjaGFpbi9jbGllbnQvZnNzcGVjLnB5) | 66.66% | [1 Missing :warning: ](https://app.codecov.io/gh/iterative/datachain/pull/375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) |
| [src/datachain/lib/arrow.py](https://app.codecov.io/gh/iterative/datachain/pull/375?src=pr&el=tree&filepath=src%2Fdatachain%2Flib%2Farrow.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#diff-c3JjL2RhdGFjaGFpbi9saWIvYXJyb3cucHk=) | 95.00% | [0 Missing and 1 partial :warning: ](https://app.codecov.io/gh/iterative/datachain/pull/375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #375      +/-   ##
==========================================
- Coverage   86.71%   86.34%   -0.37%     
==========================================
  Files          91       92       +1     
  Lines       10085    10127      +42     
  Branches     2049     2057       +8     
==========================================
- Hits         8745     8744       -1     
- Misses        985     1024      +39     
- Partials      355      359       +4     
```

| [Flag](https://app.codecov.io/gh/iterative/datachain/pull/375/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) | Coverage Δ | |
|---|---|---|
| [datachain](https://app.codecov.io/gh/iterative/datachain/pull/375/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative) | `86.31% <80.85%> (-0.33%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iterative/datachain/pull/375?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iterative).

shcheklein: Amazing! 

Coming back to the discussion we had yesterday - I wonder how easy it is to treat blobs in the `from_hf` version as `File` objects (similar to tar archives).
dberenbaum: Looking into it @shcheklein. Opened https://github.com/iterative/datachain/issues/369 to track that.

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 13b1e64..7b77663 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,8 @@ dependencies = [
   "datamodel-code-generator>=0.25",
   "Pillow>=10.0.0,<11",
   "msgpack>=1.0.4,<2",
-  "psutil"
+  "psutil",
+  "huggingface_hub"
 ]
 
 [project.optional-dependencies]
diff --git a/src/datachain/client/fsspec.py b/src/datachain/client/fsspec.py
index 298fa74..3a7c733 100644
--- a/src/datachain/client/fsspec.py
+++ b/src/datachain/client/fsspec.py
@@ -87,6 +87,7 @@ class Client(ABC):
     def get_implementation(url: str) -> type["Client"]:
         from .azure import AzureClient
         from .gcs import GCSClient
+        from .hf import HfClient
         from .local import FileClient
         from .s3 import ClientS3
 
@@ -104,6 +105,8 @@ class Client(ABC):
             return AzureClient
         if protocol == FileClient.protocol:
             return FileClient
+        if protocol == HfClient.protocol:
+            return HfClient
 
         raise NotImplementedError(f"Unsupported protocol: {protocol}")
 
diff --git a/src/datachain/client/hf.py b/src/datachain/client/hf.py
new file mode 100644
index 0000000..01c84d9
--- /dev/null
+++ b/src/datachain/client/hf.py
@@ -0,0 +1,47 @@
+import os
+import posixpath
+from typing import Any, cast
+
+from huggingface_hub import HfFileSystem
+
+from datachain.lib.file import File
+from datachain.node import Entry
+
+from .fsspec import Client
+
+
+class HfClient(Client):
+    FS_CLASS = HfFileSystem
+    PREFIX = "hf://"
+    protocol = "hf"
+
+    @classmethod
+    def create_fs(cls, **kwargs) -> HfFileSystem:
+        if os.environ.get("HF_TOKEN"):
+            kwargs["token"] = os.environ["HF_TOKEN"]
+
+        return cast(HfFileSystem, super().create_fs(**kwargs))
+
+    def convert_info(self, v: dict[str, Any], path: str) -> Entry:
+        return Entry.from_file(
+            path=path,
+            size=v["size"],
+            version=v["last_commit"].oid,
+            etag=v.get("blob_id", ""),
+            last_modified=v["last_commit"].date,
+        )
+
+    def info_to_file(self, v: dict[str, Any], path: str) -> File:
+        return File(
+            path=path,
+            size=v["size"],
+            version=v["last_commit"].oid,
+            etag=v.get("blob_id", ""),
+            last_modified=v["last_commit"].date,
+        )
+
+    async def ls_dir(self, path):
+        return self.fs.ls(path, detail=True)
+
+    def rel_path(self, path):
+        return posixpath.relpath(path, self.name)
diff --git a/src/datachain/lib/arrow.py b/src/datachain/lib/arrow.py
index 9ac8a8c..5a845d6 100644
--- a/src/datachain/lib/arrow.py
+++ b/src/datachain/lib/arrow.py
@@ -7,7 +7,9 @@ import pyarrow as pa
 from pyarrow.dataset import dataset
 from tqdm import tqdm
 
+from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import File, IndexedFile
+from datachain.lib.model_store import ModelStore
 from datachain.lib.udf import Generator
 
 if TYPE_CHECKING:
@@ -59,7 +61,13 @@ class ArrowGenerator(Generator):
                     vals = list(record.values())
                     if self.output_schema:
                         fields = self.output_schema.model_fields
-                        vals = [self.output_schema(**dict(zip(fields, vals)))]
+                        vals_dict = {}
+                        for (field, field_info), val in zip(fields.items(), vals):
+                            if ModelStore.is_pydantic(field_info.annotation):
+                                vals_dict[field] = field_info.annotation(**val)  # type: ignore[misc]
+                            else:
+                                vals_dict[field] = val
+                        vals = [self.output_schema(**vals_dict)]
                     if self.source:
                         yield [IndexedFile(file=file, index=index), *vals]
                     else:
@@ -95,15 +103,15 @@ def schema_to_output(schema: pa.Schema, col_names: Optional[Sequence[str]] = Non
         if not column:
             column = f"c{default_column}"
             default_column += 1
-        dtype = arrow_type_mapper(field.type)  # type: ignore[assignment]
-        if field.nullable:
+        dtype = arrow_type_mapper(field.type, column)  # type: ignore[assignment]
+        if field.nullable and not ModelStore.is_pydantic(dtype):
             dtype = Optional[dtype]  # type: ignore[assignment]
         output[column] = dtype
 
     return output
 
 
-def arrow_type_mapper(col_type: pa.DataType) -> type:  # noqa: PLR0911
+def arrow_type_mapper(col_type: pa.DataType, column: str = "") -> type:  # noqa: PLR0911
     """Convert pyarrow types to basic types."""
     from datetime import datetime
 
@@ -123,7 +131,15 @@ def arrow_type_mapper(col_type: pa.DataType) -> type:  # noqa: PLR0911
         return str
     if pa.types.is_list(col_type):
         return list[arrow_type_mapper(col_type.value_type)]  # type: ignore[return-value, misc]
-    if pa.types.is_struct(col_type) or pa.types.is_map(col_type):
+    if pa.types.is_struct(col_type):
+        type_dict = {}
+        for field in col_type:
+            dtype = arrow_type_mapper(field.type, field.name)
+            if field.nullable and not ModelStore.is_pydantic(dtype):
+                dtype = Optional[dtype]  # type: ignore[assignment]
+            type_dict[field.name] = dtype
+        return dict_to_data_model(column, type_dict)
+    if pa.types.is_map(col_type):
         return dict
     if isinstance(col_type, pa.lib.DictionaryType):
         return arrow_type_mapper(col_type.value_type)  # type: ignore[return-value]


Test patch:
diff --git a/tests/unit/lib/test_arrow.py b/tests/unit/lib/test_arrow.py
index f84007c..e9693c2 100644
--- a/tests/unit/lib/test_arrow.py
+++ b/tests/unit/lib/test_arrow.py
@@ -3,6 +3,7 @@ from typing import Optional
 
 import pandas as pd
 import pyarrow as pa
+import pyarrow.parquet as pq
 import pytest
 
 from datachain.lib.arrow import (
@@ -10,6 +11,7 @@ from datachain.lib.arrow import (
     arrow_type_mapper,
     schema_to_output,
 )
+from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import File, IndexedFile
 
 
@@ -55,6 +57,34 @@ def test_arrow_generator_no_source(tmp_path, catalog):
         assert o[1] == text
 
 
+def test_arrow_generator_output_schema(tmp_path, catalog):
+    ids = [12345, 67890, 34, 0xF0123]
+    texts = ["28", "22", "we", "hello world"]
+    dicts = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}, {"a": 7, "b": 8}]
+    df = pd.DataFrame({"id": ids, "text": texts, "dict": dicts})
+    table = pa.Table.from_pandas(df)
+
+    name = "111.parquet"
+    pq_path = tmp_path / name
+    pq.write_table(table, pq_path)
+    stream = File(path=pq_path.as_posix(), source="file:///")
+    stream._set_stream(catalog, caching_enabled=False)
+
+    output_schema = dict_to_data_model("", schema_to_output(table.schema))
+    func = ArrowGenerator(output_schema=output_schema)
+    objs = list(func.process(stream))
+
+    assert len(objs) == len(ids)
+    for index, (o, id, text, dict) in enumerate(zip(objs, ids, texts, dicts)):
+        assert isinstance(o[0], IndexedFile)
+        assert isinstance(o[0].file, File)
+        assert o[0].index == index
+        assert o[1].id == id
+        assert o[1].text == text
+        assert o[1].dict.a == dict["a"]
+        assert o[1].dict.b == dict["b"]
+
+
 @pytest.mark.parametrize(
     "col_type,expected",
     (
@@ -72,7 +102,6 @@ def test_arrow_generator_no_source(tmp_path, catalog):
         (pa.date32(), datetime),
         (pa.string(), str),
         (pa.large_string(), str),
-        (pa.struct({"x": pa.int32(), "y": pa.string()}), dict),
         (pa.map_(pa.string(), pa.int32()), dict),
         (pa.dictionary(pa.int64(), pa.string()), str),
         (pa.list_(pa.string()), list[str]),
@@ -82,6 +111,14 @@ def test_arrow_type_mapper(col_type, expected):
     assert arrow_type_mapper(col_type) == expected
 
 
+def test_arrow_type_mapper_struct():
+    col_type = pa.struct({"x": pa.int32(), "y": pa.string()})
+    fields = arrow_type_mapper(col_type).model_fields
+    assert list(fields.keys()) == ["x", "y"]
+    dtypes = [field.annotation for field in fields.values()]
+    assert dtypes == [Optional[int], Optional[str]]
+
+
 def test_arrow_type_error():
     col_type = pa.union(
         [pa.field("a", pa.binary(10)), pa.field("b", pa.string())],