ENH: Add compresslevel input parameter in savez_compressed() function

[geniuskey]

Proposed new feature or change:

I'm using bunch of large 2D arrays (about 10000 x 10000 ) and I need to compress and save data quickly.
The problem is that the savez_compressed() function uses only default compress level of zipfile.ZipFile class.
(The default compress level of zipfile.ZIP_DEFLATED is '6')

If I could use compress level 1 for savez_compressed() function

• w/ Compress level 1: Compression ratio: 66.8%, exec. time: about 3.2 seconds 
• w/ Compress level 6: Compression ratio: 69.5%, exec. time: about 24.7 seconds
  --> The compress level 1 is the most suitable for my project.
  Therefore I want to configure the compression level in savez_compressed() function.

Here is my solution. I could solve it with adding just few lines.

my modification code snippet of _savez() function definition of the numpy/lib/npyio.py file.

def _savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None):
    # Import is postponed to here since zipfile depends on gzip, an optional
    # component of the so-called standard library.
    import zipfile

    if not hasattr(file, 'write'):
        file = os_fspath(file)
        if not file.endswith('.npz'):
            file = file + '.npz'

    namedict = kwds
    for i, val in enumerate(args):
        key = 'arr_%d' % i
        if key in namedict.keys():
            raise ValueError(
                "Cannot use un-named variables and keyword %s" % key)
        namedict[key] = val

    if compress:
        compression = zipfile.ZIP_DEFLATED
    else:
        compression = zipfile.ZIP_STORED

    if 'compresslevel' in namedict:
        compresslevel = namedict['compresslevel']
        if not isinstance(compresslevel, int) or compresslevel < 1 or compresslevel > 9:
            compresslevel = None
    else:
        compresslevel = None

    zipf = zipfile_factory(file, mode="w", compression=compression, compresslevel=compresslevel)

    for key, val in namedict.items():
        fname = key + '.npy'
        val = np.asanyarray(val)
        # always force zip64, gh-10776
        with zipf.open(fname, 'w', force_zip64=True) as fid:
            format.write_array(fid, val,
                               allow_pickle=allow_pickle,
                               pickle_kwargs=pickle_kwargs)

    zipf.close()

It works!!

I haven't understand the github process from my idea to the merge in this kind of big open source repository cause I have no experience of github upstream activity yet. So... just let me share my idea through an issue.

thanks.

[seberg]

I feel like we had similar discussions about extending the kwargs. And wonder if it may make sense to create a np.savez_dict(file, array_mapping, *, allow_pickle=False, compression=0, ...) so that we can add keyword arguments without any worries. (The "dict" would be a misnomer, accepting any mapping.)

[bashtage]

If you want it to be more like savez, you could do:

def _savez(file, *args, *, compress=zipfile.ZIP_DEFLATED, allow_pickle=True, pickle_kwargs=None, **kwds):
    # Minimal protection for these keywords - they are now reserved
    assert isinstance(allow_pickle, bool)
    assert isinstance(pickle_kwargs, Mapping)
    assert compress in zipfile.compressor_names

This would like you then have position arguments and arbitrary variable names, and would only reserve the intended keyword arguments.

[geniuskey]

My point is that the line of zipfile_factory() function call in the _savez() function definition doesn't contain compresslevel argument.
zipf = zipfile_factory(file, mode="w", compression=compression)
Therefore, I suggest to add compresslevel argument in the function call and a pre-processing for the compresslevel argument value.

    if 'compresslevel' in namedict:
        compresslevel = namedict['compresslevel']
        if not isinstance(compresslevel, int) or compresslevel < 1 or compresslevel > 9:
            compresslevel = None
    else:
        compresslevel = None

    zipf = zipfile_factory(file, mode="w", compression=compression, compresslevel=compresslevel)

[longqzh]

Adding compresslevel argument is a good idea! I also need it and currently I edited npyio.py when I upgraded t numpy :)

[mrartemev-zero10]

This will be an amazing functionality for sure

[josephernest]

@geniuskey @seberg Is there a quick way to add this parameter on an existing Numpy distribution, without modifying/monkeypatching files in /numpy/lib/npyio.py folder?
For example by subclassing something?

More generally, are there news about this feature?

[josephernest]

Here is a short solution, that doesn't require to modify numpy files, but only monkey-patching:

import numpy as np
def zipfile_factory(file, *args, **kwargs):
    if not hasattr(file, 'read'):
        file = os_fspath(file)
    import zipfile
    kwargs['allowZip64'] = True
    kwargs['compresslevel'] = 4
    return zipfile.ZipFile(file, *args, **kwargs)
np.lib.npyio.zipfile_factory = zipfile_factory
with open("test.npz", "wb") as f:
    np.savez_compressed(f, x=np.ones(10_000_000))

See also https://stackoverflow.com/questions/78186958/how-to-monkey-patch-np-savez-compressed-to-add-compression-level-without-editin.

[geniuskey]

Here is a short solution, that doesn't require to modify numpy files, but only monkey-patching:

import numpy as np
def zipfile_factory(file, *args, **kwargs):
    if not hasattr(file, 'read'):
        file = os_fspath(file)
    import zipfile
    kwargs['allowZip64'] = True
    kwargs['compresslevel'] = 4
    return zipfile.ZipFile(file, *args, **kwargs)
np.lib.npyio.zipfile_factory = zipfile_factory
with open("test.npz", "wb") as f:
    np.savez_compressed(f, x=np.ones(10_000_000))

See also https://stackoverflow.com/questions/78186958/how-to-monkey-patch-np-savez-compressed-to-add-compression-level-without-editin.

That make sense! The idea of np.lib.npyio.zipfile_factory = zipfile_factory is brilliant. It's not common, but we can change the compression level without modifying the code.

@josephernest Is there any reason to solve this issue without any modifying or monkey-patching files in /numpy/lib/ folder? I think my original idea doesn't have side-effects.

[josephernest]

Thanks @geniuskey!

Is there any reason to solve this issue without any modifying or monkey-patching files in /numpy/lib/ folder?

Yes because if I do this, I have the burden to remember, each time I reinstall a new computer / each time I upgrade numpy, to re-modify these files. I'm sure I will forget eventually :)

[geniuskey]

Is it helpful if I make a PR of this issue for the main branch? I don't want all of numpy version support the compressionlevel parameter. I just hope the future version will support the compressionlevel parameter. 😁

[seberg]

Yes PRs are helpful, but please be aware that is also an API decision not just the code (and tests).

I.e. still think the solution is new function that doesn't use ** for the values being stored, but others may disagree. OTOH, I am not convinced that just adding compressionlevel does us much good, more things are bound to come up...