fix: ensure tools can load datasets that are exactly 4kiB

.pre-commit-config.yaml

@@ -1,5 +1,7 @@
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
+exclude: |
+  (?x)^(tests/data)
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0
@@ -10,7 +12,7 @@ repos:
       - id: check-yaml
       - id: check-added-large-files
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.0
+    rev: v0.9.9
     hooks:
       - id: ruff
         name: ruff check
@@ -24,7 +26,7 @@ repos:
         name: ruff format
         alias: format
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.391
+    rev: v1.1.396
     hooks:
       - id: pyright
         additional_dependencies: [cython, httpretty, numpy, pytest, setuptools]

cryosparc/constants.py

@@ -0,0 +1,4 @@
+ONE_MIB = 2**20
+"""
+Bytes in 1 mebibyte
+"""

cryosparc/dataset.py

@@ -51,6 +51,7 @@
 from typing_extensions import Literal, SupportsIndex
 
 from .column import Column
+from .constants import ONE_MIB
 from .core import Data, DsetType, Stream
 from .dtype import (
     NEVER_COMPRESS_FIELDS,
@@ -632,8 +633,12 @@ def _load_numpy(
     ):
         import os
 
-        # disable mmap by setting CRYOSPARC_DATASET_MMAP=false
-        if os.getenv("CRYOSPARC_DATASET_MMAP", "true").lower() == "true" and isinstance(file, (str, PurePath)):
+        # disable mmap by setting CRYOSPARC_DATASET_MMAP=false or dataset is small
+        if (
+            os.getenv("CRYOSPARC_DATASET_MMAP", "true").lower() == "true"
+            and isinstance(file, (str, PurePath))
+            and os.stat(file).st_size > ONE_MIB
+        ):
             # Use mmap to avoid loading full record array into memory
             # cast path to a string for older numpy/python
             mmap_mode, f = "r", str(file)

cryosparc/job.py

@@ -1272,8 +1272,6 @@ def add_output(
         slots: List[SlotSpec] = ...,
         passthrough: Optional[str] = ...,
         title: Optional[str] = ...,
-        *,
-        alloc: Literal[None] = None,
     ) -> str: ...
     @overload
     def add_output(
@@ -1284,7 +1282,7 @@ def add_output(
         passthrough: Optional[str] = ...,
         title: Optional[str] = ...,
         *,
-        alloc: Union[int, Dataset] = ...,
+        alloc: Union[int, Dataset],
     ) -> Dataset: ...
     def add_output(
         self,

cryosparc/tools.py

@@ -36,6 +36,7 @@
 
 from . import __version__, mrc
 from .command import CommandClient, CommandError, make_json_request, make_request
+from .constants import ONE_MIB
 from .dataset import DEFAULT_FORMAT, Dataset
 from .job import ExternalJob, Job
 from .project import Project
@@ -53,8 +54,6 @@
 from .util import bopen, noopcontext, padarray, print_table, trimarray
 from .workspace import Workspace
 
-ONE_MIB = 2**20  # bytes in one mebibyte
-
 LICENSE_REGEX = re.compile(r"[a-f\d]{8}-[a-f\d]{4}-[a-f\d]{4}-[a-f\d]{4}-[a-f\d]{12}")
 """Regular expression for matching CryoSPARC license IDs."""
 

tests/test_dataset.py

@@ -394,3 +394,10 @@ def test_allocate_many_together():
             allocated.append(Dataset(1))
         assert len(allocated) == 66_000
         del allocated
+
+
+def test_load_4k():
+    # Check that a 4kiB dataset (same as linux page size) loads correctly
+    # (numpy bug https://github.com/numpy/numpy/pull/27723 fixed in 2.2)
+    d = Dataset.load("tests/data/4k_dataset.cs")
+    assert len(d) == 0