Package: numbagg
Numbagg: Fast N-dimensional aggregation functions with NumbaFast, flexible N-dimensional array functions written withNumbaand NumPy'sgeneralized
ufuncs.Why use numbagg?PerformanceOutperforms pandasOn a single core, 2-10x faster for moving window functions, 1-2x faster for
aggregation and grouping functionsWhen parallelizing with multiple cores, 4-30x fasterOutperforms bottleneck on multiple coresOn a single core, matches bottleneckWhen parallelizing with multiple cores, 3-7x fasterOutperforms numpy on multiple coresOn a single core, matches numpyWhen parallelizing with multiple cores, 5-15x faster...though numbagg's functions are JIT compiled, so the first run is much slowerVersatilityMore functions (though bottleneck has some functions we don't have, and pandas' functions
have many more parameters)Functions work for >3 dimensions. All functions take an arbitrary axis or
tuple of axes to calculate overWritten in numba — way less code, simple to inspect, simple to improveFunctions & benchmarksSummary benchmarkTwo benchmarks summarize numbagg's performance — the first with a 1D array of 10M elements without
parallelization, and a second with a 2D array of 100x10K elements with parallelization. Numbagg's relative
performance is much higher where parallelization is possible. A wider range of arrays is
listed in the full set of benchmarks below.The values in the table are numbagg's performance as a multiple of other libraries for a
given shaped array calculated over the final axis. (so 1.00x means numbagg is equal,
higher means numbagg is faster.)func1Dpandas1Dbottleneck1Dnumpy2Dpandas2Dbottleneck2Dnumpybfill1.17x1.18xn/a12.24x4.36xn/affill1.17x1.12xn/a12.76x4.34xn/agroup_nanall1.44xn/an/a10.84xn/an/agroup_nanany1.20xn/an/a5.25xn/an/agroup_nanargmax2.88xn/an/a9.89xn/an/agroup_nanargmin2.82xn/an/a9.96xn/an/agroup_nancount1.01xn/an/a4.70xn/an/agroup_nanfirst1.39xn/an/a11.80xn/an/agroup_nanlast1.16xn/an/a5.36xn/an/agroup_nanmax1.14xn/an/a5.22xn/an/agroup_nanmean1.19xn/an/a5.64xn/an/agroup_nanmin1.13xn/an/a5.26xn/an/agroup_nanprod1.15xn/an/a4.95xn/an/agroup_nanstd1.18xn/an/a5.03xn/an/agroup_nansum_of_squares1.35xn/an/a8.11xn/an/agroup_nansum1.21xn/an/a5.95xn/an/agroup_nanvar1.19xn/an/a5.65xn/an/amove_corr19.04xn/an/a92.48xn/an/amove_cov14.58xn/an/a71.61xn/an/amove_exp_nancorr6.73xn/an/a35.30xn/an/amove_exp_nancount2.35xn/an/a10.56xn/an/amove_exp_nancov5.77xn/an/a31.75xn/an/amove_exp_nanmean2.03xn/an/a11.07xn/an/amove_exp_nanstd1.89xn/an/a10.07xn/an/amove_exp_nansum1.88xn/an/a9.70xn/an/amove_exp_nanvar1.82xn/an/a9.71xn/an/amove_mean3.82x0.87xn/a16.61x4.01xn/amove_std5.96x1.29xn/a24.52x6.04xn/amove_sum3.80x0.83xn/a15.95x3.70xn/amove_var5.78x1.27xn/a25.41x5.85xn/ananargmax[^5]2.45x1.00xn/a2.16x1.00xn/ananargmin[^5]2.19x1.01xn/a2.05x1.02xn/anancount1.40xn/a1.06x11.00xn/a4.16xnanmax[^5]3.26x1.00x0.11x3.62x3.24x0.11xnanmean2.42x0.98x2.83x13.58x4.54x13.13xnanmin[^5]3.27x1.00x0.11x3.62x3.24x0.11xnanquantile0.94xn/a0.78x5.45xn/a5.01xnanstd1.50x1.51x2.75x8.29x7.35x13.27xnansum2.28x0.97x2.52x17.71x6.24x16.05xnanvar1.50x1.49x2.81x8.18x6.97x13.32xFull benchmarksfuncshapesizepandasbottlenecknumpynumbaggpandas_ratiobottleneck_rationumpy_rationumbagg_ratiobfill(1000,)10000ms0msn/a0ms1.59x0.03xn/a1.00x(10000000,)1000000020ms20msn/a17ms1.17x1.18xn/a1.00x(100, 100000)1000000057ms20msn/a5ms12.24x4.36xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a21msn/a5msn/a4.40xn/a1.00x(100, 1000, 1000)100000000n/a248msn/a44msn/a5.70xn/a1.00xffill(1000,)10000ms0msn/a0ms1.53x0.02xn/a1.00x(10000000,)1000000020ms19msn/a17ms1.17x1.12xn/a1.00x(100, 100000)1000000056ms19msn/a4ms12.76x4.34xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a19msn/a4msn/a4.33xn/a1.00x(100, 1000, 1000)100000000n/a219msn/a42msn/a5.25xn/a1.00xgroup_nanall(1000,)10000msn/an/a0ms1.79xn/an/a1.00x(10000000,)1000000068msn/an/a47ms1.44xn/an/a1.00x(100, 100000)1000000017msn/an/a2ms10.84xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a1msn/an/an/a1.00xgroup_nanany(1000,)10000msn/an/a0ms1.78xn/an/a1.00x(10000000,)1000000068msn/an/a56ms1.20xn/an/a1.00x(100, 100000)1000000018msn/an/a3ms5.25xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanargmax(1000,)10001msn/an/a0ms17.60xn/an/a1.00x(10000000,)10000000171msn/an/a59ms2.88xn/an/a1.00x(100, 100000)1000000040msn/an/a4ms9.89xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a4msn/an/an/a1.00xgroup_nanargmin(1000,)10001msn/an/a0ms17.56xn/an/a1.00x(10000000,)10000000166msn/an/a59ms2.82xn/an/a1.00x(100, 100000)1000000041msn/an/a4ms9.96xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a4msn/an/an/a1.00xgroup_nancount(1000,)10000msn/an/a0ms1.68xn/an/a1.00x(10000000,)1000000056msn/an/a55ms1.01xn/an/a1.00x(100, 100000)1000000015msn/an/a3ms4.70xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanfirst(1000,)10000msn/an/a0ms1.88xn/an/a1.00x(10000000,)1000000063msn/an/a45ms1.39xn/an/a1.00x(100, 100000)1000000015msn/an/a1ms11.80xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a1msn/an/an/a1.00xgroup_nanlast(1000,)10000msn/an/a0ms1.87xn/an/a1.00x(10000000,)1000000062msn/an/a53ms1.16xn/an/a1.00x(100, 100000)1000000015msn/an/a3ms5.36xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a2msn/an/an/a1.00xgroup_nanmax(1000,)10000msn/an/a0ms1.89xn/an/a1.00x(10000000,)1000000066msn/an/a57ms1.14xn/an/a1.00x(100, 100000)1000000017msn/an/a3ms5.22xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanmean(1000,)10000msn/an/a0ms1.81xn/an/a1.00x(10000000,)1000000067msn/an/a57ms1.19xn/an/a1.00x(100, 100000)1000000019msn/an/a3ms5.64xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanmin(1000,)10000msn/an/a0ms1.84xn/an/a1.00x(10000000,)1000000066msn/an/a58ms1.13xn/an/a1.00x(100, 100000)1000000017msn/an/a3ms5.26xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanprod(1000,)10000msn/an/a0ms1.86xn/an/a1.00x(10000000,)1000000063msn/an/a55ms1.15xn/an/a1.00x(100, 100000)1000000016msn/an/a3ms4.95xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanstd(1000,)10000msn/an/a0ms1.73xn/an/a1.00x(10000000,)1000000070msn/an/a59ms1.18xn/an/a1.00x(100, 100000)1000000020msn/an/a4ms5.03xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a4msn/an/an/a1.00xgroup_nansum(1000,)10000msn/an/a0ms1.89xn/an/a1.00x(10000000,)1000000067msn/an/a56ms1.21xn/an/a1.00x(100, 100000)1000000019msn/an/a3ms5.95xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nanvar(1000,)10000msn/an/a0ms1.71xn/an/a1.00x(10000000,)1000000069msn/an/a58ms1.19xn/an/a1.00x(100, 100000)1000000020msn/an/a4ms5.65xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xgroup_nansum_of_squares(1000,)10000msn/an/a0ms2.36xn/an/a1.00x(10000000,)1000000075msn/an/a55ms1.35xn/an/a1.00x(100, 100000)1000000026msn/an/a3ms8.11xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a3msn/an/an/a1.00xmove_corr(1000,)10000msn/an/a0ms10.85xn/an/a1.00x(10000000,)10000000909msn/an/a48ms19.04xn/an/a1.00x(100, 100000)10000000869msn/an/a9ms92.48xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a9msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a79msn/an/an/a1.00xmove_cov(1000,)10000msn/an/a0ms10.05xn/an/a1.00x(10000000,)10000000623msn/an/a43ms14.58xn/an/a1.00x(100, 100000)10000000603msn/an/a8ms71.61xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a8msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a72msn/an/an/a1.00xmove_mean(1000,)10000ms0msn/a0ms1.84x0.03xn/a1.00x(10000000,)10000000120ms27msn/a31ms3.82x0.87xn/a1.00x(100, 100000)10000000113ms27msn/a7ms16.61x4.01xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a27msn/a7msn/a3.96xn/a1.00x(100, 1000, 1000)100000000n/a296msn/a58msn/a5.08xn/a1.00xmove_std(1000,)10000ms0msn/a0ms2.21x0.08xn/a1.00x(10000000,)10000000178ms39msn/a30ms5.96x1.29xn/a1.00x(100, 100000)10000000157ms39msn/a6ms24.52x6.04xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a39msn/a7msn/a5.88xn/a1.00x(100, 1000, 1000)100000000n/a411msn/a58msn/a7.13xn/a1.00xmove_sum(1000,)10000ms0msn/a0ms1.81x0.02xn/a1.00x(10000000,)10000000121ms26msn/a32ms3.80x0.83xn/a1.00x(100, 100000)10000000113ms26msn/a7ms15.95x3.70xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a26msn/a7msn/a3.59xn/a1.00x(100, 1000, 1000)100000000n/a281msn/a59msn/a4.77xn/a1.00xmove_var(1000,)10000ms0msn/a0ms2.04x0.08xn/a1.00x(10000000,)10000000168ms37msn/a29ms5.78x1.27xn/a1.00x(100, 100000)10000000161ms37msn/a6ms25.41x5.85xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a37msn/a6msn/a5.85xn/a1.00x(100, 1000, 1000)100000000n/a398msn/a56msn/a7.07xn/a1.00xmove_exp_nancorr(1000,)10000msn/an/a0ms7.27xn/an/a1.00x(10000000,)10000000464msn/an/a69ms6.73xn/an/a1.00x(100, 100000)10000000471msn/an/a13ms35.30xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a13msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a111msn/an/an/a1.00xmove_exp_nancount(1000,)10000msn/an/a0ms2.04xn/an/a1.00x(10000000,)1000000077msn/an/a33ms2.35xn/an/a1.00x(100, 100000)1000000069msn/an/a7ms10.56xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a6msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a59msn/an/an/a1.00xmove_exp_nancov(1000,)10000msn/an/a0ms7.07xn/an/a1.00x(10000000,)10000000298msn/an/a52ms5.77xn/an/a1.00x(100, 100000)10000000333msn/an/a10ms31.75xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a10msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a87msn/an/an/a1.00xmove_exp_nanmean(1000,)10000msn/an/a0ms1.40xn/an/a1.00x(10000000,)1000000067msn/an/a33ms2.03xn/an/a1.00x(100, 100000)1000000074msn/an/a7ms11.07xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a7msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a60msn/an/an/a1.00xmove_exp_nanstd(1000,)10000msn/an/a0ms2.33xn/an/a1.00x(10000000,)1000000088msn/an/a46ms1.89xn/an/a1.00x(100, 100000)1000000095msn/an/a9ms10.07xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a9msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a78msn/an/an/a1.00xmove_exp_nansum(1000,)10000msn/an/a0ms1.36xn/an/a1.00x(10000000,)1000000062msn/an/a33ms1.88xn/an/a1.00x(100, 100000)1000000071msn/an/a7ms9.70xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a6msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a60msn/an/an/a1.00xmove_exp_nanvar(1000,)10000msn/an/a0ms1.40xn/an/a1.00x(10000000,)1000000077msn/an/a42ms1.82xn/an/a1.00x(100, 100000)1000000084msn/an/a9ms9.71xn/an/a1.00x(10, 10, 10, 10, 1000)10000000n/an/an/a9msn/an/an/a1.00x(100, 1000, 1000)100000000n/an/an/a73msn/an/an/a1.00xnanargmax[^5](1000,)10000ms0msn/a0ms13.07x0.21xn/a1.00x(10000000,)1000000031ms12msn/a12ms2.45x1.00xn/a1.00x(100, 100000)1000000028ms13msn/a13ms2.16x1.00xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a13msn/a13msn/a1.05xn/a1.00x(100, 1000, 1000)100000000n/a133msn/a127msn/a1.05xn/a1.00xnanargmin[^5](1000,)10000ms0msn/a0ms12.72x0.21xn/a1.00x(10000000,)1000000027ms13msn/a12ms2.19x1.01xn/a1.00x(100, 100000)1000000026ms13msn/a12ms2.05x1.02xn/a1.00x(10, 10, 10, 10, 1000)10000000n/a13msn/a13msn/a1.05xn/a1.00x(100, 1000, 1000)100000000n/a135msn/a129msn/a1.05xn/a1.00xnancount(1000,)10000msn/a0ms0ms2.24xn/a0.05x1.00x(10000000,)100000005msn/a4ms3ms1.40xn/a1.06x1.00x(100, 100000)100000009msn/a3ms1ms11.00xn/a4.16x1.00x(10, 10, 10, 10, 1000)10000000n/an/a4ms1msn/an/a3.58x1.00x(100, 1000, 1000)100000000n/an/a45ms7msn/an/a6.74x1.00xnanmax[^5](1000,)10000ms0ms0ms0ms8.21x0.21x0.38x1.00x(10000000,)1000000041ms12ms1ms13ms3.26x1.00x0.11x1.00x(100, 100000)1000000045ms41ms1ms13ms3.62x3.24x0.11x1.00x(10, 10, 10, 10, 1000)10000000n/a40ms1ms12msn/a3.31x0.12x1.00x(100, 1000, 1000)100000000n/a402ms15ms121msn/a3.31x0.12x1.00xnanmean(1000,)10000ms0ms0ms0ms1.32x0.02x0.20x1.00x(10000000,)1000000023ms9ms27ms10ms2.42x0.98x2.83x1.00x(100, 100000)1000000028ms9ms27ms2ms13.58x4.54x13.13x1.00x(10, 10, 10, 10, 1000)10000000n/a9ms27ms2msn/a4.56x13.69x1.00x(100, 1000, 1000)100000000n/a91ms310ms17msn/a5.39x18.39x1.00xnanmin[^5](1000,)10000ms0ms0ms0ms8.09x0.21x0.38x1.00x(10000000,)1000000041ms12ms1ms13ms3.27x1.00x0.11x1.00x(100, 100000)1000000045ms41ms1ms13ms3.62x3.24x0.11x1.00x(10, 10, 10, 10, 1000)10000000n/a40ms1ms12msn/a3.28x0.12x1.00x(100, 1000, 1000)100000000n/a401ms15ms122msn/a3.30x0.12x1.00xnanquantile(1000,)10000msn/a0ms0ms1.46xn/a0.57x1.00x(10000000,)10000000186msn/a155ms198ms0.94xn/a0.78x1.00x(100, 100000)10000000197msn/a181ms36ms5.45xn/a5.01x1.00x(10, 10, 10, 10, 1000)10000000n/an/a425ms34msn/an/a12.50x1.00x(100, 1000, 1000)100000000n/an/a4254ms331msn/an/a12.85x1.00xnanstd(1000,)10000ms0ms0ms0ms1.06x0.06x0.46x1.00x(10000000,)1000000029ms29ms53ms19ms1.50x1.51x2.75x1.00x(100, 100000)1000000033ms29ms53ms4ms8.29x7.35x13.27x1.00x(10, 10, 10, 10, 1000)10000000n/a28ms55ms4msn/a7.25x14.43x1.00x(100, 1000, 1000)100000000n/a294ms600ms37msn/a8.02x16.35x1.00xnansum(1000,)10000ms0ms0ms0ms1.28x0.02x0.08x1.00x(10000000,)1000000022ms9ms24ms10ms2.28x0.97x2.52x1.00x(100, 100000)1000000027ms9ms24ms2ms17.71x6.24x16.05x1.00x(10, 10, 10, 10, 1000)10000000n/a9ms25ms1msn/a6.05x16.66x1.00x(100, 1000, 1000)100000000n/a90ms282ms13msn/a6.71x21.07x1.00xnanvar(1000,)10000ms0ms0ms0ms1.08x0.06x0.45x1.00x(10000000,)1000000028ms28ms53ms19ms1.50x1.49x2.81x1.00x(100, 100000)1000000033ms28ms54ms4ms8.18x6.97x13.32x1.00x(10, 10, 10, 10, 1000)10000000n/a28ms56ms4msn/a7.13x14.28x1.00x(100, 1000, 1000)100000000n/a281ms601ms32msn/a8.71x18.65x1.00x[^1][^2][^3][^4][^1]:
Benchmarks were run on a Mac M1 laptop in December 2023 on numbagg's HEAD,
pandas 2.1.1, bottleneck 1.3.7, numpy 1.25.2, withpython numbagg/test/run_benchmarks.py -- --benchmark-max-time=10. They run in CI,
though GHA's low CPU count means we don't see the full benefits of
parallelization.[^2]:
While we separate the setup and the running of the functions, pandas still
needs to do some work to create its result dataframe, and numbagg does some
checks in python which bottleneck does in C or doesn't do. So use benchmarks
on larger arrays for our summary so we can focus on the computational speed,
which doesn't asymptote away. Any contributions to improve the benchmarks are
welcome.[^3]:
In some instances, a library won't have the exact function — for example,
pandas doesn't have an equivalentmove_exp_nancountfunction, so we use
itssumfunction on an array of1s. Similarly forgroup_nansum_of_squares, we use two separate operations.[^4]:anynan&allnanare also functions in numbagg, but not listed here as they
require a different benchmark setup.[^5]:
This function is not currently parallelized, so exhibits worse performance
on parallelizable arrays.Example implementationNumbagg makes it easy to write, in pure Python/NumPy, flexible aggregation
functions accelerated by Numba. All the hard work is done by Numba's JIT
compiler and NumPy's gufunc machinery (as wrapped by Numba).For example, here is how we wrotenansum:importnumpyasnpfromnumbagg.decoratorsimportndreduce@ndreduce.wrap()defnansum(a):asum=0.0foraiina.flat:ifnotnp.isnan(ai):asum+=aireturnasumImplementation detailsNumbagg includes somewhat awkward workarounds for features missing from
NumPy/Numba:It implements its own cache for functions wrapped by Numba'sguvectorize, because that decorator is rather slow.It does itsown handling of array
transposesto handle theaxisargument in reduction functions.Itrewrites plain functions into
gufuncs,
to allow writing a traditional function while retaining the multidimensional advantages of
gufuncs.Already some of the ideas here have flowed upstream to numba (for example,an
axis parameter), and we hope
that others will follow.License3-clause BSD. Includes portions of Bottleneck, which is distributed under a
Simplified BSD license.