Read arrow files from cache

src/datachain/lib/arrow.py

@@ -46,7 +46,10 @@ def __init__(
         self.kwargs = kwargs
 
     def process(self, file: File):
-        if self.nrows:
+        if file._caching_enabled:
+            path = file.get_local_path(download=True)
+            ds = dataset(path, schema=self.input_schema, **self.kwargs)
+        elif self.nrows:
             path = _nrows_file(file, self.nrows)
             ds = dataset(path, schema=self.input_schema, **self.kwargs)
         else:

src/datachain/lib/file.py

@@ -257,14 +257,18 @@ def get_uid(self) -> UniqueId:
         dump = self.model_dump()
         return UniqueId(*(dump[k] for k in self._unique_id_keys))
 
-    def get_local_path(self) -> Optional[str]:
+    def get_local_path(self, download: bool = False) -> Optional[str]:
         """Returns path to a file in a local cache.
         Return None if file is not cached. Throws an exception if cache is not setup."""
         if self._catalog is None:
             raise RuntimeError(
                 "cannot resolve local file path because catalog is not setup"
             )
-        return self._catalog.cache.get_path(self.get_uid())
+        uid = self.get_uid()
+        if download:
+            client = self._catalog.get_client(self.source)
+            client.download(uid, callback=self._download_cb)
+        return self._catalog.cache.get_path(uid)
 
     def get_file_suffix(self):
         """Returns last part of file name with `.`."""

tests/unit/lib/test_arrow.py

@@ -15,16 +15,17 @@
 from datachain.lib.file import File, IndexedFile
 
 
-def test_arrow_generator(tmp_path, catalog):
+@pytest.mark.parametrize("cache", [True, False])
+def test_arrow_generator(tmp_path, catalog, cache):
     ids = [12345, 67890, 34, 0xF0123]
     texts = ["28", "22", "we", "hello world"]
     df = pd.DataFrame({"id": ids, "text": texts})
 
     name = "111.parquet"
     pq_path = tmp_path / name
     df.to_parquet(pq_path)
-    stream = File(path=pq_path.as_posix(), source="file:///")
-    stream._set_stream(catalog, caching_enabled=False)
+    stream = File(path=pq_path.as_posix(), source="file://")
+    stream._set_stream(catalog, caching_enabled=cache)
 
     func = ArrowGenerator()
     objs = list(func.process(stream))
@@ -46,7 +47,7 @@ def test_arrow_generator_no_source(tmp_path, catalog):
     name = "111.parquet"
     pq_path = tmp_path / name
     df.to_parquet(pq_path)
-    stream = File(path=pq_path.as_posix(), source="file:///")
+    stream = File(path=pq_path.as_posix(), source="file://")
     stream._set_stream(catalog, caching_enabled=False)
 
     func = ArrowGenerator(source=False)
@@ -67,7 +68,7 @@ def test_arrow_generator_output_schema(tmp_path, catalog):
     name = "111.parquet"
     pq_path = tmp_path / name
     pq.write_table(table, pq_path)
-    stream = File(path=pq_path.as_posix(), source="file:///")
+    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))

tests/unit/lib/test_file.py

@@ -387,3 +387,18 @@ def test_resolve_function():
 
     assert result == "resolved_file"
     mock_file.resolve.assert_called_once()
+
+
+def test_get_local_path(tmp_path, catalog):
+    file_name = "myfile"
+    data = b"some\x00data\x00is\x48\x65\x6c\x57\x6f\x72\x6c\x64\xff\xffheRe"
+
+    file_path = tmp_path / file_name
+    with open(file_path, "wb") as fd:
+        fd.write(data)
+
+    file = File(path=file_name, source=f"file://{tmp_path}")
+    file._set_stream(catalog)
+
+    assert file.get_local_path() is None
+    assert file.get_local_path(download=True) is not None