@blosc2.jit: compile control flow instead of tracing it, and rework the pandas story

CMakeLists.txt

@@ -109,7 +109,7 @@ endif()
 
 FetchContent_Declare(miniexpr
     GIT_REPOSITORY https://github.com/Blosc/miniexpr.git
-    GIT_TAG aab4b2ff030ffeddba894d0fe45ab7df6e53bd47
+    GIT_TAG 58d2d0b4a3aee3d1ac84b213712cf982744196c8
     # SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../miniexpr
 )
 FetchContent_MakeAvailable(miniexpr)

bench/bench_pandas_engine.py

@@ -25,15 +25,19 @@
 # quoted expression string, not that it wins a raw speed race. See
 # doc/guides/pandas_engine.md.
 #
-# Note: axis=1 (row-wise) is NOT a good fit for this engine. It still calls
-# the function once per row in a Python loop either way, and for a handful
-# of columns the wrapping overhead per call (building a compute-engine proxy
-# for a tiny array) is larger than the win, so engine=blosc2.jit is actually
-# *slower* than plain apply(axis=1) in that case. Use axis=0 (or restructure
-# the computation to operate on whole columns) to get the engine's benefit.
+# Row-wise (axis=1) computations are a different story. apply() cannot express
+# per-row iteration at all, so the pattern to reach for is a function taking
+# the columns as separate array parameters, called directly (no df.apply) --
+# with **df, since a DataFrame unpacks into one keyword argument per column.
+# See "Row-wise computations" in doc/guides/pandas_engine.md. bench_row_wise()
+# below measures that pattern against a plain per-row apply() and vectorized
+# NumPy on a genuine per-row-convergence problem (Kepler's equation via
+# Newton-Raphson), where a real per-row `break` beats even vectorized NumPy,
+# and sweeps both what drives that win (rows, and how unevenly rows converge).
 #
 # Each measurement is the minimum of NRUNS repetitions to reduce noise.
 
+import math
 from pathlib import Path
 from time import perf_counter
 
@@ -49,6 +53,18 @@
 
 ROW_SWEEP = (1_000, 10_000, 100_000, 1_000_000, 5_000_000)
 
+# Plain per-row apply(axis=1) is ~1000x slower than the alternatives below;
+# keep this sweep small so the benchmark finishes in a reasonable time.
+ROW_WISE_APPLY_NROWS = 2_000
+
+KEPLER_ROW_SWEEP = (10_000, 100_000, 1_000_000, 5_000_000)
+
+# Maximum orbital eccentricity: the knob controlling how *unevenly* rows
+# converge. Near-circular orbits (0.1) all converge in the same 3 iterations;
+# near-parabolic ones (0.99) leave a slow tail that vectorized NumPy must keep
+# sweeping the whole array for, while the DSL kernel's per-row break does not.
+KEPLER_ECC_SWEEP = (0.1, 0.5, 0.9, 0.99)
+
 OUT_DIR = Path(__file__).resolve().parent.parent / "doc" / "guides" / "pandas_engine"
 
 # dataviz reference palette, same values as bench/optim_tips/common.py
@@ -78,6 +94,26 @@ def yeo_johnson(col, lam=0.5):
     return np.where(col >= 0, pos, neg)
 
 
+# The same transform written with a real per-element if/else. It compiles to a
+# DSL kernel instead of being traced, so only the matching arm runs for each
+# element and no clamping is needed. lam is inlined because apply() passes the
+# column alone. The explanation lives here rather than in a docstring: a DSL
+# kernel body cannot contain a string literal.
+def yeo_johnson_branch(col):
+    if col >= 0:
+        out = (np.power(col + 1.0, 0.5) - 1.0) / 0.5
+    else:
+        out = -(np.power(-col + 1.0, 1.5) - 1.0) / 1.5
+    return out
+
+
+def yeo_johnson_scalar(x, lam=0.5):
+    """Per-element Python, the shape you would write without any engine."""
+    if x >= 0:
+        return ((x + 1.0) ** lam - 1.0) / lam
+    return -((-x + 1.0) ** (2.0 - lam) - 1.0) / (2.0 - lam)
+
+
 # The same transform as a single numexpr expression: legal, but this is what
 # the readability argument is about.
 YEO_JOHNSON_NX = (
@@ -123,6 +159,151 @@ def speedup(df, func):
     return t_plain / t_engine, t_plain, t_engine
 
 
+# Kepler's equation, solved by Newton-Raphson: a genuine per-row-convergence
+# problem (row["colname"] combines two columns, and rows converge in a
+# different number of iterations), used to benchmark the row-wise
+# "columns as direct-call parameters" pattern from doc/guides/pandas_engine.md
+# against both a plain per-row apply(axis=1) and vectorized NumPy.
+def kepler_row_scalar(row):
+    m = row["mean_anomaly"]
+    ecc = row["eccentricity"]
+    e = m + ecc * math.sin(m)
+    for _ in range(100):
+        diff = (e - ecc * math.sin(e) - m) / (1.0 - ecc * math.cos(e))
+        e = e - diff
+        if abs(diff) < 1e-12:
+            break
+    return e
+
+
+def kepler_numpy(m, ecc):
+    e = m + ecc * np.sin(m)
+    for _ in range(100):
+        diff = (e - ecc * np.sin(e) - m) / (1.0 - ecc * np.cos(e))
+        e = e - diff
+        if np.max(np.abs(diff)) < 1e-12:
+            break
+    return e
+
+
+@blosc2.jit
+def kepler_dsl(mean_anomaly, eccentricity):
+    e = mean_anomaly + eccentricity * sin(mean_anomaly)  # noqa: F821  # 'sin' resolved as a bare DSL function name
+    for _ in range(100):
+        diff = (e - eccentricity * sin(e) - mean_anomaly) / (1.0 - eccentricity * cos(e))  # noqa: F821
+        e = e - diff
+        if abs(diff) < 1e-12:
+            break
+    return e
+
+
+def make_kepler_df(nrows, ecc_max=0.95):
+    rng = np.random.default_rng(1)
+    return pd.DataFrame(
+        {
+            "mean_anomaly": rng.uniform(0, 2 * np.pi, nrows),
+            "eccentricity": rng.uniform(0.0, ecc_max, nrows),
+        }
+    )
+
+
+def kepler_max_iters(m, ecc):
+    """Iterations the slowest-converging row needs -- what vectorized NumPy
+    pays for every row, and what the DSL kernel's per-row break avoids."""
+    e = m + ecc * np.sin(m)
+    for k in range(100):
+        diff = (e - ecc * np.sin(e) - m) / (1.0 - ecc * np.cos(e))
+        e = e - diff
+        if np.max(np.abs(diff)) < 1e-12:
+            return k + 1
+    return 100
+
+
+def kepler_speedup(df):
+    """Vectorized NumPy vs the direct DSL call, returning (speedup, t_numpy, t_dsl)."""
+    m = df["mean_anomaly"].to_numpy()
+    ecc = df["eccentricity"].to_numpy()
+    t_numpy, result_numpy = timeit(lambda: kepler_numpy(m, ecc))
+    # Columns passed by keyword via **df: kernel parameters are named after the
+    # DataFrame columns, so no column has to be restated at the call site.
+    t_dsl, result_dsl = timeit(lambda: np.asarray(kepler_dsl(**df)))
+    np.testing.assert_allclose(result_dsl, result_numpy, atol=1e-9)
+    return t_numpy / t_dsl, t_numpy, t_dsl
+
+
+def bench_row_wise():
+    # Slice from one frame rather than calling make_kepler_df(n) twice with
+    # different n: a fresh same-seeded Generator's bulk draws are not
+    # guaranteed to share a common prefix across different requested sizes.
+    df_full = make_kepler_df(NROWS)
+    df_small = df_full.iloc[:ROW_WISE_APPLY_NROWS]
+    m = df_full["mean_anomaly"].to_numpy()
+    ecc = df_full["eccentricity"].to_numpy()
+
+    t_apply, result_apply = timeit(lambda: df_small.apply(kepler_row_scalar, axis=1))
+    t_numpy, result_numpy = timeit(lambda: kepler_numpy(m, ecc))
+    t_dsl, result_dsl = timeit(lambda: np.asarray(kepler_dsl(**df_full)))
+
+    # Cross-check correctness: plain apply on the small frame vs numpy on the
+    # same rows, and the direct DSL call vs numpy on the full frame.
+    np.testing.assert_allclose(
+        result_apply.to_numpy(),
+        kepler_numpy(m[:ROW_WISE_APPLY_NROWS], ecc[:ROW_WISE_APPLY_NROWS]),
+        atol=1e-9,
+    )
+    np.testing.assert_allclose(result_dsl, result_numpy, atol=1e-9)
+
+    print("\nrow-wise (axis=1), Kepler's equation via Newton-Raphson:")
+    print(f"  plain apply(axis=1),      {ROW_WISE_APPLY_NROWS:>9,} rows:  {t_apply:.4f} s")
+    print(f"  vectorized numpy,         {NROWS:>9,} rows:  {t_numpy:.4f} s")
+    print(f"  direct DSL call,          {NROWS:>9,} rows:  {t_dsl:.4f} s   {t_numpy / t_dsl:.2f}x vs numpy")
+    per_row_apply = t_apply / ROW_WISE_APPLY_NROWS
+    per_row_dsl = t_dsl / NROWS
+    print(
+        f"  per row: apply {per_row_apply * 1e6:.1f} us vs direct DSL call {per_row_dsl * 1e6:.4f} us "
+        f"(~{per_row_apply / per_row_dsl:,.0f}x)"
+    )
+
+    print("\nkepler rows sweep (speedup of the direct DSL call vs vectorized numpy):")
+    row_speedups = []
+    for nrows in KEPLER_ROW_SWEEP:
+        sp, tn, td = kepler_speedup(make_kepler_df(nrows))
+        row_speedups.append(sp)
+        print(f"  {nrows:>9,} rows:  numpy {tn:.4f} s  DSL {td:.4f} s  {sp:.2f}x")
+
+    print("\nkepler eccentricity sweep (how unevenly rows converge):")
+    ecc_speedups, ecc_iters = [], []
+    for ecc_max in KEPLER_ECC_SWEEP:
+        df = make_kepler_df(NROWS, ecc_max=ecc_max)
+        iters = kepler_max_iters(df["mean_anomaly"].to_numpy(), df["eccentricity"].to_numpy())
+        sp, tn, td = kepler_speedup(df)
+        ecc_speedups.append(sp)
+        ecc_iters.append(iters)
+        print(
+            f"  e < {ecc_max:<5}  slowest row: {iters:>2} iters   "
+            f"numpy {tn:.4f} s  DSL {td:.4f} s  {sp:.2f}x"
+        )
+
+    out_path = OUT_DIR / "kepler.png"
+    save_kepler_plot(row_speedups, ecc_speedups, ecc_iters, out_path)
+    print(f"\nplot saved to {out_path}")
+
+
+def style_speedup_axes(ax, values):
+    """Shared look for the speedup panels: break-even line, x-suffixed ticks."""
+    # Break-even: below this line the faster-looking option is a net loss.
+    ax.axhline(1.0, color=MUTED, linestyle="--", linewidth=1)
+    ax.set_ylim(0, max(values) * 1.25)
+    ax.yaxis.set_major_formatter(lambda v, _pos: f"{v:g}x")
+    ax.yaxis.grid(True, color=GRID, linewidth=0.8)
+    ax.set_axisbelow(True)
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.spines["left"].set_color(GRID)
+    ax.spines["bottom"].set_color(GRID)
+    ax.tick_params(labelsize=9, colors=MUTED)
+
+
 def save_plot(row_speedups, ops_speedups, out_path):
     import matplotlib
 
@@ -144,17 +325,7 @@ def save_plot(row_speedups, ops_speedups, out_path):
     ax_ops.set_title(f"{NROWS:,} rows x {NCOLS} columns", color=MUTED, fontsize=9)
 
     for ax, values in ((ax_rows, row_speedups), (ax_ops, ops_speedups)):
-        # Break-even: below this line the engine is a net loss.
-        ax.axhline(1.0, color=MUTED, linestyle="--", linewidth=1)
-        ax.set_ylim(0, max(values) * 1.25)
-        ax.yaxis.set_major_formatter(lambda v, _pos: f"{v:g}x")
-        ax.yaxis.grid(True, color=GRID, linewidth=0.8)
-        ax.set_axisbelow(True)
-        ax.spines["top"].set_visible(False)
-        ax.spines["right"].set_visible(False)
-        ax.spines["left"].set_color(GRID)
-        ax.spines["bottom"].set_color(GRID)
-        ax.tick_params(labelsize=9, colors=MUTED)
+        style_speedup_axes(ax, values)
 
     fig.suptitle(
         "df.apply(f, engine=blosc2.jit): when it pays off",
@@ -167,6 +338,61 @@ def save_plot(row_speedups, ops_speedups, out_path):
     plt.close(fig)
 
 
+def save_kepler_plot(row_speedups, ecc_speedups, ecc_iters, out_path):
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, (ax_rows, ax_ecc) = plt.subplots(1, 2, figsize=(8, 3.2))
+
+    ax_rows.semilogx(KEPLER_ROW_SWEEP, row_speedups, "o-", color=COLOR_TIP, linewidth=2)
+    ax_rows.set_xlabel("rows (log scale)", color=INK, fontsize=9)
+    ax_rows.set_ylabel("speedup vs vectorized NumPy", color=INK, fontsize=9)
+    ax_rows.set_title("eccentricity < 0.95", color=MUTED, fontsize=9)
+
+    ax_ecc.plot(range(len(KEPLER_ECC_SWEEP)), ecc_speedups, "o-", color=COLOR_TIP, linewidth=2)
+    ax_ecc.set_xticks(range(len(KEPLER_ECC_SWEEP)))
+    ax_ecc.set_xticklabels([f"< {e}\n({n} iters)" for e, n in zip(KEPLER_ECC_SWEEP, ecc_iters, strict=True)])
+    ax_ecc.set_xlabel("eccentricity (iterations the slowest row needs)", color=INK, fontsize=9)
+    ax_ecc.set_title(f"{NROWS:,} rows", color=MUTED, fontsize=9)
+
+    for ax, values in ((ax_rows, row_speedups), (ax_ecc, ecc_speedups)):
+        style_speedup_axes(ax, values)
+
+    fig.suptitle(
+        "Kepler by Newton-Raphson: direct DSL call vs vectorized NumPy",
+        fontsize=11,
+        color=INK,
+    )
+    fig.tight_layout(rect=[0, 0, 1, 0.90])
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(out_path, dpi=150)
+    plt.close(fig)
+
+
+SCALAR_ROWS = 50_000
+
+
+def bench_branch_vs_where(df, t_plain, result_plain):
+    """Real per-element if vs the traced np.where form, plus per-element Python.
+
+    The scalar version is timed on a smaller frame and extrapolated: at the full
+    size it takes over a second per run.
+    """
+    t_branch, result_branch = timeit(lambda: df.apply(yeo_johnson_branch, engine=blosc2.jit))
+    pd.testing.assert_frame_equal(result_branch, result_plain)
+
+    small = make_df(nrows=SCALAR_ROWS)
+    t_small, _ = timeit(lambda: small.apply(lambda col: col.map(yeo_johnson_scalar)))
+    t_scalar = t_small * (NROWS / SCALAR_ROWS)
+
+    print("\nreal if vs np.where (both under engine=blosc2.jit):")
+    print(f"  per-element Python, real if:   {t_scalar:.4f} s (extrapolated)  {t_plain / t_scalar:.2f}x")
+    print(f"  engine, real if (DSL kernel):  {t_branch:.4f} s   {t_plain / t_branch:.2f}x")
+    print("  (np.where form is the t_engine figure above)")
+
+
 def main():
     df = make_df()
 
@@ -182,6 +408,8 @@ def main():
     print(f"df.apply(f, engine=blosc2.jit):  {t_engine:.4f} s   {t_plain / t_engine:.2f}x")
     print(f"numexpr per column:              {t_numexpr:.4f} s   {t_plain / t_numexpr:.2f}x")
 
+    bench_branch_vs_where(df, t_plain, result_plain)
+
     print("\nrows sweep (speedup vs plain apply):")
     row_speedups = []
     for nrows in ROW_SWEEP:
@@ -200,6 +428,8 @@ def main():
     save_plot(row_speedups, ops_speedups, out_path)
     print(f"\nplot saved to {out_path}")
 
+    bench_row_wise()
+
 
 if __name__ == "__main__":
     main()