Repository: fsspec/filesystem_spec
Pull request #1733 (2024-10-22 19:18:55)

Consolidate block fetch requests.
Implement the TODO to consolidate consecutive block fetch requests into a single fetch request.

Adapted from the itertools example found here: https://docs.python.org/2.6/library/itertools.html#examples

martindurant: Would you mind explaining more completely what this does? It also makes that block of code look more complex (to me), so some comments inline would be appreciated.
groutr: Python's groupby works by splitting based on a change in the value of a key function. The chosen key computes the difference between an ascending range and the block numbers (which are also necessarily sorted). If the that difference changes, we know that there was a "hole" in the block numbers (ie, we already have the block cached).

Assume a file of 5 blocks, `[0, 1, 2, 3, 4]` where we have already cached block 2. The caller requests a byte range that covers  blocks 0, 1, 3.
```
need = (i for i in block_range if i not in self.blocks)   # lazily evaluated to [0, 1, 3]
groupby(enumerate(need), key=lambda x: x[0]-x[1])
    Enumerating need pairs the needed blocks with an ascending range
    (0, 0), (1, 1), (2, 3)  # Note that the actual block numbers are the second values.
    Applying the key we get values: 0, 0, 1
    Groupby splits a new group on each key value change so we end up with [(0, 0), (1, 1)] and [(2, 3)] groups.
Since we enumerated need, we need to extract the second values from each group to get our blocks, 
so we map itemgetter(1) over the group.
For the first group, we have blocks [0, 1], so we compute the offset for beginning of block 0 and the end of block 1.
One call is made to fetch for both blocks 0 and 1.

The second group is just a single block, 3. One call is made to fetch for that block.
```

I'll add an explanation via comments.

martindurant: You might be interested, that something similar happens in https://github.com/fsspec/filesystem_spec/blob/master/fsspec/utils.py#L528 (where we don't have blocks, but we do have byte ranges and possible gaps). 
groutr: @martindurant Thank you for pointing my attention to that function. I'm not sure if the approach here is directly applicable there. I'll have to spend more time reasoning thru it.

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/caching.py b/fsspec/caching.py
index a3f7a1c..bc74ad2 100644
--- a/fsspec/caching.py
+++ b/fsspec/caching.py
@@ -7,6 +7,8 @@ import math
 import os
 import threading
 import warnings
+from itertools import groupby
+from operator import itemgetter
 from concurrent.futures import Future, ThreadPoolExecutor
 from typing import (
     TYPE_CHECKING,
@@ -161,21 +163,39 @@ class MMapCache(BaseCache):
             return b""
         start_block = start // self.blocksize
         end_block = end // self.blocksize
-        need = [i for i in range(start_block, end_block + 1) if i not in self.blocks]
-        hits = [i for i in range(start_block, end_block + 1) if i in self.blocks]
-        self.miss_count += len(need)
-        self.hit_count += len(hits)
-        while need:
-            # TODO: not a for loop so we can consolidate blocks later to
-            # make fewer fetch calls; this could be parallel
-            i = need.pop(0)
-
-            sstart = i * self.blocksize
-            send = min(sstart + self.blocksize, self.size)
+        block_range = range(start_block, end_block + 1)
+        # Determine which blocks need to be fetched. This sequence is sorted by construction.
+        need = (i for i in block_range if i not in self.blocks)
+        # Count the number of blocks already cached
+        self.hit_count += sum(1 for i in block_range if i in self.blocks)
+
+        # Consolidate needed blocks.
+        # Algorithm adapted from Python 2.x itertools documentation.
+        # We are grouping an enumerated sequence of blocks. By comparing when the difference
+        # between an ascending range (provided by enumerate) and the needed block numbers
+        # we can detect when the block number skips values. The key computes this difference.
+        # Whenever the difference changes, we know that we have previously cached block(s),
+        # and a new group is started. In other words, this algorithm neatly groups
+        # runs of consecutive block numbers so they can be fetched together.
+        for _, _blocks in groupby(enumerate(need), key=lambda x: x[0] - x[1]):
+            # Extract the blocks from the enumerated sequence
+            _blocks = tuple(map(itemgetter(1), _blocks))
+            # Compute start of first block
+            sstart = _blocks[0] * self.blocksize
+            # Compute the end of the last block. Last block may not be full size.
+            send = min(_blocks[-1] * self.blocksize + self.blocksize, self.size)
+
+            # Fetch bytes (could be multiple consecutive blocks)
             self.total_requested_bytes += send - sstart
-            logger.debug(f"MMap get block #{i} ({sstart}-{send})")
+            logger.debug(
+                f"MMap get blocks {_blocks[0]}-{_blocks[-1]} ({sstart}-{send})"
+            )
             self.cache[sstart:send] = self.fetcher(sstart, send)
-            self.blocks.add(i)
+
+            # Update set of cached blocks
+            self.blocks.update(_blocks)
+            # Update cache statistics with number of blocks we had to cache
+            self.miss_count += len(_blocks)
 
         return self.cache[start:end]
 
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)
         }
diff --git a/fsspec/registry.py b/fsspec/registry.py
index c261b9b..02094ea 100644
--- a/fsspec/registry.py
+++ b/fsspec/registry.py
@@ -202,6 +202,10 @@ known_implementations = {
         "err": 'SFTPFileSystem requires "paramiko" to be installed',
     },
     "tar": {"class": "fsspec.implementations.tar.TarFileSystem"},
+    "tosfs": {
+        "class": "tosfs.TosFileSystem",
+        "err": "Install tosfs to access ByteDance volcano engine Tinder Object Storage",
+    },
     "wandb": {"class": "wandbfs.WandbFS", "err": "Install wandbfs to access wandb"},
     "webdav": {
         "class": "webdav4.fsspec.WebdavFileSystem",
diff --git a/fsspec/spec.py b/fsspec/spec.py
index 8284366..6ce338b 100644
--- a/fsspec/spec.py
+++ b/fsspec/spec.py
@@ -408,7 +408,7 @@ class AbstractFileSystem(metaclass=_Cached):
         topdown: bool (True)
             Whether to walk the directory tree from the top downwards or from
             the bottom upwards.
-        on_error: "omit", "raise", a collable
+        on_error: "omit", "raise", a callable
             if omit (default), path with exception will simply be empty;
             If raise, an underlying exception will be raised;
             if callable, it will be called with a single OSError instance as argument


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