refactor(referencefs): only override ReferenceFileSystem._open for fsspec < 2025.2.0 (#884)

Author: skshetry

pyproject.toml

@@ -24,10 +24,12 @@ dependencies = [
   "tqdm",
   "numpy>=1,<3",
   "pandas>=2.0.0",
+  "packaging",
   "pyarrow",
   "typing-extensions",
   "python-dateutil>=2",
   "attrs>=21.3.0",
+  "fsspec>=2024.2.0",
   "s3fs>=2024.2.0",
   "gcsfs>=2024.2.0",
   "adlfs>=2024.2.0",

src/datachain/fs/__init__.py

@@ -0,0 +1,21 @@
+import fsspec
+from packaging.version import Version, parse
+
+# fsspec==2025.2.0 added support for a proper `open()` in `ReferenceFileSystem`.
+# Remove this module when `fsspec` minimum version requirement can be bumped.
+if parse(fsspec.__version__) < Version("2025.2.0"):
+    from fsspec.core import split_protocol
+    from fsspec.implementations import reference
+
+    class ReferenceFileSystem(reference.ReferenceFileSystem):
+        def _open(self, path, mode="rb", *args, **kwargs):
+            # overriding because `fsspec`'s `ReferenceFileSystem._open`
+            # reads the whole file in-memory.
+            (uri,) = self.references[path]
+            protocol, _ = split_protocol(uri)
+            return self.fss[protocol].open(uri, mode, *args, **kwargs)
+else:
+    from fsspec.implementations.reference import ReferenceFileSystem  # type: ignore[no-redef]  # noqa: I001
+
+
+__all__ = ["ReferenceFileSystem"]

src/datachain/fs/reference.py

@@ -2,13 +2,12 @@
 from itertools import islice
 from typing import TYPE_CHECKING, Any, Optional
 
-import fsspec.implementations.reference
 import orjson
 import pyarrow as pa
-from fsspec.core import split_protocol
 from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm.auto import tqdm
 
+from datachain.fs.reference import ReferenceFileSystem
 from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import ArrowRow, File
 from datachain.lib.model_store import ModelStore
@@ -27,15 +26,6 @@
 DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY = b"DataChain SignalSchema"
 
 
-class ReferenceFileSystem(fsspec.implementations.reference.ReferenceFileSystem):
-    def _open(self, path, mode="rb", *args, **kwargs):
-        # overriding because `fsspec`'s `ReferenceFileSystem._open`
-        # reads the whole file in-memory.
-        (uri,) = self.references[path]
-        protocol, _ = split_protocol(uri)
-        return self.fss[protocol].open(uri, mode, *args, **kwargs)
-
-
 class ArrowGenerator(Generator):
     DEFAULT_BATCH_SIZE = 2**17  # same as `pyarrow._dataset._DEFAULT_BATCH_SIZE`