Skip to main content

Writing code the optimizer can help

The runtime optimizes array operations. Code expressed as bulk operations over arrays gives it material; code expressed as element-by-element Python does not.

Five rewrites, in rough order of how often they matter.

1. Replace scalar loops with array operations

# Interpreted, one element at a time
result = []
for v in values:
result.append(v ** 2 + 3 * v)

# One pass, vectorized
result = values ** 2 + 3 * values

Frequently one to two orders of magnitude, and shorter to read.

2. Get Python callbacks out of pandas

The most common performance mistake in dataframe code:

# Row-wise Python call: a barrier the optimizer cannot cross
df["score"] = df["value"].apply(lambda v: math.log1p(v) * 2)

# Array operation
df["score"] = np.log1p(df["value"]) * 2

If the logic genuinely cannot be expressed in array form, it may belong in a distributed task instead, so at least the Python cost is spread across workers.

3. Turn per-element branches into selects

# Divergent control flow
out = np.empty(len(x))
for i, v in enumerate(x):
out[i] = np.sqrt(v) if v > 0 else 0.0

# Compute both branches, select between them
out = np.where(x > 0, np.sqrt(np.abs(x)), 0.0)

Counter-intuitively, doing more arithmetic is faster here: np.where stays in vector form, while the loop leaves it.

4. Preallocate rather than append

# Repeated reallocation
out = []
for chunk in chunks:
out.append(transform(chunk))
result = np.concatenate(out)

# Known size, written in place
result = np.empty(total_rows, dtype=np.float32)
offset = 0
for chunk in chunks:
n = len(chunk)
result[offset:offset + n] = transform(chunk)
offset += n

Matters most when the array is large enough that reallocation copies real volume.

5. Keep memory access contiguous

Vector loads want adjacent elements.

# Row-major array, iterating columns: strided access
for j in range(cols):
process(matrix[:, j])

# Iterate the contiguous dimension
for i in range(rows):
process(matrix[i, :])

When the access pattern is fixed by the algorithm, store the data in the matching layout instead. np.ascontiguousarray after a transpose is often worth its one-time cost.

What not to do

Do not micro-optimize before measuring. Most code is not hot. Rewriting a function that accounts for 2% of runtime is effort with nothing to show. Profile first. See profile before you parallelize.

Do not trade clarity for speed in cold paths. The rewrites above mostly improve readability, which is why they are recommended. Where a fast version is genuinely uglier, apply it only where it matters and leave a comment saying why.

Do not expect this to fix algorithmic complexity. An O(n²) algorithm stays O(n²). Choosing a better algorithm remains the largest win available and no platform substitutes for it.

The short version

Instead ofWrite
for loop over scalarsarray expression
.apply(python_fn)array operation
per-element ifnp.where
list.append in a looppreallocated array
strided accesscontiguous access

Then measure rather than assume it helped.