Repository: fsspec/filesystem_spec
Pull request #1734 (2024-10-22 20:01:51)

Fix broken async reference file system `_cat_file` method
Adds missing return in `_cat_file` method of `ReferenceFileSystem`

Im not sure how to best test this because i dont think any of the tests from test_reference use an async filesystem and s3 maybe shouldnt be tested here? 

Closes https://github.com/fsspec/filesystem_spec/issues/1735

martindurant: Can we make a test? It will need to run in an async func, with an async backend (http?) - can we guarantee something that won't disappear?
mpiannucci: The async function part is easy, are you ok with pointing to s3 data via http similar to the kerchunk tests? Thats the easiest thing i can think of doing for now 
martindurant: I suppose we can reconsider if it starts to fail. The download should be small though (part of a file only), as the GHA CI runners are pretty weak.
mpiannucci: Okay I have a test added and working well!

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 48e3c95..f1f11e9 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,7 +5,7 @@ exclude: >
 repos:
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -14,7 +14,7 @@ repos:
       - id: check-yaml
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.4
+    rev: v0.6.9
     hooks:
       # Run the linter.
       - id: ruff
diff --git a/fsspec/core.py b/fsspec/core.py
index 1954667..79580be 100644
--- a/fsspec/core.py
+++ b/fsspec/core.py
@@ -329,12 +329,19 @@ def open_files(
 
 
 def _un_chain(path, kwargs):
-    x = re.compile(".*[^a-z]+.*")  # test for non protocol-like single word
-    bits = (
-        [p if "://" in p or x.match(p) else p + "://" for p in path.split("::")]
-        if "::" in path
-        else [path]
-    )
+    # Avoid a circular import
+    from fsspec.implementations.cached import CachingFileSystem
+
+    if "::" in path:
+        x = re.compile(".*[^a-z]+.*")  # test for non protocol-like single word
+        bits = []
+        for p in path.split("::"):
+            if "://" in p or x.match(p):
+                bits.append(p)
+            else:
+                bits.append(p + "://")
+    else:
+        bits = [path]
     # [[url, protocol, kwargs], ...]
     out = []
     previous_bit = None
@@ -351,10 +358,7 @@ def _un_chain(path, kwargs):
             **kws,
         )
         bit = cls._strip_protocol(bit)
-        if (
-            protocol in {"blockcache", "filecache", "simplecache"}
-            and "target_protocol" not in kw
-        ):
+        if "target_protocol" not in kw and issubclass(cls, CachingFileSystem):
             bit = previous_bit
         out.append((bit, protocol, kw))
         previous_bit = bit
diff --git a/fsspec/implementations/reference.py b/fsspec/implementations/reference.py
index d5d3f29..79bff14 100644
--- a/fsspec/implementations/reference.py
+++ b/fsspec/implementations/reference.py
@@ -41,7 +41,7 @@ def _first(d):
 
 def _prot_in_references(path, references):
     ref = references.get(path)
-    if isinstance(ref, (list, tuple)):
+    if isinstance(ref, (list, tuple)) and isinstance(ref[0], str):
         return split_protocol(ref[0])[0] if ref[0] else ref[0]
 
 
@@ -173,8 +173,11 @@ class LazyReferenceMapper(collections.abc.MutableMapping):
             """cached parquet file loader"""
             path = self.url.format(field=field, record=record)
             data = io.BytesIO(self.fs.cat_file(path))
-            df = self.pd.read_parquet(data, engine=self.engine)
-            refs = {c: df[c].to_numpy() for c in df.columns}
+            try:
+                df = self.pd.read_parquet(data, engine=self.engine)
+                refs = {c: df[c].to_numpy() for c in df.columns}
+            except IOError:
+                refs = None
             return refs
 
         self.open_refs = open_refs
@@ -803,7 +806,9 @@ class ReferenceFileSystem(AsyncFileSystem):
             return part_or_url[start:end]
         protocol, _ = split_protocol(part_or_url)
         try:
-            await self.fss[protocol]._cat_file(part_or_url, start=start, end=end)
+            return await self.fss[protocol]._cat_file(
+                part_or_url, start=start0, end=end0
+            )
         except Exception as e:
             raise ReferenceNotReachable(path, part_or_url) from e
 
@@ -871,6 +876,9 @@ class ReferenceFileSystem(AsyncFileSystem):
                 # found and on_error is "raise"
                 try:
                     u, s, e = self._cat_common(p)
+                    if not isinstance(u, (bytes, str)):
+                        # nan/None from parquet
+                        continue
                 except FileNotFoundError as err:
                     if on_error == "raise":
                         raise
diff --git a/fsspec/mapping.py b/fsspec/mapping.py
index 8fb9b9e..752eef3 100644
--- a/fsspec/mapping.py
+++ b/fsspec/mapping.py
@@ -112,7 +112,7 @@ class FSMap(MutableMapping):
             for k, v in out.items()
         }
         return {
-            key: out[k2]
+            key: out[k2] if on_error == "raise" else out.get(k2, KeyError(k2))
             for key, k2 in zip(keys, keys2)
             if on_error == "return" or not isinstance(out[k2], BaseException)
         }


Test patch:
diff --git a/fsspec/implementations/tests/test_cached.py b/fsspec/implementations/tests/test_cached.py
index 67edbc8..c9222d5 100644
--- a/fsspec/implementations/tests/test_cached.py
+++ b/fsspec/implementations/tests/test_cached.py
@@ -267,7 +267,7 @@ def test_blockcache_workflow(ftp_writable, tmp_path, force_save_pickle):
         assert f.read(5) == b"test\n"
 
 
-@pytest.mark.parametrize("impl", ["filecache", "blockcache"])
+@pytest.mark.parametrize("impl", ["filecache", "blockcache", "cached"])
 def test_workflow(ftp_writable, impl):
     host, port, user, pw = ftp_writable
     fs = FTPFileSystem(host, port, user, pw)
@@ -295,7 +295,7 @@ def test_workflow(ftp_writable, impl):
         )  # new value, because we overwrote the cached location
 
 
-@pytest.mark.parametrize("impl", ["simplecache", "blockcache"])
+@pytest.mark.parametrize("impl", ["simplecache", "blockcache", "cached"])
 def test_glob(ftp_writable, impl):
     host, port, user, pw = ftp_writable
     fs = FTPFileSystem(host, port, user, pw)
@@ -622,7 +622,7 @@ def test_metadata_save_blocked(ftp_writable, caplog):
     assert "Cache save failed due to interpreter shutdown" in caplog.text
 
 
-@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache"])
+@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache", "cached"])
 def test_local_filecache_creates_dir_if_needed(impl):
     import tempfile
 
@@ -875,7 +875,7 @@ def test_filecache_with_checks():
     assert fs.cat(f1) == data * 2  # changed, since origin changed
 
 
-@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache"])
+@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache", "cached"])
 @pytest.mark.parametrize("fs", ["local", "multi"], indirect=["fs"])
 def test_filecache_takes_fs_instance(impl, fs):
     origin = tempfile.mkdtemp()
@@ -889,7 +889,7 @@ def test_filecache_takes_fs_instance(impl, fs):
     assert fs2.cat(f1) == data
 
 
-@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache"])
+@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache", "cached"])
 @pytest.mark.parametrize("fs", ["local", "multi"], indirect=["fs"])
 def test_filecache_serialization(impl, fs):
     fs1 = fsspec.filesystem(impl, fs=fs)
@@ -1031,7 +1031,9 @@ def test_multi_cache(protocol):
             assert f.read() == b"hello"
 
 
-@pytest.mark.parametrize("protocol", ["simplecache", "filecache", "blockcache"])
+@pytest.mark.parametrize(
+    "protocol", ["simplecache", "filecache", "blockcache", "cached"]
+)
 def test_multi_cat(protocol, ftp_writable):
     host, port, user, pw = ftp_writable
     fs = FTPFileSystem(host, port, user, pw)
@@ -1064,7 +1066,9 @@ def test_multi_cache_chain(protocol):
         assert files[0].read() == b"hello"
 
 
-@pytest.mark.parametrize("protocol", ["blockcache", "simplecache", "filecache"])
+@pytest.mark.parametrize(
+    "protocol", ["blockcache", "cached", "simplecache", "filecache"]
+)
 def test_strip(protocol):
     fs = fsspec.filesystem(protocol, target_protocol="memory")
     url1 = "memory://afile"
@@ -1235,9 +1239,11 @@ def test_cache_dir_auto_deleted(temp_cache, tmpdir):
         assert local.exists(cache_dir)
 
 
-@pytest.mark.parametrize("protocol", ["filecache", "blockcache", "simplecache"])
+@pytest.mark.parametrize(
+    "protocol", ["filecache", "blockcache", "cached", "simplecache"]
+)
 def test_cache_size(tmpdir, protocol):
-    if win and protocol == "blockcache":
+    if win and protocol in {"blockcache", "cached"}:
         pytest.skip("Windows file locking affects blockcache size tests")
 
     source = os.path.join(tmpdir, "source")
diff --git a/fsspec/implementations/tests/test_reference.py b/fsspec/implementations/tests/test_reference.py
index d82dc17..0063d11 100644
--- a/fsspec/implementations/tests/test_reference.py
+++ b/fsspec/implementations/tests/test_reference.py
@@ -513,6 +513,35 @@ def test_cat_file_ranges(m):
     assert fs.cat_file("d", 1, -3) == other[4:10][1:-3]
 
 
+@pytest.mark.asyncio
+async def test_async_cat_file_ranges():
+    fsspec.get_filesystem_class("http").clear_instance_cache()
+    fss = fsspec.filesystem("https", asynchronous=True)
+    session = await fss.set_session()
+
+    fs = fsspec.filesystem(
+        "reference",
+        fo={
+            "version": 1,
+            "refs": {
+                "reference_time/0": [
+                    "https://noaa-nwm-retro-v2-0-pds.s3.amazonaws.com/full_physics/2017/201704010000.CHRTOUT_DOMAIN1.comp",
+                    39783,
+                    12,
+                ],
+            },
+        },
+        fs={"https": fss},
+        remote_protocol="https",
+        asynchronous=True,
+    )
+
+    assert (
+        await fs._cat_file("reference_time/0") == b"x^K0\xa9d\x04\x00\x03\x13\x01\x0f"
+    )
+    await session.close()
+
+
 @pytest.mark.parametrize(
     "fo",
     [
@@ -792,3 +821,44 @@ def test_deep_parq(m, engine):
         "instant/one/.zarray",
         "instant/one/0",
     ]
+
+
+def test_parquet_no_data(m):
+    zarr = pytest.importorskip("zarr")
+    lz = fsspec.implementations.reference.LazyReferenceMapper.create(
+        "memory://out.parq", fs=m
+    )
+
+    g = zarr.open_group(lz, mode="w")
+    arr = g.create_dataset(
+        name="one",
+        dtype="int32",
+        shape=(10,),
+        chunks=(5,),
+        compression=None,
+        fill_value=1,
+    )
+    lz.flush()
+
+    assert (arr[:] == 1).all()
+
+
+def test_parquet_no_references(m):
+    zarr = pytest.importorskip("zarr")
+    lz = fsspec.implementations.reference.LazyReferenceMapper.create(
+        "memory://out.parq", fs=m
+    )
+
+    g = zarr.open_group(lz, mode="w")
+    arr = g.create_dataset(
+        name="one",
+        dtype="int32",
+        shape=(),
+        chunks=(),
+        compression=None,
+        fill_value=1,
+    )
+    lz.flush()
+    arr[...]
+
+    assert arr[...].tolist() == 1  #  scalar, equal to fill value