Repository navigation
Add SafeDisk.with_numpy() to store numpy arrays without pickling them - #34
Merged
Merged
Conversation
Documentation build overview
|
…orged entries with UnpicklingError.
…keep float NaN on every disk.
… with_types and with_options in any order.
…type name, and read past large strings and buffers in the write-time pickle scan.
…iting copies it once, and refuse BUILD on a stored array by test.
…Y.md and the compatibility contract from CONTRIBUTING.md.
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
SafeDiskrefuses numpy, because unpickling an array is not memory-safe, so a library that caches arrays can't use it. Today each one has to convert arrays to bytes by hand, and numpy scalars such asnumpy.float64, which pandas results are full of, raiseTypeErrorat write time.SafeDisk.with_numpy()returns aSafeDisksubclass that stores arrays without pickling them. It followswith_typesandwith_options, so the three combine in any order:SafeDisk.with_numpy().with_types(Result).with_options(on_error='miss')reads as "safe pickle, plus numpy, plus my class, and treat errors as misses". The pickler and unpickler live in a newdiskcache_community/arrays.py.SafeDiskitself still refuses numpy,UNSAFE_MODULESstays, and thewith_typeserror now points towith_numpy()..npyheader and their data, anywhere in a value: inside any containerSafeDiskaccepts, and inside classes allowed withwith_types, such as a dataclass holding arrays.Diskdoesn't change what code gets back.Diskboth ways. An array that is not contiguous or not writeable is copied once more first.ndarraysubclasses,datetime64,timedelta64andlongdoublescalars, and arrays the reader would refuse, such as a dtype with no width or a header over numpy's limit, raiseTypeErrorat write time.numpy.savewould silently drop a mask, so subclasses are refused rather than stored.SafeJSONDisk.with_numpy()raisesTypeError, since JSON holds no arrays.This also fixes a bug on every disk: a float NaN stored as a value read back as
None, because SQLite stores NaN asNULL. It is now stored as a pickle, which 5.6.3 reads too.For existing SafeDisk users
Nothing changes unless
with_numpy()is called. A plainSafeDiskwrites the same rows as before and refuses numpy as before, with the sameTypeError. The three new names on the class arewith_numpy, thesafe_numpyattribute, which isFalse, and the private_plain_unpickler. The one visible change for everyone is the NaN fix above, whereNonebecomesnan. A directory written bySafeDisk.with_numpy()is read by a plainSafeDiskexcept for the entries holding numpy data, which raiseUnpicklingErroror read as a miss withon_error='miss'.On-disk format
This is the one part that can't change after release. A row holding numpy data is a pickle of protocol 5 or higher containing persistent IDs of
('npy', <bytes of the .npy header>, <bytearray of the array data>)or('scalar', <type name>, <value>). Keeping the header apart from the data is what lets pickle write the array's memory in place.CONTRIBUTING.mdnow lists this in the compatibility contract, andtest_row_format_is_pinnedbuilds such rows by hand. A value outside the standard-library types is pickled at protocol 5 or higher whateverdisk_pickle_protocolsays, so that array data can be stored in place; other values use the configured protocol.An entry holding numpy data can't be read by another disk or by 5.6.3. They open the directory and raise
UnpicklingErrorfor that entry, and a disk withon_error='miss'deletes it. Every process that shares a directory needs this disk, and the tutorial says so. Plain entries in the same directory stay readable by everything. Arrays pickled earlier by the defaultDiskare refused like any other pickle outside the allowlist.Reading forged entries
Reading is a new path for bytes that may have been tampered with:
.npyheader are checked against the data before anything is built from it, object dtypes are refused there, a zero-width dtype can't claim a shape, and a header with bytes after it is refused.BUILD, which callsndarray.__setstate__, the call that makes numpy pickles memory-unsafe. It fails withTypeErrorbecause that call accepts only anumpy.dtypein its dtype slot and nothing the reader can build is one: not an allowlisted name, not a persistent ID. A test pins this with five forgeries, and a comment on the unpickler says why it holds.UnpicklingError, or reads as a miss withon_error='miss'. That includes a crafted header that makes the Python parser raiseMemoryError.('scalar', 'bytes', n)can't ask numpy fornzero bytes.SafeDiskreader and a persistent ID there is refused.ImportErroreven withon_error='miss'and keeps the entry, while a pickle naming a removed class still reads as a miss as documented.Also in this PR
core.pygets small hooks:Disk._store_pickleandSafeDisk._check_pickle,SafeDisk._unpicklernow takes a file so file-backed rows are streamed, and_ReaderMissing.core.pyimportsarrays.pyat its end, sincearrays.pyimports the classes above.with_typesclass peaks at the same 2.1x as a bare one rather than 4x. This applies to everySafeDisk.with_typeswrite. The pylint module limit forcore.pyrises from 4000 to 4100 lines.DISKCACHE_SAFE_TYPESandDISKCACHE_SAFE_MODULESnow apply whenDISKCACHE_DISKnames anySafeDisksubclass, so a class built withwith_numpy()in your own module, such asmyapp.cache:ArrayDisk, can be configured from the environment. There is no new environment variable.numpy.saverecipe,SECURITY.mdandCONTRIBUTING.mdare updated, andCHANGELOG.mdhas entries underUnreleased.Limits
TypeError, includingmemoizearguments.cache.incrdoesn't work on a numpy integer stored as a value, as on the defaultDisk.Disk: 14 ms and +96 MB peak RSS against 17 ms and +96 MB, each in a fresh process. Reading matches too: 6.0 ms against 5.8 ms, with no copy beyond the row.TypeError..npyheader, so many small arrays are slower than pickling: a dict of eight 100-element arrays reads in 78 µs against 7 µs on the defaultDisk, and 10,000 eight-element arrays in one value write in 127 ms and read in 90 ms, against 7 ms and 11 ms.numpy.float64in a list takes 186 ms, against 70 ms. Writing is about the same, 667 ms against 630 ms.SafeDiskand write within noise.Checked
arrays.pyat 100%), with numpy 2.2 and 2.5.tests/test_safedisk_numpy.pyalso passes on numpy 1.26 on Python 3.12, and on numpy 1.24 on Python 3.8.UnpicklingError, and read as a clean miss withon_error='miss'. Read-only, strided, windowed, broadcast and Fortran-strided arrays round-trip. 1001 scalar IDs naming one 100 KB string decode to one numpy string, and a 1.9 GB decode of a 9.5 MB row before this check is now 10 MB..npyinputs across four Python and numpy pairs were either accepted or raisedUnpicklingError, with no crash, and accepted arrays were read back in full.MemoryError. A sample of 80 of those did the same with arrays disabled, so the ones checked don't come from this code.Diskor onSafeDiskacross the classic, sizes, hit-rate, process, thread and Django workloads, apart from the NaN fix costing about 0.3 µs per NaN read.UnpicklingErrorfor the array entry.bluecheck,doc8,docs,flake8,isortcheck,mypyandpylintpass, and so do the coexistence and golden tests with upstreamdiskcache5.6.3 installed alongside.