Skip to content

Commit

Permalink
Add a safer mmap reading strategy
Browse files Browse the repository at this point in the history
  • Loading branch information
birchkwok committed Sep 23, 2024
1 parent ae6f1eb commit 573a169
Showing 1 changed file with 10 additions and 8 deletions.
18 changes: 10 additions & 8 deletions lynse/utils/utils.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import re
import time
from functools import wraps
Expand Down Expand Up @@ -221,14 +222,15 @@ def safe_mmap_reader(path, ids=None):
ids (list): The slices to read from the file.
Returns:
np.ndarray: The numpy ndarray.
np.ndarray or np.memmap: If the system is Windows, the file will be directly loaded into memory.
Otherwise, the file will be memory-mapped.
"""
mmap_file = np.load(path, "r")
if ids is None:
array = np.asarray(memoryview(mmap_file))
if os.name == 'nt':
mmap_mode = None
else:
array = np.asarray(memoryview(mmap_file[ids]))
mmap_mode = 'r'

if ids is None:
return np.load(path, mmap_mode=mmap_mode)

# Close the mmap file
mmap_file._mmap.close()
return array
return np.load(path, mmap_mode=mmap_mode)[ids]

0 comments on commit 573a169

Please sign in to comment.