feat: Add mmap_mode parameter to NpyRef.load()

Support memory-mapped loading for large arrays:
- Local filesystem stores: mmap directly, no download
- Remote stores: download to cache, then mmap

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: dimitri-yatsenko

src/datajoint/builtin_codecs.py

@@ -940,28 +940,77 @@ def is_loaded(self) -> bool:
         """True if array data has been downloaded and cached."""
         return self._cached is not None
 
-    def load(self):
+    def load(self, mmap_mode=None):
         """
         Download and return the array.
 
+        Parameters
+        ----------
+        mmap_mode : str, optional
+            Memory-map mode for lazy, random-access loading of large arrays:
+
+            - ``'r'``: Read-only
+            - ``'r+'``: Read-write
+            - ``'c'``: Copy-on-write (changes not saved to disk)
+
+            If None (default), loads entire array into memory.
+
         Returns
         -------
-        numpy.ndarray
-            The array data.
+        numpy.ndarray or numpy.memmap
+            The array data. Returns ``numpy.memmap`` if mmap_mode is specified.
 
         Notes
         -----
-        The array is cached after first load. Subsequent calls return
-        the cached copy without additional I/O.
+        When ``mmap_mode`` is None, the array is cached after first load.
+
+        For local filesystem stores, memory mapping accesses the file directly
+        with no download. For remote stores (S3, etc.), the file is downloaded
+        to a local cache (``{tempdir}/datajoint_mmap/``) before memory mapping.
+
+        Examples
+        --------
+        Standard loading::
+
+            arr = ref.load()  # Loads entire array into memory
+
+        Memory-mapped for random access to large arrays::
+
+            arr = ref.load(mmap_mode='r')
+            slice = arr[1000:2000]  # Only reads the needed portion from disk
         """
         import io
 
         import numpy as np
 
-        if self._cached is None:
-            buffer = self._backend.get_buffer(self.path)
-            self._cached = np.load(io.BytesIO(buffer), allow_pickle=False)
-        return self._cached
+        if mmap_mode is None:
+            # Standard loading with caching
+            if self._cached is None:
+                buffer = self._backend.get_buffer(self.path)
+                self._cached = np.load(io.BytesIO(buffer), allow_pickle=False)
+            return self._cached
+        else:
+            # Memory-mapped loading
+            if self._backend.protocol == "file":
+                # Local filesystem - mmap directly, no download needed
+                local_path = self._backend._full_path(self.path)
+                return np.load(local_path, mmap_mode=mmap_mode, allow_pickle=False)
+            else:
+                # Remote storage - download to local cache first
+                import hashlib
+                import tempfile
+                from pathlib import Path
+
+                path_hash = hashlib.md5(self.path.encode()).hexdigest()[:12]
+                cache_dir = Path(tempfile.gettempdir()) / "datajoint_mmap"
+                cache_dir.mkdir(exist_ok=True)
+                cache_path = cache_dir / f"{path_hash}.npy"
+
+                if not cache_path.exists():
+                    buffer = self._backend.get_buffer(self.path)
+                    cache_path.write_bytes(buffer)
+
+                return np.load(str(cache_path), mmap_mode=mmap_mode, allow_pickle=False)
 
     def __array__(self, dtype=None):
         """

tests/integration/test_npy_codec.py

@@ -139,6 +139,74 @@ class MockBackend:
         with pytest.raises(TypeError, match="0-dimensional"):
             len(ref)
 
+    def test_npy_ref_mmap_local_filesystem(self, tmp_path):
+        """NpyRef mmap_mode should work directly on local filesystem."""
+        # Create a real .npy file
+        test_array = np.arange(100, dtype=np.float64)
+        npy_path = tmp_path / "test.npy"
+        np.save(npy_path, test_array)
+
+        metadata = {
+            "path": "test.npy",
+            "store": None,
+            "dtype": "float64",
+            "shape": [100],
+        }
+
+        # Mock backend that simulates local filesystem
+        class MockFileBackend:
+            protocol = "file"
+
+            def _full_path(self, path):
+                return str(tmp_path / path)
+
+            def get_buffer(self, path):
+                return (tmp_path / path).read_bytes()
+
+        ref = NpyRef(metadata, MockFileBackend())
+
+        # Load with mmap_mode
+        mmap_arr = ref.load(mmap_mode="r")
+
+        # Should be a memmap
+        assert isinstance(mmap_arr, np.memmap)
+        np.testing.assert_array_equal(mmap_arr, test_array)
+
+        # Standard load should still work and cache
+        regular_arr = ref.load()
+        assert isinstance(regular_arr, np.ndarray)
+        assert not isinstance(regular_arr, np.memmap)
+        np.testing.assert_array_equal(regular_arr, test_array)
+
+    def test_npy_ref_mmap_remote_storage(self, tmp_path):
+        """NpyRef mmap_mode should download to cache for remote storage."""
+        # Create test data
+        test_array = np.array([1, 2, 3, 4, 5], dtype=np.int32)
+        npy_buffer = np.save(tmp_path / "temp.npy", test_array)
+        npy_bytes = (tmp_path / "temp.npy").read_bytes()
+
+        metadata = {
+            "path": "remote/path/data.npy",
+            "store": "s3-store",
+            "dtype": "int32",
+            "shape": [5],
+        }
+
+        # Mock backend that simulates remote storage
+        class MockS3Backend:
+            protocol = "s3"
+
+            def get_buffer(self, path):
+                return npy_bytes
+
+        ref = NpyRef(metadata, MockS3Backend())
+
+        # Load with mmap_mode - should download to cache
+        mmap_arr = ref.load(mmap_mode="r")
+
+        assert isinstance(mmap_arr, np.memmap)
+        np.testing.assert_array_equal(mmap_arr, test_array)
+
 
 class TestNpyCodecUnit:
     """Unit tests for NpyCodec without database."""