memmap.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337
  1. import numpy as np
  2. from .numeric import uint8, ndarray, dtype
  3. from numpy.compat import (
  4. os_fspath, contextlib_nullcontext, is_pathlib_path
  5. )
  6. from numpy.core.overrides import set_module
  7. __all__ = ['memmap']
  8. dtypedescr = dtype
  9. valid_filemodes = ["r", "c", "r+", "w+"]
  10. writeable_filemodes = ["r+", "w+"]
  11. mode_equivalents = {
  12. "readonly":"r",
  13. "copyonwrite":"c",
  14. "readwrite":"r+",
  15. "write":"w+"
  16. }
  17. @set_module('numpy')
  18. class memmap(ndarray):
  19. """Create a memory-map to an array stored in a *binary* file on disk.
  20. Memory-mapped files are used for accessing small segments of large files
  21. on disk, without reading the entire file into memory. NumPy's
  22. memmap's are array-like objects. This differs from Python's ``mmap``
  23. module, which uses file-like objects.
  24. This subclass of ndarray has some unpleasant interactions with
  25. some operations, because it doesn't quite fit properly as a subclass.
  26. An alternative to using this subclass is to create the ``mmap``
  27. object yourself, then create an ndarray with ndarray.__new__ directly,
  28. passing the object created in its 'buffer=' parameter.
  29. This class may at some point be turned into a factory function
  30. which returns a view into an mmap buffer.
  31. Flush the memmap instance to write the changes to the file. Currently there
  32. is no API to close the underlying ``mmap``. It is tricky to ensure the
  33. resource is actually closed, since it may be shared between different
  34. memmap instances.
  35. Parameters
  36. ----------
  37. filename : str, file-like object, or pathlib.Path instance
  38. The file name or file object to be used as the array data buffer.
  39. dtype : data-type, optional
  40. The data-type used to interpret the file contents.
  41. Default is `uint8`.
  42. mode : {'r+', 'r', 'w+', 'c'}, optional
  43. The file is opened in this mode:
  44. +------+-------------------------------------------------------------+
  45. | 'r' | Open existing file for reading only. |
  46. +------+-------------------------------------------------------------+
  47. | 'r+' | Open existing file for reading and writing. |
  48. +------+-------------------------------------------------------------+
  49. | 'w+' | Create or overwrite existing file for reading and writing. |
  50. +------+-------------------------------------------------------------+
  51. | 'c' | Copy-on-write: assignments affect data in memory, but |
  52. | | changes are not saved to disk. The file on disk is |
  53. | | read-only. |
  54. +------+-------------------------------------------------------------+
  55. Default is 'r+'.
  56. offset : int, optional
  57. In the file, array data starts at this offset. Since `offset` is
  58. measured in bytes, it should normally be a multiple of the byte-size
  59. of `dtype`. When ``mode != 'r'``, even positive offsets beyond end of
  60. file are valid; The file will be extended to accommodate the
  61. additional data. By default, ``memmap`` will start at the beginning of
  62. the file, even if ``filename`` is a file pointer ``fp`` and
  63. ``fp.tell() != 0``.
  64. shape : tuple, optional
  65. The desired shape of the array. If ``mode == 'r'`` and the number
  66. of remaining bytes after `offset` is not a multiple of the byte-size
  67. of `dtype`, you must specify `shape`. By default, the returned array
  68. will be 1-D with the number of elements determined by file size
  69. and data-type.
  70. order : {'C', 'F'}, optional
  71. Specify the order of the ndarray memory layout:
  72. :term:`row-major`, C-style or :term:`column-major`,
  73. Fortran-style. This only has an effect if the shape is
  74. greater than 1-D. The default order is 'C'.
  75. Attributes
  76. ----------
  77. filename : str or pathlib.Path instance
  78. Path to the mapped file.
  79. offset : int
  80. Offset position in the file.
  81. mode : str
  82. File mode.
  83. Methods
  84. -------
  85. flush
  86. Flush any changes in memory to file on disk.
  87. When you delete a memmap object, flush is called first to write
  88. changes to disk.
  89. See also
  90. --------
  91. lib.format.open_memmap : Create or load a memory-mapped ``.npy`` file.
  92. Notes
  93. -----
  94. The memmap object can be used anywhere an ndarray is accepted.
  95. Given a memmap ``fp``, ``isinstance(fp, numpy.ndarray)`` returns
  96. ``True``.
  97. Memory-mapped files cannot be larger than 2GB on 32-bit systems.
  98. When a memmap causes a file to be created or extended beyond its
  99. current size in the filesystem, the contents of the new part are
  100. unspecified. On systems with POSIX filesystem semantics, the extended
  101. part will be filled with zero bytes.
  102. Examples
  103. --------
  104. >>> data = np.arange(12, dtype='float32')
  105. >>> data.resize((3,4))
  106. This example uses a temporary file so that doctest doesn't write
  107. files to your directory. You would use a 'normal' filename.
  108. >>> from tempfile import mkdtemp
  109. >>> import os.path as path
  110. >>> filename = path.join(mkdtemp(), 'newfile.dat')
  111. Create a memmap with dtype and shape that matches our data:
  112. >>> fp = np.memmap(filename, dtype='float32', mode='w+', shape=(3,4))
  113. >>> fp
  114. memmap([[0., 0., 0., 0.],
  115. [0., 0., 0., 0.],
  116. [0., 0., 0., 0.]], dtype=float32)
  117. Write data to memmap array:
  118. >>> fp[:] = data[:]
  119. >>> fp
  120. memmap([[ 0., 1., 2., 3.],
  121. [ 4., 5., 6., 7.],
  122. [ 8., 9., 10., 11.]], dtype=float32)
  123. >>> fp.filename == path.abspath(filename)
  124. True
  125. Flushes memory changes to disk in order to read them back
  126. >>> fp.flush()
  127. Load the memmap and verify data was stored:
  128. >>> newfp = np.memmap(filename, dtype='float32', mode='r', shape=(3,4))
  129. >>> newfp
  130. memmap([[ 0., 1., 2., 3.],
  131. [ 4., 5., 6., 7.],
  132. [ 8., 9., 10., 11.]], dtype=float32)
  133. Read-only memmap:
  134. >>> fpr = np.memmap(filename, dtype='float32', mode='r', shape=(3,4))
  135. >>> fpr.flags.writeable
  136. False
  137. Copy-on-write memmap:
  138. >>> fpc = np.memmap(filename, dtype='float32', mode='c', shape=(3,4))
  139. >>> fpc.flags.writeable
  140. True
  141. It's possible to assign to copy-on-write array, but values are only
  142. written into the memory copy of the array, and not written to disk:
  143. >>> fpc
  144. memmap([[ 0., 1., 2., 3.],
  145. [ 4., 5., 6., 7.],
  146. [ 8., 9., 10., 11.]], dtype=float32)
  147. >>> fpc[0,:] = 0
  148. >>> fpc
  149. memmap([[ 0., 0., 0., 0.],
  150. [ 4., 5., 6., 7.],
  151. [ 8., 9., 10., 11.]], dtype=float32)
  152. File on disk is unchanged:
  153. >>> fpr
  154. memmap([[ 0., 1., 2., 3.],
  155. [ 4., 5., 6., 7.],
  156. [ 8., 9., 10., 11.]], dtype=float32)
  157. Offset into a memmap:
  158. >>> fpo = np.memmap(filename, dtype='float32', mode='r', offset=16)
  159. >>> fpo
  160. memmap([ 4., 5., 6., 7., 8., 9., 10., 11.], dtype=float32)
  161. """
  162. __array_priority__ = -100.0
  163. def __new__(subtype, filename, dtype=uint8, mode='r+', offset=0,
  164. shape=None, order='C'):
  165. # Import here to minimize 'import numpy' overhead
  166. import mmap
  167. import os.path
  168. try:
  169. mode = mode_equivalents[mode]
  170. except KeyError as e:
  171. if mode not in valid_filemodes:
  172. raise ValueError(
  173. "mode must be one of {!r} (got {!r})"
  174. .format(valid_filemodes + list(mode_equivalents.keys()), mode)
  175. ) from None
  176. if mode == 'w+' and shape is None:
  177. raise ValueError("shape must be given")
  178. if hasattr(filename, 'read'):
  179. f_ctx = contextlib_nullcontext(filename)
  180. else:
  181. f_ctx = open(os_fspath(filename), ('r' if mode == 'c' else mode)+'b')
  182. with f_ctx as fid:
  183. fid.seek(0, 2)
  184. flen = fid.tell()
  185. descr = dtypedescr(dtype)
  186. _dbytes = descr.itemsize
  187. if shape is None:
  188. bytes = flen - offset
  189. if bytes % _dbytes:
  190. raise ValueError("Size of available data is not a "
  191. "multiple of the data-type size.")
  192. size = bytes // _dbytes
  193. shape = (size,)
  194. else:
  195. if not isinstance(shape, tuple):
  196. shape = (shape,)
  197. size = np.intp(1) # avoid default choice of np.int_, which might overflow
  198. for k in shape:
  199. size *= k
  200. bytes = int(offset + size*_dbytes)
  201. if mode in ('w+', 'r+') and flen < bytes:
  202. fid.seek(bytes - 1, 0)
  203. fid.write(b'\0')
  204. fid.flush()
  205. if mode == 'c':
  206. acc = mmap.ACCESS_COPY
  207. elif mode == 'r':
  208. acc = mmap.ACCESS_READ
  209. else:
  210. acc = mmap.ACCESS_WRITE
  211. start = offset - offset % mmap.ALLOCATIONGRANULARITY
  212. bytes -= start
  213. array_offset = offset - start
  214. mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start)
  215. self = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm,
  216. offset=array_offset, order=order)
  217. self._mmap = mm
  218. self.offset = offset
  219. self.mode = mode
  220. if is_pathlib_path(filename):
  221. # special case - if we were constructed with a pathlib.path,
  222. # then filename is a path object, not a string
  223. self.filename = filename.resolve()
  224. elif hasattr(fid, "name") and isinstance(fid.name, str):
  225. # py3 returns int for TemporaryFile().name
  226. self.filename = os.path.abspath(fid.name)
  227. # same as memmap copies (e.g. memmap + 1)
  228. else:
  229. self.filename = None
  230. return self
  231. def __array_finalize__(self, obj):
  232. if hasattr(obj, '_mmap') and np.may_share_memory(self, obj):
  233. self._mmap = obj._mmap
  234. self.filename = obj.filename
  235. self.offset = obj.offset
  236. self.mode = obj.mode
  237. else:
  238. self._mmap = None
  239. self.filename = None
  240. self.offset = None
  241. self.mode = None
  242. def flush(self):
  243. """
  244. Write any changes in the array to the file on disk.
  245. For further information, see `memmap`.
  246. Parameters
  247. ----------
  248. None
  249. See Also
  250. --------
  251. memmap
  252. """
  253. if self.base is not None and hasattr(self.base, 'flush'):
  254. self.base.flush()
  255. def __array_wrap__(self, arr, context=None):
  256. arr = super(memmap, self).__array_wrap__(arr, context)
  257. # Return a memmap if a memmap was given as the output of the
  258. # ufunc. Leave the arr class unchanged if self is not a memmap
  259. # to keep original memmap subclasses behavior
  260. if self is arr or type(self) is not memmap:
  261. return arr
  262. # Return scalar instead of 0d memmap, e.g. for np.sum with
  263. # axis=None
  264. if arr.shape == ():
  265. return arr[()]
  266. # Return ndarray otherwise
  267. return arr.view(np.ndarray)
  268. def __getitem__(self, index):
  269. res = super(memmap, self).__getitem__(index)
  270. if type(res) is memmap and res._mmap is None:
  271. return res.view(type=ndarray)
  272. return res