I spent a day writing a NumPy kernel to replace some pandas code. It was the good
kind of optimisation: operate on the block, not the column. Compute the moments in
one pass over a 2-D array instead of looping over Series. Skip the pandas
machinery entirely.
Then I benchmarked it against the pandas code it replaced, and deleted it.
The thing I built
The library this is from, edaprep, profiles a DataFrame before deciding how to
preprocess it. Profiling means computing per-column statistics -- mean, standard
deviation, skewness, kurtosis, quantiles -- over every numeric column.
The obvious pandas version does that column by column. The obvious optimisation is
to stop doing that: take the numeric columns as one contiguous 2-D array, compute
the moments along an axis in a single vectorised pass, and hand back the results.
One trip through memory instead of N.
That reasoning is correct, and the code was faster in the microbenchmark I wrote
while developing it. I was pleased with it.
The measurement
Real benchmark, real frame shapes, candidates interleaved, minimum of five runs:
| implementation | time | peak memory |
|---|---|---|
| hand-written NumPy block kernel | 799 ms | 298 MiB |
| the pandas code it replaced | 557 ms | 1.6 MiB |
Slower, and using nearly two hundred times the memory.
The memory number is the tell, and it explains the time. Getting "the numeric
columns as one contiguous 2-D array" requires materialising that array. pandas
stores a frame as blocks; the moment your columns are not already one homogeneous
block -- different dtypes, a nullable column, anything -- .to_numpy() copies the
entire numeric half of the frame into a fresh allocation. On a 100,000-row frame
that is hundreds of megabytes of memcpy before a single arithmetic operation runs.
pandas never pays that. It computes each column's moments against the block that
already exists, in place, and allocates almost nothing.
I had optimised the arithmetic, which was never the bottleneck, by adding a copy,
which was.
The part that stung
There was a second version. I kept a branch that chose between per-column and
frame-level based on width, on the theory that wide frames would amortise the
copy. The branch was inverted. On a narrow frame it picked frame-level and took
1684 ms and 378 MiB, against 543 ms and 3.1 MiB for the per-column path.
So the fix was not to correct the branch. It was to notice that the branch had no
winning side and remove it. profiling/statistics.py now delegates to pandas, and
it is the fastest, smallest and shortest of the three versions I wrote.
Why it is in the docs and not just in the git history
The temptation is to quietly delete the branch and let the README claim the good
numbers. I put it in section 1 of docs/performance.md instead, because a
performance document that only contains wins is not evidence, it is marketing. The
same benchmark harness that produced the wins produced this, and the fact that I
published the one that went against me is the reason you might believe the others:
| operation | edaprep | baseline |
|---|---|---|
Scaler (standard) |
6.9 ms | sklearn StandardScaler 26.7 ms |
MissingValueHandler (median) |
5.2 ms | sklearn SimpleImputer 17.0 ms |
OutlierHandler (IQR clip) |
25.2 ms | the usual IQR block 41.4 ms |
The library's contributing guide now says the same thing as a rule: performance
claims need a benchmark, and be willing to delete.
Three things I would tell myself a day earlier
-
Benchmark the replacement against the thing it replaces, on realistic
shapes, not against a strawman loop on a toy frame. My microbenchmark used a
single-dtype frame, which is the one case where
.to_numpy()is free. - Measure memory alongside time. The time result was ambiguous enough that I might have argued with it. 298 MiB against 1.6 MiB is not arguable, and it explained the time.
- "Vectorised" is not a synonym for "fast." It is a synonym for "the arithmetic is not the bottleneck." If you had to copy the data to get there, you moved the bottleneck rather than removing it.
The library is edaprep -- MIT, pure
Python, pip install edaprep. It does EDA and leakage-safe preprocessing, and it
prints the reason for every decision it makes with the measurement attached, which
is the same instinct as publishing this benchmark.
There is still no compiled code in it. The one place a hand-written kernel looked
promising was slower than pandas.
Top comments (0)