ENH: Improve performance of np.save for small arrays

[ohadravid]

Improve performance of np.savefor small arrays by removing the call to _filter_header
when writing the header, as it is known to be done in Python 3.

---

We use np.save to save a lot of small arrays (which are 3D point, so are just 3 ints/floats),
and we found out that np.save is about 20 times slower than pickle.dumps.

Profiling shows that most of the time is spent in _filter_header,
which (as I understand it) is not really needed anymore.

The result is about 10x improvement on my machine:

import numpy as np
from numpy.lib.format import write_array
import io
array = np.array([2, 3, 4])

%%timeit
buffer = io.BytesIO()
write_array(buffer, array, allow_pickle=False)

Before:
140 µs ± 10.7 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

After:
14.4 µs ± 497 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

For reference, using pickle.dumps is:
6.27 µs ± 142 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Edit: This PR also changes reading an array to only call this function for older versions which might have been created by Python 2 before this function was added in 2014 (which are still used by default),
resulting in a nice improvement as well:

Reading a v1 array (same as today):
141 µs ± 2.47 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Reading a v3 array now:
56 µs ± 13.6 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
37.3 µs ± 1.11 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) (after more little improvements)

[seberg]

Thanks! The documentation of _filter_header already allures to Python 2 support in reading, so I think that is good as is.

Since this seems worth it... there isn't a version number that we can rely on having stripped headers when reading? I am just curious if we could skip this even on reading most of the time. IIRC, we do have different version numbers but only use the new one in the rare event that the old version is not able to write the array correctly? If that is the case, then it seems like its not worth skipping the header cleanup on read.

[eric-wieser]

I agree, this line was pointless on Python3. I haven't checked to see if the CI failure is real.

[ohadravid]

Hi @seberg 👋
Looking at the release and git history, I think that the header is filtered since 2014, but versions (1,0) and (2,0) existed before and might have included Ls.
Versions (3,0) comes after 1.16 which was the last release supporting Python 2.7, so we could check if version >= (3,0), but
since np.save defaults to saving with the lowest compat this won't help too much I think, but it won't hurt :)

[charris]

This could use a release note. See doc/release/upcoming_changes/ for examples. If there are potential compatibility problems the release note should probably be in the compatibility section, otherwise improvements. Maybe two notes if both :)

[charris]

@seberg @eric-wieser Ping.

[seberg]

Happy, with putting it in. Found something to nitpick, but there always is :).

We could name the improvement.rst also performance.rst, since we actually got a "performance" category in the release notes right now.

[seberg]

OK, I doubt it is very useful, but its simple enough. If you are enthusiastic, maybe add a comment why we only do this on older "versions", though?
(Would be happy to just skip as well).

[ohadravid]

It's useful because it lets me pair a write_array(..., version=(3,0) with read_array(..., version=(3,0)) and get the nice perf boost 😎

[seberg]

Hah, fair enough :). Thanks @ohadravid!

[seberg]

Removing the sorted above looks correct to me. But I think d.keys() now prints slightly less nice? OTOH, this is de-facto an internal error so I am OK with that.

[ohadravid]

Added the sorted bash when generating the error :)