perf: avoid building a UnionArray in ak.where for numeric leaves

[ikrommyd]

When both branches are plain numeric buffers, ak.where built a UnionArray of the two and immediately had simplified merge it away. We can call nplike.where directly instead. Whether the two are mergeable is still decided by ak._do.mergeable, so the result type is the same as before.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.90%. Comparing base (60390ac) to head (373270d).

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/operations/ak_where.py	96.77% <100.00%> (+1.21%)	⬆️

... and 2 files with indirect coverage changes

[github-actions]

The documentation preview is ready to be viewed at https://awkward-array.org/doc/pr/4335/

[ikrommyd]

🤖 AI text below 🤖

CPU timings, main vs this branch, median of 7 repeats per case.

case	base (us)	branch (us)	change
where(cond, f8, 0.0) [100000]	9,954.9	2,894.5	-70.9%
where(cond, i8, f4) [100000]	9,249.0	3,247.6	-64.9%
where(cond, f8, f8) [100000]	6,482.5	2,789.8	-57.0%
where(cond, f8, f8) [1000]	186.7	105.5	-43.5%
where(cond, cond, i8, mergebool=False) [1000]	192.2	198.0	+3.0%

Median across all cases: -33.6% (10 cases total, showing the four best and the worst).

benchmark script

"""Benchmark for PR #4335: the numeric-leaf fast path in ak.where."""

import statistics
import timeit

import numpy as np

import awkward as ak


def report(label, stmt, glb, *, number, repeats=7):
    ts = timeit.repeat(stmt, number=number, repeat=repeats, globals=glb)
    us = sorted(t / number * 1e6 for t in ts)
    print(
        f"{label:<45} {statistics.median(us):10.3f} us  min {us[0]:10.3f}  sd {statistics.stdev(us):8.3f}  (n={number}x{repeats})"
    )


print(f"awkward {ak.__version__} from {ak.__file__}")


def make(n_events, avg=10, seed=0):
    rng = np.random.default_rng(seed)
    counts = rng.integers(0, 2 * avg + 1, n_events)
    total = int(counts.sum())
    return {
        "counts": counts,
        "cond": ak.unflatten(rng.random(total) < 0.5, counts),
        "f8": ak.unflatten(rng.random(total), counts),
        "f4": ak.unflatten(rng.random(total).astype(np.float32), counts),
        "i8": ak.unflatten(rng.integers(0, 100, total), counts),
        "strings": ak.unflatten(np.array(["abc"] * total), counts),
    }


for n_events, number in ((1_000, 200), (100_000, 10)):
    d = make(n_events)
    total = int(d["counts"].sum())
    size = f"{n_events} events / {total} items"
    print(f"\n--- {size} ---")
    glb = {"ak": ak, **d}

    report(f"where(cond, f8, f8)  [{n_events}]", "ak.where(cond, f8, f8)", glb, number=number)
    report(f"where(cond, i8, f4)  [{n_events}]", "ak.where(cond, i8, f4)", glb, number=number)
    report(f"where(cond, f8, 0.0) [{n_events}]", "ak.where(cond, f8, 0.0)", glb, number=number)
    # must fall back to the UnionArray path: a string leaf is not numeric
    report(f"where(cond, f8, str) [{n_events}]", "ak.where(cond, f8, strings)", glb, number=number)
    # must fall back: mergebool=False keeps bool and int in separate buffers
    report(
        f"where(cond, cond, i8, mergebool=False) [{n_events}]",
        "ak.where(cond, cond, i8, mergebool=False)",
        glb,
        number=number,
    )