diff --git a/CMakeLists.txt b/CMakeLists.txt index 9b0294d06..1895b275a 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/bench/bench_pandas_engine.py b/bench/bench_pandas_engine.py index 898be4dce..cb03fdfab 100644 --- a/bench/bench_pandas_engine.py +++ b/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() diff --git a/bench/ndarray/jit-dsl-mandelbrot.py b/bench/ndarray/jit-dsl-mandelbrot.py new file mode 100644 index 000000000..6e9dfd5a4 --- /dev/null +++ b/bench/ndarray/jit-dsl-mandelbrot.py @@ -0,0 +1,103 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +# Compares a NumPy-vectorized Mandelbrot escape-time kernel against the same +# kernel run through @blosc2.jit's DSL (control-flow) dispatch route, directly +# on NumPy operands. Tracing would silently drop the per-pixel loop/break, so +# jit compiles the whole function with miniexpr instead; see +# doc/guides/optimization_tips.md ("Let @blosc2.jit compile control flow +# instead of tracing it"). +# +# Return paths are equalized (both calls end in a plain NumPy array): any +# non-None jit() kwarg flips the return from `retval[()]` to `.compute()`, +# which would otherwise skew the comparison. + +from __future__ import annotations + +import argparse +import statistics +import time + +import numpy as np + +import blosc2 + + +@blosc2.jit +def mandelbrot_jit(cr, ci, max_iter): + zr = 0.0 + zi = 0.0 + n = 0 + for _i in range(max_iter): + if zr * zr + zi * zi > 4.0: + break + new_zr = zr * zr - zi * zi + cr + zi = 2 * zr * zi + ci + zr = new_zr + n = n + 1 + return n + + +def mandelbrot_numpy(cr, ci, max_iter): + zr = np.zeros_like(cr) + zi = np.zeros_like(ci) + n = np.zeros(cr.shape, dtype=np.int64) + active = np.ones(cr.shape, dtype=bool) + for _i in range(max_iter): + mag = zr * zr + zi * zi + active &= mag <= 4.0 + new_zr = zr * zr - zi * zi + cr + new_zi = 2 * zr * zi + ci + zr = np.where(active, new_zr, zr) + zi = np.where(active, new_zi, zi) + n = np.where(active, n + 1, n) + return n + + +def _grid(height: int, width: int) -> tuple[np.ndarray, np.ndarray]: + cr = np.linspace(-2.0, 1.0, width, dtype=np.float64)[None, :] * np.ones((height, 1)) + ci = np.linspace(-1.5, 1.5, height, dtype=np.float64)[:, None] * np.ones((1, width)) + return cr, ci + + +def _bench(fn, reps: int, warmup: int) -> float: + for _ in range(warmup): + fn() + times = [] + for _ in range(reps): + t0 = time.perf_counter() + fn() + times.append(time.perf_counter() - t0) + return statistics.median(times) + + +def main(): + parser = argparse.ArgumentParser(description="NumPy vs @blosc2.jit (DSL route) Mandelbrot benchmark.") + parser.add_argument("--height", type=int, default=800) + parser.add_argument("--width", type=int, default=1200) + parser.add_argument("--max-iter", type=int, default=100) + parser.add_argument("--reps", type=int, default=5) + parser.add_argument("--warmup", type=int, default=1) + args = parser.parse_args() + + cr, ci = _grid(args.height, args.width) + print(f"grid: {args.height}x{args.width}, max_iter={args.max_iter}") + + ref = mandelbrot_numpy(cr, ci, args.max_iter) + got = mandelbrot_jit(cr, ci, args.max_iter) + assert np.array_equal(ref, got), "jit DSL result does not match the NumPy reference" + + numpy_med = _bench(lambda: mandelbrot_numpy(cr, ci, args.max_iter), args.reps, args.warmup) + jit_med = _bench(lambda: mandelbrot_jit(cr, ci, args.max_iter), args.reps, args.warmup) + + print(f"numpy vectorized: {numpy_med:.6f} s") + print(f"blosc2.jit (DSL route): {jit_med:.6f} s") + print(f"speedup: {numpy_med / jit_med:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/doc/guides/optimization_tips.md b/doc/guides/optimization_tips.md index e01e2b41d..b567b9505 100644 --- a/doc/guides/optimization_tips.md +++ b/doc/guides/optimization_tips.md @@ -58,6 +58,27 @@ The output passes light uniformity checks against NumPy's PCG64 (the benchmark s *Benchmark for this tip: [`tip_11_dsl_random.py`](https://github.com/Blosc/python-blosc2/blob/main/bench/optim_tips/tip_11_dsl_random.py)* +## Let `@blosc2.jit` compile control flow instead of tracing it + +{func}`@blosc2.jit ` normally works by *tracing*: it calls your function once with proxy operands to record a `LazyExpr` string, so an `if`/`for`/`while` in the body only ever sees one (traced) path — the rest is silently lost. When the body contains control flow **and** it fits the [DSL grammar](../reference/dsl_syntax.md), `jit` detects this at decoration time and instead compiles the whole function with the same miniexpr engine that powers `@blosc2.dsl_kernel`, so every branch and loop runs as written, once per chunk, on NumPy or NDArray operands directly (no conversion copy). + +```python +# Without this detection, jit would call the function once, record only the +# branch that one call happened to take, and reuse it for every pixel — not +# a real per-pixel escape-time loop. +@blosc2.jit +def mandelbrot(cr, ci, max_iter): + zr, zi, n = 0.0, 0.0, 0 + for _ in range(max_iter): + if zr * zr + zi * zi > 4.0: + break + zr, zi = zr * zr - zi * zi + cr, 2 * zr * zi + ci + n += 1 + return n # jit detects the control flow and compiles this kernel whole +``` + +Functions **without** control flow always trace, even when they happen to be DSL-valid: tracing plus vectorized `ne_evaluate`/miniexpr is faster than whole-kernel miniexpr for pure elementwise expressions, so `jit` only takes the DSL route when tracing would silently lose branches/loops. Use `jit(strict=True)` to force the DSL route regardless (raises at decoration time if the function can't be compiled), or `jit(strict=False)` to force tracing even with control flow (only correct when branches depend on plain Python values, not on the arrays themselves). + ## Align your reads with the double partition blosc2 arrays are partitioned twice: the array is split into **chunks** (the unit of storage and compression), and each chunk is subdivided into **blocks** (the unit of decompression, sized to fit CPU caches). A read that lands exactly on a partition boundary decompresses only the chunk or block it needs, while the same-sized read shifted off-grid straddles (and decompresses) extra ones. diff --git a/doc/guides/pandas_engine.md b/doc/guides/pandas_engine.md index d80dd5fd5..25d90c556 100644 --- a/doc/guides/pandas_engine.md +++ b/doc/guides/pandas_engine.md @@ -1,20 +1,27 @@ -# Using Blosc2 as a pandas Engine +# Using Blosc2 with pandas -pandas' `DataFrame.apply` and `Series.map` accept an `engine=` argument, and -`blosc2.jit` is one such engine. Instead of running your function once per -element in a Python loop, pandas hands it the **whole column at once**, and -Blosc2 evaluates the entire function body in a single multi-threaded pass -over the data. +There are two ways to make a pandas computation faster with Blosc2, and which +one you want depends on a single question: **does your function work on one +column at a time, or does it combine several columns per row?** -The result is typically **2-3x faster than a plain `apply`**, with the -function itself left exactly as you wrote it. +| your function | use | typical win | +| --- | --- | --- | +| transforms one column (`axis=0`) | `df.apply(f, engine=blosc2.jit)` | 2-3x | +| combines columns, one result per row | `f(**df)` — no `apply` at all | 5-7x | + +Both need a decent amount of data and a decent amount of arithmetic to be +worth it; the sections below say how much. -## An example +## One column at a time: `engine=blosc2.jit` -Yeo-Johnson is the power transform behind scikit-learn's `PowerTransformer`, -used to make skewed features more normally distributed. Its parameter is -fitted per column, so applying it to every column of a feature matrix is -precisely what you want: +`DataFrame.apply` and `Series.map` accept an `engine=` argument, and +`blosc2.jit` is one such engine. pandas hands your function the **whole column +at once**, and Blosc2 evaluates the entire function body in a single +multi-threaded pass. + +Yeo-Johnson — the power transform behind scikit-learn's `PowerTransformer` — +is a good fit: its parameter is fitted per column, so you apply it to every +column of a feature matrix. ```python import numpy as np @@ -37,159 +44,199 @@ def yeo_johnson(col, lam=0.5): result = df.apply(yeo_johnson, engine=blosc2.jit) ``` -`result` is a DataFrame of the same shape and column names as `df`, with the -transform applied to every column — the same thing plain `df.apply(yeo_johnson)` -returns, only computed differently. On the machine below it takes 0.056 s -instead of 0.119 s, a **2.1x** speedup. - -## Why it is faster +You get back the same DataFrame plain `df.apply(yeo_johnson)` would return, +computed differently: 0.058 s instead of 0.121 s, a **2.1x** speedup. -Evaluated by NumPy, that function body is a sequence of separate steps: raise -to a power, subtract, divide, do it again for the negative branch, then select. -Each step walks the full column and allocates a new full-size temporary array -to hold its result. +The win comes from not materialising intermediates. NumPy runs that body one +operation at a time, allocating a full-size temporary array at each step. +Blosc2 captures the whole expression first, then makes one pass over the data, +computing every step on a small piece while it is still in cache and spreading +the pieces across cores. -Blosc2 does not execute the steps one at a time. It captures the whole -expression first, then makes a single pass over the data, computing every step -on one small piece while that piece is still in cache, and spreading the pieces -across cores. The intermediate arrays are never created. - -## When it pays off - -Two things have to be true, and the plot below measures each one on its own. +### When it pays off ![Speedup vs rows and vs number of fused operations](pandas_engine/speedup.png) -**Enough rows.** Setting up the compute engine costs a fixed amount per call. -At a few hundred thousand rows that setup is still larger than anything it -saves, and the engine is a net loss — at 100,000 rows it runs at 0.70x. Break-even -falls between 100,000 and 1,000,000 rows. - -**Enough arithmetic.** The more operations there are to fuse into one pass, the -more temporaries are avoided and the bigger the win — from 2.1x for a single -operation up to 3.8x for five. +**Enough rows** (left): setup costs a fixed amount per call, so break-even +falls between 100,000 and 1,000,000 rows. Below that, use a plain `apply`. +Beyond a few million the win eases off — the data no longer fits in cache and +memory bandwidth becomes the limit. -Beyond a few million rows the speedup flattens and then eases off (1.9x at 5M -above): the arrays no longer fit in cache and the whole computation becomes -limited by memory bandwidth, which fusion can reduce but not eliminate. +**Enough arithmetic** (right): the more operations there are to fuse, the more +temporaries are skipped — 1.9x for a single operation, 3.5x for five. One +cheap operation over a big array wins nothing at all +(`df.apply(lambda col: col + 1, engine=blosc2.jit)` runs at *half* the speed +of a plain `apply`). Reach for the engine when the function does real work per +element. -## When not to use it +`pd.eval(..., engine="numexpr")` fuses expressions the same way and is +somewhat faster here; the reason to prefer `engine=blosc2.jit` is that you +write a Python function instead of a quoted expression string. See +[Details](#details). -**Trivial expressions.** Arithmetic intensity matters more than the raw number -of operations. A single cheap operation over a large array is limited by memory -bandwidth, not by computation, so there is nothing for the engine to win back: +## Several columns per row: skip `apply`, pass the columns -```python -result = df.apply(lambda col: col + 1, engine=blosc2.jit) # 0.53x — slower! -``` - -That runs at roughly **half** the speed of a plain `apply`. Reach for the engine -when the function does real work per element, such as transcendental functions -or several chained operations. - -**Small frames.** See the plot above: below a few hundred thousand rows, use a -plain `apply`. +This is the case pandas 3 highlights `engine=` for, and it is where `apply` is +the wrong shape: its contract is one call per row. A `@blosc2.jit` function +takes one parameter per column and is called directly — and since a DataFrame +unpacks into one keyword argument per column, that call is just `f(**df)`. -## Compared to `pd.eval` and numexpr +Kepler's equation, solved per row by Newton-Raphson: -pandas' own [Enhancing performance](https://pandas.pydata.org/docs/user_guide/enhancingperf.html) -guide describes `pd.eval(..., engine="numexpr")`, which fuses expressions using -the same underlying idea. It is worth being straightforward about how the two -compare on the example above: - -| approach | time | vs plain apply | -| --- | --- | --- | -| plain `df.apply(f)` | 0.1194 s | 1.00x | -| `df.apply(f, engine=blosc2.jit)` | 0.0561 s | 2.13x | -| `numexpr.evaluate(...)` per column | 0.0380 s | 3.14x | +```python +rng = np.random.default_rng(0) +orbits = pd.DataFrame( + { + "mean_anomaly": rng.uniform(0, 2 * np.pi, 1_000_000), + "eccentricity": rng.uniform(0.0, 0.95, 1_000_000), + } +) -On an in-memory DataFrame, numexpr is somewhat faster. The reason to prefer -`engine=blosc2.jit` is not raw speed but that **you write a Python function -rather than a quoted string**. The numexpr equivalent of `yeo_johnson` has to -become a single expression: -```python -"where(c >= 0, " -"((maximum(c, 0.0) + 1.0) ** lam - 1.0) / lam, " -"-((maximum(-c, 0.0) + 1.0) ** (2.0 - lam) - 1.0) / (2.0 - lam))" +@blosc2.jit +def kepler(mean_anomaly, eccentricity): + e = mean_anomaly + eccentricity * sin(mean_anomaly) + for _ in range(100): + diff = (e - eccentricity * sin(e) - mean_anomaly) / ( + 1.0 - eccentricity * cos(e) + ) + e = e - diff + if abs(diff) < 1e-12: + break + return e + + +orbits["E"] = kepler(**orbits) ``` -The Blosc2 version keeps its intermediate variables and its name, can call -helper functions, can be unit-tested and reused elsewhere under `@blosc2.jit`, -and is checked by your editor and linters. That is the trade being offered. - -Note also that Blosc2's characteristic strength — computing directly over -compressed, potentially larger-than-memory arrays — does not come into play on -this path, because pandas materialises each column as a plain NumPy array -before the engine ever sees it. +No `apply`, no `engine=`. The parameter names match the column names, so `**` +does the wiring — by name, so the column order in the frame is irrelevant. +Each column arrives as a pandas Series, which the kernel accepts like any +array (zero-copy for ordinary numeric dtypes). `**df` passes +*every* column, so subset first if the frame has more: +`kepler(**orbits[["mean_anomaly", "eccentricity"]])`. + +On 1,000,000 rows that runs in 0.027 s against 0.165 s for fully vectorized +NumPy — **6.2x** — and about 164x faster per row than a plain +`df.apply(..., axis=1)`. + +### Why it beats even vectorized NumPy + +Vectorized NumPy has to keep looping until the *worst* row converges: the loop +is at the whole-array level, so every row pays for the slowest one. The +kernel's `for`/`if`/`break` compile to a real per-element loop, so each row +stops as soon as *it* converges — and the whole thing still runs as one fused, +multi-threaded pass. + +![Kepler speedup vs rows and vs eccentricity](pandas_engine/kepler.png) + +That explanation predicts the right-hand panel, which varies orbital +eccentricity — i.e. how much harder the worst row is than a typical one. With +near-circular orbits every row converges in the same 3 iterations, there is +nothing for `break` to skip, and the win falls to 2.7x. With near-parabolic +ones the slowest row needs 10 iterations while the average needs 4, and the +win climbs to 7.4x. **The more uneven the work per row, the more this +pattern wins.** The left panel is the familiar row-count story: break-even +around 20,000 rows, 3.5x at 100,000. + +Note that the kernel never learns it was handed a DataFrame. The same call +works with polars, xarray, h5py — or a `blosc2.NDArray`, which is how you +reach compressed, larger-than-memory operands. Only the `**df` spelling is +pandas-specific. + +### If your function has no per-row loop + +`df.apply(f, engine=blosc2.jit, axis=1)` does work for functions that merely +combine columns by name (`row["a"] + row["b"]`): they are dispatched to a +single whole-column call rather than a per-row Python loop. Add a `for` or +`while` to that idiom and it raises a `TypeError` pointing back here — that +shape can only be compiled, not traced. ## Gotchas -**`std()` silently changes meaning.** A plain `apply` passes your function a -pandas Series, whose `.std()` defaults to `ddof=1`. The engine passes a NumPy -array, whose `.std()` defaults to `ddof=0`. The same source code therefore -computes slightly different numbers depending on the engine. Pass `ddof` -explicitly if it matters: - -```python -z = (col - col.mean()) / col.std(ddof=0) -``` - -**No Python `if` on values.** The function is traced rather than executed -statement by statement, so branching on array contents raises -`ValueError: The truth value of an array ... is ambiguous`. Use `np.where`, -nesting it where you would have used `elif`. (This is the same restriction -numexpr has.) Branching on a scalar *parameter* is fine. +**Your function normally runs only once.** The engine calls it a single time +with stand-in objects that record operations rather than compute them, then +evaluates the recorded expression over the real data (this is *tracing*). So a +plain `if` on column values has nothing to look at and raises `ValueError: The +truth value of an array ... is ambiguous`. Use `np.where`, or write the +function so it compiles instead — see below. + +**Real `if`/`for`/`while` works, if the function fits the DSL.** When your +function branches or loops over column values *and* fits Blosc2's +[DSL grammar](../reference/dsl_syntax.md), it is compiled and runs as written, +branches and all — that is what makes the Kepler kernel above possible. If it +doesn't fit the grammar, it silently falls back to tracing (hence the +`ValueError`); if it looks DSL-shaped but calls an unsupported function, you +get a `RuntimeError` naming it. + +**`np.where` evaluates both arms.** Both branches are computed over the whole +column before one is selected, so each runs on values it was never meant to +see — negative bases, divisions by zero, `RuntimeWarning`s. The answer is +still correct, but clamp each arm to its own domain (as `yeo_johnson` does +with `np.maximum`) to keep the noise and wasted work away. + +**Don't put a reduction inside per-element control flow.** In a DSL kernel, +`sum`/`max`/`min` collapse the whole block being evaluated to a single value, +not one per row. So `if max(abs(diff)) < 1e-12: break` compiles, runs, and +silently gives wrong results for every row but the first. Write the +per-element form: `if abs(diff) < 1e-12: break`. + +**`std()` changes meaning.** A plain `apply` passes a pandas Series +(`ddof=1`); the engine passes a NumPy array (`ddof=0`). Pass `ddof` explicitly +if it matters. -**`np.where` evaluates both arms.** Unlike a real `if`, both branches are -computed over the whole column and only then selected between, so each one runs -on values it was never meant to see. In `yeo_johnson` above, `np.power(col + 1.0, -0.5)` would hit a negative base wherever `col < -1` — around 159,000 elements in -a million-row standard normal — producing NaNs and a -`RuntimeWarning: invalid value encountered in power`. The final answer is still -correct, because `np.where` discards them, but the warning is noise and the work -is wasted. That is why each arm is clamped to its own domain with `np.maximum`. -The same caveat applies to numexpr's `where()`. +## Limitations -**`np.sign` is not supported** on the traced expressions and raises a -`TypeError`. Express it with `np.where` instead. +- Numeric dtypes only; anything else raises `ValueError`. +- `na_action="ignore"` is not supported for `map` (`NotImplementedError`): + there is no per-element step at which to skip a value. +- Only `DataFrame.apply` and `Series.map` reach the engine. pandas 3's + `Series.apply` doesn't accept `engine=` for non-string functions, and + `DataFrame.map` doesn't forward it — both are pandas-side limits. +- `engine=` always gets auto-detection: pandas' protocol requires the plain + `blosc2.jit` object, so `strict=True`/`strict=False` cannot be passed + through it (a direct call can). -## Row-wise (`axis=1`) +## Details -`axis=0`, the default, calls the function once per column and is where the -benefit lies. `axis=1` calls it once per row — the same Python-level loop plain -pandas would run, except each call now also pays the cost of wrapping a tiny -array for the compute engine. For a handful of columns that overhead outweighs -any gain, making `engine=blosc2.jit` with `axis=1` typically *slower* than a -plain `apply(axis=1)`. +Two things worth knowing once you are actually using this, neither needed to +get started. -Use `axis=0`, or restructure the computation so it works column-wise. +**Against numexpr.** pandas' own +[Enhancing performance](https://pandas.pydata.org/docs/user_guide/enhancingperf.html) +guide describes `pd.eval(..., engine="numexpr")`, which fuses expressions on +the same principle. On the Yeo-Johnson example: -## Limitations - -- Only numeric dtypes are supported. A non-numeric (e.g. object-dtype or - string) column raises a `ValueError` naming the limitation rather than - attempting the computation. -- `na_action="ignore"` is not supported for `map` and raises - `NotImplementedError` — the vectorized-call contract means there is no - per-element step at which to skip a value. -- `Series.apply(func, engine=...)` and `DataFrame.map(func, engine=...)` do - not reach `blosc2.jit` at all: pandas 3's `Series.apply` does not accept - an `engine` keyword for non-string functions, and `DataFrame.map` doesn't - forward `engine` to a dispatch mechanism at all. These are limitations of - the pandas-side API surface, not of the Blosc2 engine. The two entry - points that do reach the engine are `DataFrame.apply` and `Series.map`. - -`Series.map(func, engine=blosc2.jit)` works the same way as `DataFrame.apply`: -`func` is called once with the Series' full underlying array. +| approach | time | vs plain apply | +| --- | --- | --- | +| plain `df.apply(f)` | 0.121 s | 1.00x | +| `df.apply(f, engine=blosc2.jit)` | 0.058 s | 2.07x | +| `numexpr.evaluate(...)` per column | 0.039 s | 3.09x | + +numexpr is faster on an in-memory frame, so the trade is not raw speed: its +version of `yeo_johnson` has to collapse into one quoted string, +`"where(c >= 0, ((maximum(c, 0.0) + 1.0) ** lam - 1.0) / lam, ...)"`, while +the Blosc2 version stays a Python function — intermediate variables, a name, +helper calls, unit tests, and your linter. Note also that neither +`engine=blosc2.jit` nor numexpr touches compressed data here: pandas +materialises each column as a plain NumPy array first. Only the direct-call +pattern reaches `blosc2.NDArray` operands. + +**What column extraction costs.** `df["col"]` is a zero-copy view for ordinary +numeric dtypes, including in mixed-dtype frames: pulling out one column never +triggers the whole-frame upcast that `df.to_numpy()` does (which is what +`engine=blosc2.jit` uses internally for `axis=0`), and each column keeps its +own dtype. Nullable (`Float64`/`Int64`), Arrow-backed and tz-aware columns are +the exception — they allocate a converted array once, which is noise next to +iterative work like the Kepler kernel. ## Reproducing these numbers All figures on this page come from `bench/bench_pandas_engine.py`, measured on -an Apple M4 with pandas 3.0.3 and 8 threads. Run it yourself with: +an Apple M4 with pandas 3.0.3 and 8 threads: ``` python bench/bench_pandas_engine.py ``` -It prints the comparison table, both sweeps, and regenerates the plot above. +It prints every table on this page and regenerates both plots. diff --git a/doc/guides/pandas_engine/kepler.png b/doc/guides/pandas_engine/kepler.png new file mode 100644 index 000000000..74dcdd413 Binary files /dev/null and b/doc/guides/pandas_engine/kepler.png differ diff --git a/doc/guides/pandas_engine/speedup.png b/doc/guides/pandas_engine/speedup.png index b26d37a3e..c86cfda43 100644 Binary files a/doc/guides/pandas_engine/speedup.png and b/doc/guides/pandas_engine/speedup.png differ diff --git a/doc/reference/dsl_syntax.md b/doc/reference/dsl_syntax.md index 3a5acafb7..82097546d 100644 --- a/doc/reference/dsl_syntax.md +++ b/doc/reference/dsl_syntax.md @@ -18,6 +18,15 @@ def kernel(x, y): Use Python-style indentation and always return a value on the paths you execute. +`@blosc2.jit` auto-detects this DSL: a decorated function whose body contains an +`if`/`for`/`while` and that compiles under this grammar is dispatched here +automatically, so its branches and loops actually run, once per chunk, instead +of `jit`'s normal approach of calling the function only once to record a single +expression — which would otherwise capture just whichever branch that one call +happened to take, silently dropping the rest. `@blosc2.dsl_kernel` remains the +explicit form — it always requires the DSL to compile, equivalent to +`jit(strict=True)`. + ## Program shape - Exactly one top-level `def ...:` function is expected. @@ -169,6 +178,21 @@ Rules: - `while` condition is a regular DSL expression. - Runtime iteration cap is enforced by `ME_DSL_WHILE_MAX_ITERS`. +### Reductions inside control flow + +`sum`, `max`, `min` and other block reductions collapse the whole chunk being +evaluated to one value, not one value per element. Using one as the +condition of `if`/`while`, or assigning one to a local that per-element code +later reads (e.g. `y = max(x)` followed by `if x > 0: y = y + 1`), does +**not** raise a compile-time or runtime error -- it compiles and runs, and +produces results that are only correct for element 0 of the block; every +other element sees a stale/zero value where the reduction result should be. +This is a rough edge in the underlying +[miniexpr](https://github.com/Blosc/miniexpr) compiler, not something this +Python layer validates today. Write the per-element form instead (drop the +reduction, e.g. `if abs(diff) < tol` rather than `if max(abs(diff)) < tol`) +whenever the intent is a per-element, not whole-block, decision. + ## `print(...)` `print` is supported as a DSL statement. @@ -271,3 +295,6 @@ These Python features are not part of this DSL: - Ternary expression: `a if cond else b` - `for ... else` and `while ... else` - Keyword-argument calls and other call forms outside the supported subset +- Docstrings (or any other bare string-literal statement) inside the kernel + body -- this is a compile-time parse error at the miniexpr level, not a + silently-ignored statement. diff --git a/src/blosc2/blosc2_ext.pyx b/src/blosc2/blosc2_ext.pyx index 1ddce5a82..ed325a565 100644 --- a/src/blosc2/blosc2_ext.pyx +++ b/src/blosc2/blosc2_ext.pyx @@ -801,6 +801,8 @@ ctypedef struct me_input_cache_s: ctypedef struct me_udata: b2nd_array_t** inputs + uint8_t** np_data # per-input raw base pointer; NULL entry = b2nd input + int32_t* np_typesizes # per-input itemsize; valid where np_data[i] != NULL me_input_cache_s* input_chunk_caches int ninputs me_eval_params* eval_params @@ -2409,6 +2411,10 @@ cdef class SChunk: free(me_data.input_chunk_caches) if me_data.inputs != NULL: free(me_data.inputs) + if me_data.np_data != NULL: + free(me_data.np_data) + if me_data.np_typesizes != NULL: + free(me_data.np_typesizes) if me_data.miniexpr_handle != NULL: # XXX do we really need the conditional? me_free(me_data.miniexpr_handle) if me_data.eval_params != NULL: @@ -2515,6 +2521,15 @@ cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock, cdef int64_t start_ndim[B2ND_MAX_DIM] cdef int64_t stop_ndim[B2ND_MAX_DIM] cdef int64_t buffershape[B2ND_MAX_DIM] + # Raw-NumPy-input gather (odometer copy of a block out of a C-order buffer) + cdef int64_t counter[B2ND_MAX_DIM] + cdef int64_t shape_strides[B2ND_MAX_DIM] + cdef int64_t blockshape_strides[B2ND_MAX_DIM] + cdef int32_t np_ts + cdef int np_ndim + cdef c_bool all_pad + cdef int64_t ext, row_items, row_bytes, src_flat, dst_flat + cdef int dd cdef b2nd_array_t* ndarr cdef int rc @@ -2566,6 +2581,68 @@ cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock, return 0 for i in range(udata.ninputs): + if udata.np_data != NULL and udata.np_data[i] != NULL: + # Raw NumPy input: gather block (nchunk, nblock) from a C-order buffer. + # All geometry comes from the output array, valid because the Python + # gates guarantee every operand shares the output's shape and grid. + np_ts = udata.np_typesizes[i] + blocknitems = udata.array.blocknitems + block_nbytes = blocknitems * np_ts + if expected_blocknitems == -1: + expected_blocknitems = blocknitems + elif blocknitems != expected_blocknitems: + raise ValueError("miniexpr: inconsistent block element counts across inputs") + input_buffers[i] = malloc(block_nbytes) + if input_buffers[i] == NULL: + raise MemoryError("miniexpr: cannot allocate input block buffer") + memset(input_buffers[i], 0, block_nbytes) # zero padding, matches b2nd semantics + + np_ndim = udata.array.ndim + blosc2_unidim_to_multidim(np_ndim, udata.chunks_in_array, nchunk, chunk_ndim) + blosc2_unidim_to_multidim(np_ndim, udata.blocks_in_chunk, nblock, block_ndim) + + all_pad = False + for dd in range(np_ndim): + start_ndim[dd] = chunk_ndim[dd] * udata.array.chunkshape[dd] + block_ndim[dd] * udata.array.blockshape[dd] + ext = udata.array.shape[dd] - start_ndim[dd] + stop_ndim[dd] = udata.array.blockshape[dd] if ext >= udata.array.blockshape[dd] else ext + if stop_ndim[dd] <= 0: + all_pad = True # fully-padded block: buffer stays zeroed + + if not all_pad: + # C-order element strides for the full array shape and for the block shape. + shape_strides[np_ndim - 1] = 1 + blockshape_strides[np_ndim - 1] = 1 + for dd in range(np_ndim - 2, -1, -1): + shape_strides[dd] = shape_strides[dd + 1] * udata.array.shape[dd + 1] + blockshape_strides[dd] = blockshape_strides[dd + 1] * udata.array.blockshape[dd + 1] + + row_items = stop_ndim[np_ndim - 1] + row_bytes = row_items * np_ts + for dd in range(np_ndim): + counter[dd] = 0 + # Odometer over the outer np_ndim-1 dims; the innermost dim is + # copied in bulk per row (1-D collapses to a single memcpy). + while True: + src_flat = 0 + dst_flat = 0 + for dd in range(np_ndim): + src_flat += (start_ndim[dd] + counter[dd]) * shape_strides[dd] + dst_flat += counter[dd] * blockshape_strides[dd] + memcpy( input_buffers[i] + dst_flat * np_ts, + udata.np_data[i] + src_flat * np_ts, + row_bytes) + dd = np_ndim - 2 + while dd >= 0: + counter[dd] += 1 + if counter[dd] < stop_ndim[dd]: + break + counter[dd] = 0 + dd -= 1 + else: + break + continue + ndarr = udata.inputs[i] if ndarr.sc.storage.urlpath == NULL: src = ndarr.sc.data[nchunk] @@ -2651,11 +2728,10 @@ cdef int aux_miniexpr(me_udata *udata, int64_t nchunk, int32_t nblock, if rc < 0: raise ValueError("miniexpr: error decompressing the chunk") # For reduction operations, we need to track which block we're processing - # The linear_block_index should be based on the INPUT array structure, not the output array - # Get the first input array's chunk and block structure - cdef b2nd_array_t* first_input = udata.inputs[0] + # The linear_block_index should be based on the same grid the output shares + # with every input (raw NumPy inputs have no b2nd_array_t to read ndim from). cdef int nblocks_per_chunk = 1 - for i in range(first_input.ndim): + for i in range(udata.array.ndim): nblocks_per_chunk *= udata.blocks_in_chunk[i] # Calculate the global linear block index: nchunk * blocks_per_chunk + nblock # This works because blocks never span chunks (chunks are padded to block boundaries) @@ -3913,6 +3989,8 @@ cdef class NDArray: cdef me_udata *udata = calloc(1, sizeof(me_udata)) cdef me_eval_params* eval_params cdef b2nd_array_t** inputs_ + cdef uint8_t** np_data + cdef int32_t* np_typesizes cdef me_input_cache_s* input_chunk_caches cdef void* aux_reduc_ptr = NULL cdef int i @@ -3925,14 +4003,32 @@ cdef class NDArray: if udata == NULL: raise MemoryError("Cannot allocate miniexpr user data") inputs_ = NULL + np_data = NULL + np_typesizes = NULL if ninputs > 0: inputs_ = malloc(ninputs * sizeof(b2nd_array_t*)) if inputs_ == NULL: free(udata) raise MemoryError("Cannot allocate miniexpr input table") + np_data = calloc(ninputs, sizeof(uint8_t*)) + np_typesizes = calloc(ninputs, sizeof(int32_t)) + if np_data == NULL or np_typesizes == NULL: + free(inputs_) + free(np_data) + free(np_typesizes) + free(udata) + raise MemoryError("Cannot allocate miniexpr raw-input tables") for i, operand in enumerate(operands): - inputs_[i] = operand.c_array + if isinstance(operand, np.ndarray): + # Caller (fast_eval) guarantees C-contiguous, native-endian, non-scalar. + inputs_[i] = NULL + np_data[i] = np.PyArray_DATA( operand) + np_typesizes[i] = ( operand).itemsize + else: + inputs_[i] = operand.c_array udata.inputs = inputs_ + udata.np_data = np_data + udata.np_typesizes = np_typesizes udata.ninputs = ninputs input_chunk_caches = NULL if ninputs > 0: diff --git a/src/blosc2/dsl_kernel.py b/src/blosc2/dsl_kernel.py index 85dcb0de0..82c11b373 100644 --- a/src/blosc2/dsl_kernel.py +++ b/src/blosc2/dsl_kernel.py @@ -16,6 +16,8 @@ from io import StringIO from typing import ClassVar +import numpy + _PRINT_DSL_KERNEL = os.environ.get("PRINT_DSL_KERNEL", "").strip().lower() _PRINT_DSL_KERNEL = _PRINT_DSL_KERNEL not in ("", "0", "false", "no", "off") _DSL_USAGE_DOC_URL = "https://github.com/Blosc/python-blosc2/blob/main/doc/reference/dsl_syntax.md" @@ -25,6 +27,43 @@ class DSLSyntaxError(ValueError): """Raised when a @dsl_kernel function uses unsupported DSL syntax.""" +# NumPy function names that the DSL grammar recognizes only under a different +# (but semantically identical, same-arity) name. Verified individually against +# the NumPy function they alias -- not a general "closest match" mapping, so +# names with subtle semantic differences (e.g. `np.mod`/`np.remainder`'s sign +# convention vs C's `fmod`) are deliberately left out. +_NUMPY_TO_DSL_FUNC_ALIASES = { + "maximum": "fmax", + "minimum": "fmin", + "absolute": "abs", +} + + +class _NumpyAttrCallRewriter(ast.NodeTransformer): + """Rewrite `alias.foo(...)` calls to the bare `foo(...)` form the DSL grammar + requires, for every *alias* bound to the real NumPy module. Also applies + `_NUMPY_TO_DSL_FUNC_ALIASES` for the handful of functions the DSL knows + under a different name. + """ + + def __init__(self, aliases: set[str]): + self._aliases = aliases + self.rewrote_any = False + + def visit_Call(self, node: ast.Call) -> ast.AST: + self.generic_visit(node) + func = node.func + if ( + isinstance(func, ast.Attribute) + and isinstance(func.value, ast.Name) + and func.value.id in self._aliases + ): + dsl_name = _NUMPY_TO_DSL_FUNC_ALIASES.get(func.attr, func.attr) + node.func = ast.copy_location(ast.Name(id=dsl_name, ctx=ast.Load()), func) + self.rewrote_any = True + return node + + def _normalize_miniexpr_scalar(value): # NumPy scalar-like values expose .item(); plain Python scalars do not. # Do not call .item() on non-scalar arrays; for Blosc2 arrays this can be expensive. @@ -338,8 +377,6 @@ def _args(self, func_node: ast.FunctionDef): args = func_node.args if args.vararg or args.kwarg or args.kwonlyargs: self._err(args, "DSL kernel does not support *args/**kwargs/kwonly args") - if args.defaults or args.kw_defaults: - self._err(args, "DSL kernel does not support default arguments") def _check_input_assign(self, target: ast.Name): # G2: miniexpr forbids reassigning an input parameter (inputs alias operand buffers). @@ -531,9 +568,10 @@ def __init__(self, func): dsl_source = None input_names = None self.dsl_error = e - except Exception: + except Exception as e: dsl_source = None input_names = None + self.dsl_error = e self.dsl_source = dsl_source self.input_names = input_names @@ -560,6 +598,11 @@ def _extract_dsl(self, func, validate: bool = True): if dsl_func is None: raise ValueError("No function definition found in sliced DSL source") input_names = self._input_names_from_signature(dsl_func) + + dsl_source, dsl_tree, dsl_func = self._rewrite_numpy_attr_calls( + func, dsl_source, dsl_tree, dsl_func, input_names + ) + if validate: DSLValidator(dsl_source, input_names=input_names).validate(dsl_func) if _PRINT_DSL_KERNEL: @@ -568,6 +611,34 @@ def _extract_dsl(self, func, validate: bool = True): print(dsl_source) return dsl_source, input_names + @staticmethod + def _rewrite_numpy_attr_calls(func, dsl_source, dsl_tree, dsl_func, input_names): + """Rewrite `np.foo(...)` calls to bare `foo(...)`, for every name in *func*'s + defining scope that is bound to the real NumPy module (typically `np`, but + any alias, including a bare `numpy` import, is honored). The DSL grammar + only accepts bare function-name calls. No-op, returning the inputs + unchanged, when there is nothing to rewrite (including when the alias is + shadowed by one of the kernel's own parameter names). + """ + aliases = { + name + for name, value in getattr(func, "__globals__", {}).items() + if value is numpy and name not in input_names + } + if not aliases: + return dsl_source, dsl_tree, dsl_func + + rewriter = _NumpyAttrCallRewriter(aliases) + rewritten = rewriter.visit(ast.parse(dsl_source)) + if not rewriter.rewrote_any: + return dsl_source, dsl_tree, dsl_func + + ast.fix_missing_locations(rewritten) + new_source = ast.unparse(rewritten) + new_tree = ast.parse(new_source) + new_func = next((node for node in new_tree.body if isinstance(node, ast.FunctionDef)), None) + return new_source, new_tree, new_func + @staticmethod def _slice_function_source(source: str, func_node: ast.FunctionDef) -> str: lines = source.splitlines() @@ -584,8 +655,6 @@ def _input_names_from_signature(func_node: ast.FunctionDef) -> list[str]: args = func_node.args if args.vararg or args.kwarg or args.kwonlyargs: raise ValueError("DSL kernel does not support *args/**kwargs/kwonly args") - if args.defaults or args.kw_defaults: - raise ValueError("DSL kernel does not support default arguments") return [a.arg for a in (args.posonlyargs + args.args)] def __call__(self, inputs_tuple, output, offset=None): @@ -614,7 +683,13 @@ def __call__(self, inputs_tuple, output, offset=None): def dsl_kernel(func): - """Decorator to wrap a function in a DSLKernel.""" + """Decorator to wrap a function in a DSLKernel. + + Default argument values in *func*'s signature are accepted as ordinary named + inputs, but they are honored (filled in when omitted) only through the + ``@blosc2.jit`` call path. Calling a :class:`DSLKernel` directly (e.g. via + :func:`lazyudf`) requires every input to be passed positionally. + """ return DSLKernel(func) diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index cd9cd29e4..515c7cd4e 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -1495,6 +1495,15 @@ def _js_dtypes_ok(operands, kwargs) -> bool: ) +def _trace_js_backend(expression): + """BLOSC_ME_JIT_TRACE counterpart for the JS bridge, which never reaches + miniexpr's trace point in `fast_eval` (see there for the message format).""" + if os.environ.get("BLOSC_ME_JIT_TRACE", "").lower() in ("1", "true", "on"): + source = getattr(expression, "dsl_source", None) or expression + expr_short = str(source)[:120].replace("\n", " ") + print(f"[blosc2] engine=js expr={expr_short}", flush=True) + + def _maybe_js_backend(expression, jit, jit_backend, reduce_args, operands, kwargs, shape=None): """Resolve the JS backend for a DSL kernel. @@ -1524,7 +1533,9 @@ def _maybe_js_backend(expression, jit, jit_backend, reduce_args, operands, kwarg 'jit_backend="js" requires a floating-point output dtype ' f"(got {np.dtype(out_dtype)}); drop jit_backend to use miniexpr" ) - return _as_js_udf(expression, shape), None, None + bridge = _as_js_udf(expression, shape) + _trace_js_backend(expression) + return bridge, None, None prefer_js = ( jit is not False # jit=True/None prefer the best JIT (js); only jit=False forces interpreter and jit_backend is None @@ -1541,6 +1552,7 @@ def _maybe_js_backend(expression, jit, jit_backend, reduce_args, operands, kwarg bridge = _as_js_udf(expression, shape) # transpiles; raises on any unsupported construct except Exception: return expression, jit, jit_backend # fall back to miniexpr, no regression + _trace_js_backend(expression) return bridge, None, None @@ -1735,22 +1747,36 @@ def fast_eval( # noqa: C901 expr_string_miniexpr = _apply_jit_backend_pragma( expr_string_miniexpr, operands_miniexpr, jit_backend ) - all_ndarray_miniexpr = all( - isinstance(value, blosc2.NDArray) and value.shape != () for value in operands_miniexpr.values() - ) - # Require aligned NDArray operands with identical chunk/block grid. - same_shape = all(hasattr(op, "shape") and op.shape == shape for op in operands_miniexpr.values()) - same_chunks = all(hasattr(op, "chunks") and op.chunks == chunks for op in operands_miniexpr.values()) - same_blocks = all(hasattr(op, "blocks") and op.blocks == blocks for op in operands_miniexpr.values()) - if not (same_shape and same_chunks and same_blocks): + + def _miniexpr_eligible_operand(op): + if isinstance(op, blosc2.NDArray): + return op.shape != () and op.shape == shape and op.chunks == chunks and op.blocks == blocks + if isinstance(op, np.ndarray): + # Raw NumPy operands are only miniexpr-eligible for DSL kernels (the + # jit control-flow dispatch route). Plain string expressions and traced + # `jit` calls keep the old NDArray-only gate, so their numeric behavior + # (e.g. transcendentals matching numexpr bit-for-bit) is unchanged. + return ( + is_dsl + and op.ndim > 0 + and op.shape == shape + and op.dtype.isnative + and op.dtype.kind in "biufc" + ) + return False + + all_eligible_miniexpr = all(_miniexpr_eligible_operand(op) for op in operands_miniexpr.values()) + if not all_eligible_miniexpr: use_miniexpr = False if is_dsl and dsl_disable_reason is None: - dsl_disable_reason = "all DSL operands must share shape/chunks/blocks." - if not (all_ndarray_miniexpr and out is None): + dsl_disable_reason = ( + "all DSL operands must be NDArray or NumPy inputs sharing shape/chunks/blocks." + ) + if not (all_eligible_miniexpr and out is None): use_miniexpr = False if is_dsl and dsl_disable_reason is None: dsl_disable_reason = ( - "DSL kernels require NDArray inputs and do not support the `out` argument." + "DSL kernels require NDArray or NumPy inputs and do not support the `out` argument." ) has_complex = any( isinstance(op, blosc2.NDArray) and blosc2.isdtype(op.dtype, "complex floating") @@ -1779,7 +1805,11 @@ def fast_eval( # noqa: C901 print(f"[blosc2] engine={engine} {jit_info} expr={expr_short}", flush=True) if use_miniexpr: - cparams = kwargs.pop("cparams", blosc2.CParams()) + cparams = kwargs.pop("cparams", None) + if cparams is None: + # getitem output is throwaway scratch (returned as a NumPy array and + # discarded), so compressing it buys nothing but a round trip. + cparams = blosc2.CParams(clevel=0) if getitem else blosc2.CParams() # All values will be overwritten, so we can use an uninitialized array res_eval = blosc2.uninit(shape, dtype, chunks=chunks, blocks=blocks, cparams=cparams, **kwargs) prefilter_set = False @@ -1787,6 +1817,11 @@ def fast_eval( # noqa: C901 # Fuse where(cond, x, y) into the expression for miniexpr _pref_expr = expr_string_miniexpr _pref_ops = operands_miniexpr + if any(isinstance(v, np.ndarray) for v in _pref_ops.values()): + _pref_ops = { + k: (np.ascontiguousarray(v) if isinstance(v, np.ndarray) else v) + for k, v in _pref_ops.items() + } if where is not None and len(where) == 2: _pref_expr = f"where({_pref_expr}, _where_x, _where_y)" # _cb_anchor keeps the contiguous bitmap alive for the whole @@ -1883,7 +1918,11 @@ def fast_eval( # noqa: C901 if callable(expression): if _is_dsl_kernel_expression(expression): _raise_dsl_miniexpr_required( - "internal fallback attempted to execute the DSL kernel directly in Python." + "DSL kernels require the miniexpr fast path, and it was unavailable for this " + "evaluation. Common causes: operands with mismatched chunks/blocks or " + "non-contiguous/non-native-byte-order NumPy arrays, an explicit `out=` argument, " + "or a `where(cond, x, y)` with cardinality-changing semantics (len(where) == 1) — " + "all unsupported for DSL kernels." ) if _in_place: expression(tuple(chunk_operands.values()), out, offset=offset) @@ -2219,7 +2258,9 @@ def slices_eval( # noqa: C901 if callable(expression): if _is_dsl_kernel_expression(expression): _raise_dsl_miniexpr_required( - "internal sliced fallback attempted to execute the DSL kernel directly in Python." + "DSL kernels only support evaluating the full array (compute() or [()]) through " + "the miniexpr fast path; slicing a DSL computation (e.g. `lexpr[1:5]`) is not " + "supported." ) if _in_place: # presumably the user knows what they're doing # edit out in-place @@ -2373,7 +2414,8 @@ def slices_eval_getitem( if callable(expression): if _is_dsl_kernel_expression(expression): _raise_dsl_miniexpr_required( - "internal getitem fallback attempted to execute the DSL kernel directly in Python." + "DSL kernels only support evaluating the full array (compute() or [()]) through " + "the miniexpr fast path; sliced getitem (e.g. `lexpr[1:5]`) is not supported." ) offset = tuple(0 if s is None else s.start for s in _slice_bcast) # offset for the udf if _in_place: @@ -2775,7 +2817,8 @@ def reduce_slices( # noqa: C901 if callable(expression): if _is_dsl_kernel_expression(expression): _raise_dsl_miniexpr_required( - "internal reduction fallback attempted to execute the DSL kernel directly in Python." + "DSL kernels do not support reductions (e.g. sum()/mean() over an axis); write " + "the reduction outside the kernel, over its computed result." ) # TODO: Implement the reductions for UDFs (and test them) result = np.empty(cslice_shape, dtype=out.dtype) @@ -3018,7 +3061,7 @@ def _eval_zero_input_dsl_if_needed( return True, full_res -def chunked_eval( +def chunked_eval( # noqa: C901 expression: str | Callable[[tuple, np.ndarray, tuple[int]], None], operands: dict, item=(), **kwargs ): """ @@ -3120,6 +3163,19 @@ def chunked_eval( return slices_eval(expression, operands, getitem=getitem, _slice=item, shape=shape, **kwargs) fast_path = full_slice and fast_path + if not fast_path and full_slice and _is_dsl_kernel_expression(expression) and operands: + # All-NumPy DSL operands: validate_inputs only sees NDinputs, so it never + # sets fast_path for this case; reroute it here instead. Scalar operands + # (e.g. a Python int/float parameter) don't participate in the shape/grid + # check, mirroring validate_inputs' own raw_inputs filtering. + raw_ops = [v for v in operands.values() if not _isscalar(v)] + if ( + raw_ops + and all(isinstance(v, np.ndarray) and v.ndim > 0 for v in raw_ops) + and len({v.shape for v in raw_ops}) == 1 + ): + fast_path = True + if fast_path: # necessarily item is () if getitem: # When using getitem, taking the fast path is always possible @@ -4532,12 +4588,45 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None, ne_args=No return new_expr +def _align_dsl_operand_grids(inputs): + """Put every NDArray operand of a DSL kernel on a single chunks/blocks grid. + + Blocks are sized in bytes, so same-shaped operands of different itemsize get + different grids by default (float32 vs int64 over 1M elements: blocks of + 31250 vs 15625 elements). miniexpr needs one common grid, and a DSL kernel + has no slow path to fall back on, so evaluation would fail outright with a + confusing "slicing is not supported" error. Copy the odd operands onto the + grid of the widest dtype: its blocks hold the fewest elements, so every + operand still fits within the cache budget the heuristic aimed at. + + The copies happen once, at construction, and only when the grids actually + disagree -- whether they do depends on the array size and on the platform's + cache detection, which is why this used to fail only on some CI runners. + """ + nd = [x for x in inputs if isinstance(x, blosc2.NDArray) and x.ndim > 0] + if len(nd) < 2 or len({(x.shape, x.chunks, x.blocks) for x in nd}) < 2: + return inputs + ref = max(nd, key=lambda x: x.dtype.itemsize) + aligned = [] + for x in inputs: + misaligned = ( + isinstance(x, blosc2.NDArray) + and x.ndim > 0 + and x.shape == ref.shape + and (x.chunks, x.blocks) != (ref.chunks, ref.blocks) + ) + aligned.append(x.copy(chunks=ref.chunks, blocks=ref.blocks) if misaligned else x) + return aligned + + class LazyUDF(LazyArray): def __init__( self, func, inputs, dtype, shape=None, chunked_eval=True, jit=None, jit_backend=None, **kwargs ): # After this, all the inputs should be np.ndarray or NDArray objects self.inputs = convert_inputs(inputs) + if isinstance(func, DSLKernel): + self.inputs = _align_dsl_operand_grids(self.inputs) # Get res shape if shape is None: self._shape = compute_broadcast_shape(self.inputs) @@ -4576,7 +4665,16 @@ def __init__( # DSL kernels are using input names that are extracted from params as a list, # and we need to use them for matching variables in miniexpr # (instead of the 'o{%d}' notation). - self.inputs_dict = dict(zip(self.func.input_names, self.inputs, strict=True)) + names = self.func.input_names + if len(names) != len(self.inputs): + # Otherwise this surfaces as a bare "zip() argument 2 is longer + # than argument 1", which names neither the kernel nor the counts. + udf_name = getattr(self.func.func, "__name__", self.func.__name__) + raise ValueError( + f"DSL kernel {udf_name!r} takes {len(names)} operand(s) " + f"({', '.join(names)}), but {len(self.inputs)} were passed." + ) + self.inputs_dict = dict(zip(names, self.inputs, strict=True)) else: self.inputs_dict = {f"o{i}": obj for i, obj in enumerate(self.inputs)} diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 4d5306688..ad661353f 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -116,7 +116,12 @@ np.floor: "floor", np.ceil: "ceil", np.trunc: "trunc", + np.sign: "sign", np.signbit: "signbit", + np.square: "square", + np.negative: "negative", + np.positive: "positive", + np.reciprocal: "reciprocal", np.round: "round", } diff --git a/src/blosc2/proxy.py b/src/blosc2/proxy.py index 7f346d36d..a827f00b1 100644 --- a/src/blosc2/proxy.py +++ b/src/blosc2/proxy.py @@ -5,7 +5,10 @@ # SPDX-License-Identifier: BSD-3-Clause ####################################################################### +import ast import asyncio +import inspect +import textwrap from abc import ABC, abstractmethod from collections.abc import Sequence @@ -18,11 +21,19 @@ import numpy as np import blosc2 +from blosc2.dsl_kernel import DSLKernel, DSLSyntaxError # Default Proxy.afetch concurrency cap for remote sources (e.g. C2Array), # where fetches are dominated by round-trip latency, not local CPU/IO. REMOTE_MAX_CONCURRENCY = 8 +# `jit` kwargs that tune *how* an expression is evaluated, not what container the +# result is stored in. Unlike storage kwargs (`cparams`, `chunks`, `urlpath`, ...), +# these must not by themselves flip the return type from a plain NumPy array to +# an NDArray -- wanting a faster JIT backend has nothing to do with wanting a +# compressed/persisted container back. +_JIT_EXECUTION_TUNING_KWARGS = frozenset({"jit", "jit_backend", "fp_accuracy"}) + class ProxyNDSource(ABC): """ @@ -751,28 +762,302 @@ def as_simpleproxy(*arrs: Sequence[blosc2.Array]) -> tuple[SimpleProxy | blosc2. return out[0] if len(out) == 1 else out -def jit(func=None, *, out=None, disable=False, **kwargs): +class _PandasRowProxy(blosc2.Operand): + """Row proxy for `PandasUdfEngine.apply`'s axis=1 route. + + Stands in for "the current row" the way the textbook `axis=1` idiom + expects (`row["colname"]`), but is backed by whole *columns*: `row["a"] + + row["b"]` traces to one fused expression over the whole column set in + a single call, instead of looping over rows in Python. Columns are + extracted lazily (and cached) from the original DataFrame, not from a + whole-frame NumPy array, so per-column dtypes are preserved. + """ + + def __init__(self, df): + self._df = df + self._cache = {} + + def __getitem__(self, key): + if not isinstance(key, str): + raise TypeError( + f"row[{key!r}]: axis=1 row proxies only support column access by " + "name (a string). Positional or iterable row access is not " + "supported; for row-wise computations, call your @blosc2.jit " + "function directly with the DataFrame columns as separate " + "arguments instead, e.g. func(df['a'], df['b'])." + ) + if key in self._cache: + return self._cache[key] + n_matches = int((self._df.columns == key).sum()) + if n_matches == 0: + raise KeyError(f"row[{key!r}]: no such column in the DataFrame") + if n_matches > 1: + raise KeyError( + f"row[{key!r}]: column label is duplicated ({n_matches} matches); " + "axis=1 row proxies require unique column labels" + ) + col = self._df[key].to_numpy() + if col.dtype.kind not in "biufc": + raise ValueError( + f"row[{key!r}]: column has dtype {col.dtype!r}, which is not numeric. " + "The Blosc2 engine only supports vectorized numeric computations." + ) + proxy = SimpleProxy(col) + self._cache[key] = proxy + return proxy + + def __getattr__(self, name): + raise AttributeError( + f"row.{name}: axis=1 row proxies only support column access via " + f"row[{name!r}]; attribute access, iteration and per-row methods " + "(e.g. row.isna()) are not supported. For per-row computations that " + "need more than combining columns (e.g. per-row branching), call " + "your @blosc2.jit function directly with the columns as separate " + "array arguments instead of through df.apply(..., axis=1)." + ) + + +def _undecorated(func): + """The original function behind a @blosc2.jit wrapper, or *func* itself. + + Source inspection has to see what the user wrote, not the wrapper. + """ + return getattr(func, "_blosc2_jit_wrapped", func) + + +def _decorate_once(func, decorator): + """Apply *decorator* unless *func* is already a @blosc2.jit wrapper. + + Decorating twice used to break the DSL route: the outer (tracing) wrapper + replaces array arguments with SimpleProxy operands, so the inner DSL kernel + saw no array at all and failed asking for `shape=`. + """ + return func if hasattr(func, "_blosc2_jit_wrapped") else decorator(func) + + +def _analyze_row_func(func) -> tuple[bool, bool]: + """Inspect *func* for the two signals `PandasUdfEngine.apply`'s axis=1 + route needs to pick a dispatch strategy: whether it subscripts its first + parameter with a string literal anywhere in its body (the `row["colname"]` + idiom), and whether its body contains a `for`/`while` loop. + + Both default to False (the historical per-row loop) if the source can't + be inspected, e.g. a dynamically built function. + """ + try: + sig = inspect.signature(func) + params = list(sig.parameters.values()) + if not params: + return False, False + row_name = params[0].name + source = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(source) + except (OSError, TypeError, SyntaxError, ValueError): + return False, False + nodes = list(ast.walk(tree)) + uses_subscript = any( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Name) + and node.value.id == row_name + and isinstance(node.slice, ast.Constant) + and isinstance(node.slice.value, str) + for node in nodes + ) + has_loop = any(isinstance(node, ast.For | ast.While) for node in nodes) + return uses_subscript, has_loop + + +def _has_control_flow(source: str | None) -> bool: + """Whether *source* (a DSL-extracted function source, or None) contains a + branch or loop that tracing cannot observe.""" + if source is None: + return False + tree = ast.parse(source) + return any(isinstance(node, ast.If | ast.For | ast.While) for node in ast.walk(tree)) + + +def _wide_frame_hint(err: BaseException, func_name: str, params) -> str | None: + """Guidance to append when a call gets a keyword the function doesn't take. + + The usual cause is the `kernel(**df)` idiom (see doc/guides/pandas_engine.md) + against a frame carrying more columns than the kernel has parameters. Extra + keywords are rejected rather than dropped, so that a keyword meant to do + something -- a typo, a stale argument name -- never goes silently unused. + """ + if not isinstance(err, TypeError) or "unexpected keyword argument" not in str(err): + return None + params = list(params) + if not params: + return None + cols = ", ".join(repr(p) for p in params) + return ( + f"If you are calling {func_name}(**df), subset the frame to the " + f"kernel's parameters: {func_name}(**df[[{cols}]])" + ) + + +def _signature_params(func) -> list: + """Parameter names of *func*, or an empty list if it cannot be introspected.""" + try: + return list(inspect.signature(func).parameters) + except (TypeError, ValueError): + return [] + + +def _jit_dsl_wrapper(kernel: DSLKernel, out, decorator_kwargs: dict): + """Build the call wrapper for the DSL (control-flow) dispatch route of `jit`. + + Unlike the tracing `wrapper` (which calls `func` once to record a single + expression, losing any branch not taken on that one call), this calls + `kernel` once per invocation through `blosc2.lazyudf`, so every branch and + loop in the kernel body is compiled and actually runs, once per chunk. + """ + + def dsl_wrapper(*args, **func_kwargs): + sig = kernel._sig + if sig is None: + raise TypeError(f"@blosc2.jit: cannot introspect the signature of {kernel.__name__!r}") + try: + bound = sig.bind(*args, **func_kwargs) + except TypeError as e: + # sig.bind's message names no function; prefix it, and point at the + # subsetting fix when a wide DataFrame was unpacked into the call. + hint = _wide_frame_hint(e, kernel.__name__, kernel.input_names or sig.parameters) + raise TypeError(f"{kernel.__name__}() {e}" + (f"\n{hint}" if hint else "")) from None + bound.apply_defaults() + values = tuple(bound.arguments[name] for name in kernel.input_names) + # Accept array-protocol operands (pandas Series, polars Series, ...) the + # same way the tracing route already does; zero-copy when the source is + # numpy-backed. + values = tuple( + np.asarray(v) + if not isinstance(v, np.ndarray | blosc2.NDArray) + and hasattr(v, "__array__") + and getattr(v, "ndim", 0) > 0 + else v + for v in values + ) + + array_shapes = { + v.shape + for v in values + if isinstance(v, np.ndarray | blosc2.NDArray) and getattr(v, "ndim", 0) > 0 + } + if not array_shapes: + shape = decorator_kwargs.get("shape") + if shape is None: + raise TypeError( + "@blosc2.jit DSL kernels with only scalar inputs require `shape=` " + "(passed to the jit decorator) to determine the result shape." + ) + elif len(array_shapes) > 1: + raise TypeError( + "blosc2.jit DSL kernels do not support broadcasting; all array arguments " + f"must share one shape, got {sorted(array_shapes)}" + ) + else: + (shape,) = array_shapes + + # Execution-tuning kwargs (jit/jit_backend/fp_accuracy) are baked into the + # LazyUDF at construction, so they take effect on *both* the getitem + # (NumPy) and compute (NDArray) return paths below. Storage kwargs + # (cparams, chunks, urlpath, ...) are applied once, only at the return + # step -- passing them here too would e.g. apply `urlpath=` twice and raise. + exec_kwargs = { + k: v for k, v in decorator_kwargs.items() if k in _JIT_EXECUTION_TUNING_KWARGS and v is not None + } + storage_kwargs = {k: v for k, v in decorator_kwargs.items() if k not in _JIT_EXECUTION_TUNING_KWARGS} + lexpr = blosc2.lazyudf(kernel, values, dtype=None, shape=shape, **exec_kwargs) + + if out is not None: + if isinstance(out, blosc2.NDArray): + raise NotImplementedError( + "blosc2.jit does not support an NDArray `out` on the DSL (control-flow) " + "dispatch route; use lexpr.compute(urlpath=..., mode='w') to persist a " + "result chunk-by-chunk instead." + ) + if not isinstance(out, np.ndarray): + raise TypeError(f"blosc2.jit `out` must be a NumPy array or NDArray, got {type(out)!r}") + if out.shape != shape: + raise TypeError(f"`out` shape {out.shape} does not match operand shape {shape}") + res = lexpr.compute(cparams=blosc2.CParams(clevel=0)) + if out.dtype != res.dtype: + raise TypeError( + f"`out` dtype {out.dtype} does not match the inferred result dtype {res.dtype}" + ) + if out.flags.c_contiguous: + res.get_slice_numpy(out, (tuple(0 for _ in res.shape), tuple(res.shape))) + else: + np.copyto(out, res[()], casting="no") + return out + + if storage_kwargs and any(v is not None for v in storage_kwargs.values()): + return lexpr.compute(**decorator_kwargs) + return lexpr[()] + + return dsl_wrapper + + +def jit(func=None, *, out=None, disable=False, strict=None, **kwargs): # noqa: C901 """ Prepare a function so that it can be used with the Blosc2 compute engine. The inputs of the function can be any combination of NumPy/NDArray arrays - and scalars. The function will be called with the NumPy arrays replaced by - :ref:`SimpleProxy` objects, whereas NDArray objects will be used as is. - - The returned value will be a NDArray if appropriate kwargs are provided - (e.g. `cparams=`). Else, the return value will be a NumPy array - (if the function returns a NumPy array). If `out` is provided, - the result will be computed and stored in the `out` array + and scalars. By default, the function is *traced*: it is called once with + the NumPy arrays replaced by :ref:`SimpleProxy` objects (NDArray objects are + used as is) to record a single expression, which is then what actually gets + evaluated. Because tracing only calls the function once, an ``if``/``for``/ + ``while`` in the body only ever takes the one path that single call + happened to follow — see `strict` below for when `jit` instead compiles the + function whole, so every branch and loop genuinely runs. + + The returned value will be a NDArray if a *storage* kwarg is provided (e.g. + `cparams=`, `chunks=`, `urlpath=` — anything that only makes sense for a + compressed/persisted container). Else, the return value will be a NumPy + array (if the function returns a NumPy array). Execution-tuning kwargs + (`jit=`, `jit_backend=`, `fp_accuracy=`) do not by themselves trigger this — + they take effect either way, without changing the return type. If `out` is + provided, the result will be computed and stored in the `out` array. Parameters ---------- func: callable The function to be prepared for the Blosc2 compute engine. out: np.ndarray, NDArray, optional - The output array where the result will be stored. + The output array where the result will be stored. On the DSL + (control-flow) dispatch route, a NumPy `out` is filled in place + (directly when C-contiguous, else via a copy); an NDArray `out` is not + supported there — use ``compute(urlpath=..., mode="w")`` instead. disable: bool, optional If True, the decorator is disabled and the original function is returned unchanged. Default is False. + strict: bool, optional + Control which evaluation route is used: + + - ``None`` (default): if *func*'s body contains an ``if``/``for``/``while`` + and it compiles as a DSL kernel, dispatch to the DSL route (miniexpr + runs the whole function, so branches/loops behave as written); a + control-flow function that fails DSL extraction still falls back to + tracing, but a subsequent tracing failure is annotated with the DSL + extraction error. Functions without control flow always trace, even + if they happen to be DSL-valid (tracing is faster for pure elementwise + expressions). + - ``True``: always use the DSL route, raising + :class:`~blosc2.dsl_kernel.DSLSyntaxError` at decoration time if + *func*'s source cannot be **parsed** as a DSL kernel. Note the + guarantee is exactly that -- parsing -- and not that the kernel will + compile: a function that is DSL-shaped but calls something miniexpr + does not implement passes here and fails later, at call time, with a + ``RuntimeError``. See the DSL syntax reference for what the grammar + accepts. (Unrelated to :func:`blosc2.dsl_kernel`, which builds a + :class:`DSLKernel` object rather than an evaluating wrapper.) + + This also works as a pandas engine, which is the only way to reach + ``strict`` through that entry point: + ``df.apply(f, engine=blosc2.jit(strict=True))``. + - ``False``: always use the tracing route, even if *func* has control + flow (this only works when branches/loops depend on plain Python + values, not on traced arrays). **kwargs: dict, optional Additional keyword arguments supported by the :func:`empty` constructor. @@ -788,25 +1073,83 @@ def jit(func=None, *, out=None, disable=False, **kwargs): (e.g. when using a reduction as the last function). In this case, you can still use the `out` parameter of the reduction function for some custom control over the output. + * DSL-route kernels do not support broadcasting: every array argument must + share the same shape. Examples -------- >>> import numpy as np >>> import blosc2 >>> @blosc2.jit - >>> def compute_expression(a, b, c): - >>> return np.sum(((a ** 3 + np.sin(a * 2)) > 2 * c) & (b > 0), axis=1) + ... def compute_expression(a, b, c): + ... return np.sum(((a ** 3 + np.sin(a * 2)) > 2 * c) & (b > 0), axis=1) >>> a = np.arange(20, dtype=np.float32).reshape(4, 5) >>> b = np.arange(20).reshape(4, 5) >>> c = np.arange(5) >>> compute_expression(a, b, c) - [5 5 5 5] + array([3, 5, 5, 5]) + + With ``strict=True`` the function is compiled as a DSL kernel, so a real + per-element ``if`` runs as written -- only the matching arm is evaluated: + + >>> @blosc2.jit(strict=True) + ... def clamp(x): + ... if x < 0.0: + ... out = 0.0 + ... else: + ... out = x + ... return out + >>> clamp(np.array([-1.5, 2.0, -0.5])) + array([0., 2., 0.]) + + The guarantee is that the source *parses* as DSL, checked at decoration + time. A body the grammar does not accept is rejected right away, rather + than silently falling back to tracing: + + >>> @blosc2.jit(strict=True) # doctest: +IGNORE_EXCEPTION_DETAIL + ... def not_dsl(x): + ... return np.where(x >= 0, x.mean(), x) + Traceback (most recent call last): + ... + blosc2.dsl_kernel.DSLSyntaxError: Unsupported call target in DSL ... """ - def decorator(func): + def decorator(func): # noqa: C901 if disable: return func + kernel = DSLKernel(func) + has_cf = _has_control_flow(kernel.dsl_source) + dsl_ok = kernel.dsl_source is not None and kernel.dsl_error is None + if strict is True and not dsl_ok: + # One condition, one exception type: DSLSyntaxError (a ValueError) + # whether the source failed to parse as DSL or could not be read at + # all (a lambda, a C function). The message avoids naming the + # decorator spelling, since `strict=True` also arrives through + # `df.apply(..., engine=blosc2.jit(strict=True))`. + raise kernel.dsl_error or DSLSyntaxError( + f"strict=True: could not extract a DSL kernel from {func.__name__!r}" + ) + use_dsl = strict is True or (strict is None and has_cf and dsl_ok) + + if use_dsl: + dsl_wrapper = _jit_dsl_wrapper(kernel, out, kwargs) + dsl_wrapper._blosc2_jit_wrapped = func + return dsl_wrapper + + _trace_hint = None + if strict is None and has_cf and not dsl_ok: + _trace_hint = ( + f"Note: {func.__name__!r} contains control flow (if/for/while) but could not be " + f"compiled as a DSL kernel: {kernel.dsl_error or 'source unavailable'}. See " + "doc/reference/dsl_syntax.md for the DSL syntax reference." + ) + + exec_kwargs = { + k: v for k, v in kwargs.items() if k in _JIT_EXECUTION_TUNING_KWARGS and v is not None + } + storage_kwargs = {k: v for k, v in kwargs.items() if k not in _JIT_EXECUTION_TUNING_KWARGS} + def wrapper(*args, **func_kwargs): # Get some kwargs in decorator for SimpleProxy constructor proxy_kwargs = {"chunks": kwargs.get("chunks"), "blocks": kwargs.get("blocks")} @@ -825,13 +1168,28 @@ def wrapper(*args, **func_kwargs): func_kwargs[key] = SimpleProxy(value, **proxy_kwargs) # Call function with the new arguments - retval = func(*new_args, **func_kwargs) + try: + retval = func(*new_args, **func_kwargs) + except Exception as e: + hints = [ + hint + for hint in ( + _wide_frame_hint( + e, getattr(func, "__name__", "the function"), _signature_params(func) + ), + _trace_hint, + ) + if hint is not None + ] + if hints: + raise type(e)("\n".join([str(e), *hints])) from e + raise # Treat return value # If it is a numpy array, return it as is if isinstance(retval, np.ndarray): - if kwargs and any(kwargs[key] is not None for key in kwargs): - # But if kwargs are provided, return a NDArray instead + if storage_kwargs and any(v is not None for v in storage_kwargs.values()): + # But if storage kwargs are provided, return a NDArray instead return blosc2.asarray(retval, **kwargs) return retval @@ -843,13 +1201,23 @@ def wrapper(*args, **func_kwargs): # If the return value is a LazyExpr, compute it if out is not None: return retval.compute(out=out, **kwargs) - if kwargs and any(kwargs[key] is not None for key in kwargs): + if storage_kwargs and any(v is not None for v in storage_kwargs.values()): return retval.compute(**kwargs) - # If no kwargs are provided, return a numpy array - return retval[()] + # No storage kwargs: return a NumPy array (like retval[()]), but still + # honor any execution-tuning kwargs (jit/jit_backend/fp_accuracy). + return retval.compute(_getitem=True, **exec_kwargs) + # Lets callers (notably the pandas engine below) tell an already-jitted + # function from a plain one, so it is not decorated a second time. + wrapper._blosc2_jit_wrapped = func return wrapper + # Carry the engine on the decorator too, so a configured call such as + # `blosc2.jit(strict=True)` is accepted by `df.apply(..., engine=...)`: + # pandas gates on hasattr(engine, "__pandas_udf__") and then uses the engine + # object itself as the decorator. + decorator.__pandas_udf__ = PandasUdfEngine + if func is None: return decorator else: @@ -886,7 +1254,7 @@ def map(cls, data, func, args, kwargs, decorator, skip_na): if skip_na: raise NotImplementedError("The Blosc2 engine does not support na_action='ignore' in map.") values = cls._ensure_numpy_data(data) - func = decorator(func) + func = _decorate_once(func, decorator) return func(values, *args, **kwargs) @classmethod @@ -899,7 +1267,11 @@ def apply(cls, data, func, args, kwargs, decorator, axis): """ orig = data values = cls._ensure_numpy_data(data) - func = decorator(func) + func_name = getattr(func, "__name__", "the function") + uses_subscript, has_loop = ( + _analyze_row_func(_undecorated(func)) if hasattr(orig, "columns") else (False, False) + ) + func = _decorate_once(func, decorator) if values.ndim == 1 or axis is None: # pandas Series.apply or pipe result = func(values, *args, **kwargs) @@ -908,9 +1280,46 @@ def apply(cls, data, func, args, kwargs, decorator, axis): result = [func(values[:, col_idx], *args, **kwargs) for col_idx in range(values.shape[1])] result = np.vstack(result).transpose() elif axis in (1, "columns"): - # pandas apply(axis=1) row-wise - result = [func(values[row_idx, :], *args, **kwargs) for row_idx in range(values.shape[0])] - result = np.vstack(result) + if uses_subscript and has_loop: + # row["colname"] combined with for/while: tracing would unroll + # the loop eagerly at call time, growing the traced expression + # with every iteration (a real per-row iteration count, like a + # Newton-Raphson loop, blows this up well past practical). No + # existing dispatch route can run this well; point at the one + # that can instead of hanging or crashing confusingly. + raise TypeError( + f"@blosc2.jit engine=... axis=1: {func_name!r} " + 'combines row["colname"] access with a for/while loop, which cannot be ' + "traced efficiently per-row. Call your @blosc2.jit function directly with " + "the DataFrame columns as separate array arguments instead: name its " + "parameters after the columns and call kernel(**df) -- see " + "doc/guides/pandas_engine.md." + ) + if uses_subscript: + # The `row["colname"]` idiom: replace the per-row Python loop + # with one call over whole per-column arrays (row-proxy, see + # `_PandasRowProxy`), extracted from the original DataFrame so + # per-column dtypes survive. + row_proxy = _PandasRowProxy(orig) + result = func(row_proxy, *args, **kwargs) + if not ( + isinstance(result, np.ndarray) + and result.ndim == 1 + and result.shape[0] == values.shape[0] + ): + raise TypeError( + '@blosc2.jit engine=... axis=1: functions using row["colname"] must ' + f"return one scalar per row (shape ({values.shape[0]},)); got " + f"{result!r}. Returning multiple values per row is not supported here." + ) + else: + # pandas apply(axis=1) row-wise: the historical per-row loop. + # Fine for functions treating the row as a plain array (e.g. + # `row + 1`); functions using row["colname"] are dispatched + # above instead, since this loop hands each call a positional + # ndarray row that does not support string subscripting. + result = [func(values[row_idx, :], *args, **kwargs) for row_idx in range(values.shape[0])] + result = np.vstack(result) else: raise NotImplementedError(f"Unknown axis '{axis}'. Use one of 0, 1 or None.") diff --git a/tests/ndarray/test_dsl_kernels.py b/tests/ndarray/test_dsl_kernels.py index 0439015cf..a001a474e 100644 --- a/tests/ndarray/test_dsl_kernels.py +++ b/tests/ndarray/test_dsl_kernels.py @@ -1157,3 +1157,176 @@ def other(a, b): st = blosc2.validate_dsl_jit(other, [np.float64, np.float64], np.float64) assert st["compiled"] assert not st["jit"] + + +def _dsl_reference(kernel, operands, dtype=None): + """Evaluate *kernel* over NDArray copies of *operands* (same engine, same buffers).""" + nd_operands = tuple(blosc2.asarray(op) if isinstance(op, np.ndarray) else op for op in operands) + return blosc2.lazyudf(kernel, nd_operands, dtype=dtype)[()] + + +@blosc2.dsl_kernel +def _numpy_operand_kernel(x, y): + return x * 2.0 + y + + +@pytest.mark.parametrize( + "shape", + [ + (10_007,), # 1-D: partial last chunk and block + (101, 67), # 2-D: odd shape + (13, 17, 19), # 3-D: odd shape + ], +) +def test_dsl_kernel_numpy_operands_match_ndarray_reference(shape): + rng = np.random.default_rng(0) + a = rng.random(shape).astype(np.float64) + b = rng.random(shape).astype(np.float64) + res = blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None)[()] + np.testing.assert_array_equal(res, _dsl_reference(_numpy_operand_kernel, (a, b))) + + +def test_dsl_kernel_numpy_operands_mixed_dtype_promotes_output(): + rng = np.random.default_rng(1) + a = (rng.random(10_007) * 10).astype(np.float32) + b = (rng.random(10_007) * 10).astype(np.int64) + res = blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None)[()] + ref = _dsl_reference(_numpy_operand_kernel, (a, b)) + assert res.dtype == ref.dtype + np.testing.assert_array_equal(res, ref) + + +def test_dsl_kernel_ndarray_operands_with_different_itemsize(): + # Blocks are sized in bytes, so a float32 and an int64 operand get different + # chunks/blocks by default; the DSL path has no slow fallback, so it used to + # raise "slicing is not supported" whenever the grids diverged (which depends + # on array size and on the platform's cache detection). + n = 1_000_000 + a = (np.arange(n) % 7).astype(np.float32) + b = (np.arange(n) % 5).astype(np.int64) + A, B = blosc2.asarray(a), blosc2.asarray(b) + assert (A.chunks, A.blocks) != (B.chunks, B.blocks) + res = blosc2.lazyudf(_numpy_operand_kernel, (A, B), dtype=None)[()] + np.testing.assert_array_equal(res, a * 2.0 + b) + + +def test_dsl_kernel_mixed_ndarray_and_numpy_operand(): + shape = (20, 10) + a = np.arange(np.prod(shape), dtype=np.float64).reshape(shape) + b = blosc2.asarray(np.arange(np.prod(shape), dtype=np.float64).reshape(shape) * 2) + res = blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None)[()] + ref = blosc2.lazyudf(_numpy_operand_kernel, (blosc2.asarray(a), b), dtype=None)[()] + np.testing.assert_array_equal(res, ref) + + +def test_dsl_kernel_numpy_operands_f_ordered_and_strided(): + shape = (20, 10) + b = np.arange(np.prod(shape), dtype=np.float64).reshape(shape) + ref = _dsl_reference(_numpy_operand_kernel, (b, b)) + + a_f = np.asfortranarray(b) + res_f = blosc2.lazyudf(_numpy_operand_kernel, (a_f, b), dtype=None)[()] + np.testing.assert_array_equal(res_f, ref) + + a_strided = np.arange(2 * np.prod(shape), dtype=np.float64).reshape(40, 10)[::2] + ref_strided = _dsl_reference(_numpy_operand_kernel, (np.ascontiguousarray(a_strided), b)) + res_strided = blosc2.lazyudf(_numpy_operand_kernel, (a_strided, b), dtype=None)[()] + np.testing.assert_array_equal(res_strided, ref_strided) + + +def test_dsl_kernel_numpy_operand_non_native_endian_requires_miniexpr(): + a = np.arange(100, dtype=">f8").reshape(10, 10) + b = np.arange(100, dtype=np.float64).reshape(10, 10) + with pytest.raises(RuntimeError, match="NDArray or NumPy inputs"): + blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None)[()] + + +def test_dsl_kernel_zero_input_dummy_operand_injection_still_works(): + @blosc2.dsl_kernel + def ramp(start, step): + return start + step * _i0 # noqa: F821 # DSL index symbol resolved by miniexpr + + res = blosc2.lazyudf(ramp, (1.0, 2.0), dtype=np.float64, shape=(100,))[()] + expected = 1.0 + 2.0 * np.arange(100, dtype=np.float64) + np.testing.assert_allclose(res, expected) + + +def test_dsl_kernel_numpy_out_matches_compute_and_honors_explicit_cparams(): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + lexpr = blosc2.lazyudf(_numpy_operand_kernel, (a, b), dtype=None) + + res_getitem = lexpr[()] + res_compute = lexpr.compute()[:] + np.testing.assert_array_equal(res_getitem, res_compute) + + res_explicit = lexpr.compute(cparams=blosc2.CParams(clevel=5)) + assert res_explicit.schunk.cparams.clevel == 5 + + +def test_dsl_kernel_numpy_attribute_calls_are_rewritten_to_bare_names(): + @blosc2.dsl_kernel + def k(x, y): + if x >= 0: + return np.sin(x) + y + else: + return -np.sin(-x) + y + + assert k.dsl_error is None + assert "np.sin" not in k.dsl_source + assert "sin(" in k.dsl_source + + x = np.linspace(-2, 2, 1000) + y = np.ones_like(x) + res = blosc2.lazyudf(k, (x, y), dtype=None)[()] + expected = np.where(x >= 0, np.sin(x) + y, -np.sin(-x) + y) + np.testing.assert_allclose(res, expected) + + +@pytest.mark.parametrize( + ("numpy_call", "expected_dsl_name"), + [ + # np.power keeps its name: the DSL accepts `power` as an alias of `pow`, + # so only the `np.` prefix is stripped. + ("np.power(x, 2.0)", "power"), + ("np.maximum(x, 0.5)", "fmax"), + ("np.minimum(x, -0.5)", "fmin"), + ("np.absolute(x)", "abs"), + ], +) +def test_dsl_kernel_numpy_func_aliases_map_to_dsl_names(numpy_call, expected_dsl_name): + src = f"def k(x):\n if x >= 0:\n return {numpy_call}\n else:\n return -x\n" + k = kernel_from_source(src) + + assert k.dsl_error is None + assert f"{expected_dsl_name}(" in k.dsl_source + + x = np.linspace(-2, 2, 1000) + res = blosc2.lazyudf(k, (x,), dtype=None)[()] + expected = np.where(x >= 0, eval(numpy_call, {"np": np, "x": x}), -x) + np.testing.assert_allclose(res, expected) + + +def test_dsl_kernel_numpy_alias_not_rewritten_when_shadowed_by_parameter(): + # A parameter literally named "np" shadows the module -- the rewrite must + # not mistake a per-call NDArray/scalar input for the NumPy module. + src = "def k(np, y):\n return np * y\n" + k = kernel_from_source(src) + + assert k.dsl_error is None + a = np.linspace(1, 2, 100) + b = np.linspace(2, 3, 100) + res = blosc2.lazyudf(k, (a, b), dtype=None)[()] + np.testing.assert_allclose(res, a * b) + + +def test_dsl_kernel_numpy_call_without_alias_left_untouched(): + # No import of numpy bound in the kernel's defining scope -- nothing to + # rewrite, and the plain bare-name form still works unaffected. + @blosc2.dsl_kernel + def k(x): + return sin(x) # noqa: F821 # 'sin' resolved as a bare DSL function name + + a = np.linspace(-2, 2, 1000) + res = blosc2.lazyudf(k, (a,), dtype=None)[()] + np.testing.assert_allclose(res, np.sin(a)) diff --git a/tests/ndarray/test_elementwise_funcs.py b/tests/ndarray/test_elementwise_funcs.py index c82705f2c..91bb6655f 100644 --- a/tests/ndarray/test_elementwise_funcs.py +++ b/tests/ndarray/test_elementwise_funcs.py @@ -355,3 +355,13 @@ def test_binary_funcs_torch_proxy(np_func, blosc_func, dtype, shape, chunkshape) @pytest.mark.parametrize(("shape", "chunkshape"), SHAPES_CHUNKS_HEAVY) def test_binary_funcs_heavy(np_func, blosc_func, dtype, shape, chunkshape): _test_binary_func_impl(np_func, blosc_func, dtype, shape, chunkshape) + + +@pytest.mark.parametrize("np_func", [np.sign, np.square, np.negative, np.positive, np.reciprocal]) +@pytest.mark.parametrize("dtype", [blosc2.int64, blosc2.float64, blosc2.complex128]) +def test_ufunc_dispatch(np_func, dtype): + # These must build a lazy expression, not raise TypeError from __array_ufunc__. + # No zero in the input: reciprocal(0) warns, and that is not what is under test. + a = np.array([-2, 1, 3], dtype=dtype) + b = blosc2.asarray(a) + np.testing.assert_allclose(np_func(b)[()], np_func(a)) diff --git a/tests/ndarray/test_jit.py b/tests/ndarray/test_jit.py index 0dbcdf820..ba09af30b 100644 --- a/tests/ndarray/test_jit.py +++ b/tests/ndarray/test_jit.py @@ -177,3 +177,31 @@ def reduc_std_jit_cparams(a, b, c): assert d_jit.schunk.cparams.clevel == 1 assert d_jit.schunk.cparams.codec == blosc2.Codec.LZ4 assert d_jit.schunk.cparams.filters == [blosc2.Filter.BITSHUFFLE] + [blosc2.Filter.NOFILTER] * 5 + + +def test_jit_execution_tuning_kwarg_alone_keeps_numpy_return(): + # jit/jit_backend/fp_accuracy tune *how* an expression runs, not what + # container the result comes back in -- they must not by themselves flip + # the return type from NumPy to NDArray (unlike storage kwargs). + @blosc2.jit(jit=False) + def f(a, b): + return a * 2.0 + b + + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + res = f(a, b) + assert isinstance(res, np.ndarray) + np.testing.assert_allclose(res, a * 2.0 + b) + + +def test_jit_execution_tuning_kwarg_with_storage_kwarg_still_returns_ndarray(): + @blosc2.jit(jit=False, cparams=blosc2.CParams(clevel=2)) + def f(a, b): + return a * 2.0 + b + + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + res = f(a, b) + assert isinstance(res, blosc2.NDArray) + assert res.schunk.cparams.clevel == 2 + np.testing.assert_allclose(res[:], a * 2.0 + b) diff --git a/tests/ndarray/test_jit_dsl_dispatch.py b/tests/ndarray/test_jit_dsl_dispatch.py new file mode 100644 index 000000000..6a7ba8751 --- /dev/null +++ b/tests/ndarray/test_jit_dsl_dispatch.py @@ -0,0 +1,269 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +import numpy as np +import pytest + +import blosc2 + + +def _mandel_numpy(cr, ci, max_iter): + zr = np.zeros_like(cr) + zi = np.zeros_like(ci) + n = np.zeros(cr.shape, dtype=np.int64) + active = np.ones(cr.shape, dtype=bool) + for _ in range(max_iter): + mag = zr * zr + zi * zi + active = active & ~(active & (mag > 4.0)) + new_zr = zr * zr - zi * zi + cr + new_zi = 2 * zr * zi + ci + zr = np.where(active, new_zr, zr) + zi = np.where(active, new_zi, zi) + n = np.where(active, n + 1, n) + return n + + +def _mandel_grid(): + h, w = 12, 16 + cr = np.linspace(-2, 1, w).astype(np.float64)[None, :] * np.ones((h, 1)) + ci = np.linspace(-1, 1, h).astype(np.float64)[:, None] * np.ones((1, w)) + return cr, ci + + +def test_jit_control_flow_dispatches_to_dsl_and_matches_numpy(monkeypatch, capsys): + @blosc2.jit + def mandel(cr, ci, max_iter): + zr = 0.0 + zi = 0.0 + n = 0 + for _i in range(max_iter): + if zr * zr + zi * zi > 4.0: + break + new_zr = zr * zr - zi * zi + cr + zi = 2 * zr * zi + ci + zr = new_zr + n = n + 1 + return n + + cr, ci = _mandel_grid() + monkeypatch.setenv("BLOSC_ME_JIT_TRACE", "1") + res = mandel(cr, ci, 30) + captured = capsys.readouterr() + # Under WebAssembly this kernel is transpiled to the JS bridge instead of + # going through miniexpr (see _maybe_js_backend); both engines trace. + assert f"engine={'js' if blosc2.IS_WASM else 'miniexpr'}" in captured.out + assert "def mandel" in captured.out + np.testing.assert_array_equal(res, _mandel_numpy(cr, ci, 30)) + + +def test_jit_control_flow_with_default_argument(): + @blosc2.jit + def mandel(cr, ci, max_iter=30): + zr = 0.0 + zi = 0.0 + n = 0 + for _i in range(max_iter): + if zr * zr + zi * zi > 4.0: + break + new_zr = zr * zr - zi * zi + cr + zi = 2 * zr * zi + ci + zr = new_zr + n = n + 1 + return n + + cr, ci = _mandel_grid() + res = mandel(cr, ci) + np.testing.assert_array_equal(res, _mandel_numpy(cr, ci, 30)) + + +def test_jit_elementwise_function_still_traces(monkeypatch): + calls = [] + real_lazyudf = blosc2.lazyudf + + def spy_lazyudf(*args, **kwargs): + calls.append((args, kwargs)) + return real_lazyudf(*args, **kwargs) + + monkeypatch.setattr(blosc2, "lazyudf", spy_lazyudf) + + @blosc2.jit + def elemwise(a, b): + return a * 2.0 + b + + a = np.arange(100, dtype=np.float64) + b = np.arange(100, dtype=np.float64) * 0.5 + res = elemwise(a, b) + np.testing.assert_allclose(res, a * 2.0 + b) + assert calls == [] # no control flow -> never routed through the DSL/lazyudf path + + +def test_jit_strict_true_on_elementwise_dsl_valid_function_uses_dsl(monkeypatch): + calls = [] + import blosc2.proxy as proxy_mod + + real_wrapper = proxy_mod._jit_dsl_wrapper + + def spy(*args, **kwargs): + calls.append(True) + return real_wrapper(*args, **kwargs) + + monkeypatch.setattr(proxy_mod, "_jit_dsl_wrapper", spy) + + @blosc2.jit(strict=True) + def elemwise(a, b): + return a * 2.0 + b + + a = np.arange(100, dtype=np.float64) + b = np.arange(100, dtype=np.float64) * 0.5 + res = elemwise(a, b) + np.testing.assert_allclose(res, a * 2.0 + b) + assert calls # dispatched through the DSL wrapper + + +def test_jit_strict_true_on_non_dsl_function_raises_at_decoration_time(): + with pytest.raises(Exception, match="axis"): + + @blosc2.jit(strict=True) + def bad(a): + return np.sum(a, axis=1) + + +def test_jit_strict_false_on_control_flow_traces(): + @blosc2.jit(strict=False) + def cf_func(a, b): + if True: + return a + b + return a - b + + a = np.arange(100, dtype=np.float64) + b = np.arange(100, dtype=np.float64) * 0.5 + res = cf_func(a, b) + np.testing.assert_allclose(res, a + b) + + +def test_jit_control_flow_on_python_scalar_flag_still_traces(): + @blosc2.jit + def scalar_flag(a, b, flag): + if flag: + return a + b + return a - b + + a = np.arange(100, dtype=np.float64) + b = np.arange(100, dtype=np.float64) * 0.5 + np.testing.assert_allclose(scalar_flag(a, b, True), a + b) + np.testing.assert_allclose(scalar_flag(a, b, False), a - b) + + +def test_jit_dsl_route_rejects_broadcasting(): + @blosc2.jit + def kernel(a, b, n): + acc = 0.0 + for _i in range(n): + acc = acc + a + b + return acc + + with pytest.raises(TypeError, match="broadcasting"): + kernel(np.zeros((10,)), np.zeros((20,)), 2) + + +def _kernel_src(a, b, n): + acc = 0.0 + for _i in range(n): + acc = acc + a + b + return acc + + +def test_jit_dsl_route_out_numpy_c_contiguous_filled_in_place(): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + out = np.empty(1000, dtype=np.float64) + jit_f = blosc2.jit(out=out)(_kernel_src) + res = jit_f(a, b, 3) + assert res is out + np.testing.assert_allclose(out, (a + b) * 3) + + +def test_jit_dsl_route_out_numpy_non_contiguous_uses_copyto_fallback(): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + out = np.empty(2000, dtype=np.float64)[::2] + assert not out.flags.c_contiguous + jit_f = blosc2.jit(out=out)(_kernel_src) + res = jit_f(a, b, 3) + assert res is out + np.testing.assert_allclose(out, (a + b) * 3) + + +def test_jit_dsl_route_out_mismatched_shape_or_dtype_raises_typeerror(): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + + with pytest.raises(TypeError, match="shape"): + blosc2.jit(out=np.empty(500, dtype=np.float64))(_kernel_src)(a, b, 3) + + with pytest.raises(TypeError, match="dtype"): + blosc2.jit(out=np.empty(1000, dtype=np.float32))(_kernel_src)(a, b, 3) + + +def test_jit_dsl_route_ndarray_out_raises_not_implemented_mentioning_urlpath(): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + nd_out = blosc2.zeros((1000,), dtype=np.float64) + with pytest.raises(NotImplementedError, match="urlpath"): + blosc2.jit(out=nd_out)(_kernel_src)(a, b, 3) + + +def test_jit_dsl_route_compute_urlpath_persists_result(tmp_path): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + urlpath = str(tmp_path / "persisted.b2nd") + jit_f = blosc2.jit(urlpath=urlpath, mode="w")(_kernel_src) + res = jit_f(a, b, 3) + assert isinstance(res, blosc2.NDArray) + reopened = blosc2.open(urlpath) + np.testing.assert_allclose(reopened[:], (a + b) * 3) + + +def test_jit_dsl_route_ndarray_operands_match_numpy_operands(): + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + na = blosc2.asarray(a) + nb = blosc2.asarray(b) + jit_f = blosc2.jit()(_kernel_src) + res_numpy = jit_f(a, b, 3) + res_ndarray = jit_f(na, nb, 3) + np.testing.assert_array_equal(res_numpy, res_ndarray) + + +def test_jit_dsl_route_execution_tuning_kwarg_alone_keeps_numpy_return(): + # Same rule as the tracing route: jit/jit_backend/fp_accuracy tune execution, + # not the return container, so they must not force an NDArray on their own. + jit_f = blosc2.jit(jit=False)(_kernel_src) + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + res = jit_f(a, b, 3) + assert isinstance(res, np.ndarray) + np.testing.assert_allclose(res, (a + b) * 3) + + +def test_jit_dsl_route_execution_tuning_kwarg_with_storage_kwarg_still_returns_ndarray(): + jit_f = blosc2.jit(jit=False, cparams=blosc2.CParams(clevel=2))(_kernel_src) + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + res = jit_f(a, b, 3) + assert isinstance(res, blosc2.NDArray) + assert res.schunk.cparams.clevel == 2 + np.testing.assert_allclose(res[:], (a + b) * 3) + + +def test_jit_dsl_route_accepts_array_protocol_operands(): + pd = pytest.importorskip("pandas") + a = np.arange(1000, dtype=np.float64) + b = np.arange(1000, dtype=np.float64) * 0.5 + df = pd.DataFrame({"a": a, "b": b}) + jit_f = blosc2.jit()(_kernel_src) + np.testing.assert_array_equal(jit_f(df["a"], df["b"], 3), jit_f(a, b, 3)) diff --git a/tests/test_pandas_udf_engine.py b/tests/test_pandas_udf_engine.py index 7441345f9..f9df6bc16 100644 --- a/tests/test_pandas_udf_engine.py +++ b/tests/test_pandas_udf_engine.py @@ -180,3 +180,200 @@ def test_apply_object_dtype_raises_clear_error(self): df = pd.DataFrame({"a": ["x", "y"]}) with pytest.raises(ValueError, match="numeric dtype"): df.apply(lambda x: x + 1, engine=blosc2.jit) + + def test_apply_axis1_row_subscript_idiom_matches_default_engine(self): + def add_people(row): + return row["max_people"] + row["max_children"] + + df = pd.DataFrame({"max_people": [4, 2, 8], "max_children": [1, 0, 3]}) + expected = df.apply(add_people, axis=1) + result = df.apply(add_people, engine=blosc2.jit, axis=1) + pd.testing.assert_series_equal(result, expected) + + def test_apply_axis1_row_subscript_args_kwargs_forwarded(self): + def combine(row, num1, num2=0): + return row["a"] + row["b"] + num1 + num2 + + df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]}) + expected = df.apply(combine, axis=1, args=(10,), num2=100) + result = df.apply(combine, engine=blosc2.jit, axis=1, args=(10,), num2=100) + pd.testing.assert_series_equal(result, expected) + + def test_apply_axis1_row_subscript_preserves_column_dtype(self): + # a mixed-dtype frame would be upcast by DataFrame.values; the row + # proxy must extract columns from the original frame instead. + def add(row): + return row["i"] + row["f"] + + df = pd.DataFrame({"i": np.array([1, 2, 3], dtype=np.int64), "f": [0.5, 0.5, 0.5]}) + result = df.apply(add, engine=blosc2.jit, axis=1) + np.testing.assert_allclose(result.to_numpy(), [1.5, 2.5, 3.5]) + + def test_apply_axis1_row_subscript_with_loop_raises_clear_error(self): + def kepler_row(row): + m, ecc = row["m"], row["ecc"] + e = m + ecc * np.sin(m) + for _ in range(50): + diff = (e - ecc * np.sin(e) - m) / (1.0 - ecc * np.cos(e)) + e = e - diff + return e + + df = pd.DataFrame({"m": [0.1, 0.5], "ecc": [0.1, 0.2]}) + with pytest.raises(TypeError, match="for/while loop"): + df.apply(kepler_row, engine=blosc2.jit, axis=1) + + def test_apply_axis1_row_subscript_duplicate_column_raises(self): + def add(row): + return row["a"] + 1 + + df = pd.DataFrame(np.ones((2, 2)), columns=["a", "a"]) + with pytest.raises(KeyError, match="duplicated"): + df.apply(add, engine=blosc2.jit, axis=1) + + def test_apply_axis1_row_subscript_attribute_access_raises(self): + def bad(row): + return row["a"] + row.b + + df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]}) + with pytest.raises(AttributeError, match="row\\['b'\\]"): + df.apply(bad, engine=blosc2.jit, axis=1) + + def test_apply_axis1_row_subscript_non_numeric_column_raises(self): + # Whole-frame numeric-dtype validation (`_ensure_numpy_data`) already + # gates this ahead of row-proxy dispatch; `_PandasRowProxy` carries + # its own per-column check too, for callers that construct it + # directly. + def bad(row): + return row["a"] + len(row["b"]) + + df = pd.DataFrame({"a": [1.0, 2.0], "b": ["x", "y"]}) + with pytest.raises(ValueError, match="numeric dtype"): + df.apply(bad, engine=blosc2.jit, axis=1) + + def test_apply_axis1_positional_idiom_still_uses_per_row_loop(self): + # No `row["..."]` subscript: falls back to the historical per-row + # loop, unaffected by the row-proxy dispatch added for the subscript + # idiom above. + df = pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}) + expected = df.apply(lambda row: row * 2, axis=1) + result = df.apply(lambda row: row * 2, engine=blosc2.jit, axis=1) + pd.testing.assert_frame_equal(result, expected) + + def test_apply_already_jitted_function_is_not_decorated_twice(self): + # Decorating and passing engine= both request the same thing. Applying + # the decorator a second time used to wrap the array in a SimpleProxy + # before the inner DSL kernel saw it, which then failed asking for + # `shape=`. Traced functions tolerated it, so only branches broke. + def branch(col): + if col >= 0: + out = col + 1.0 + else: + out = col - 1.0 + return out + + df = pd.DataFrame({"a": [-2.0, 1.0, 3.0], "b": [4.0, -5.0, 6.0]}) + expected = df.apply(lambda col: np.where(col >= 0, col + 1.0, col - 1.0)) + + for func in (branch, blosc2.jit(branch)): + result = df.apply(func, engine=blosc2.jit) + pd.testing.assert_frame_equal(result, expected) + + def test_map_already_jitted_function_is_not_decorated_twice(self): + def branch(col): + if col >= 0: + out = col * 2.0 + else: + out = col + return out + + s = pd.Series([-2.0, 1.0, 3.0]) + expected = np.where(s.to_numpy() >= 0, s.to_numpy() * 2.0, s.to_numpy()) + + for func in (branch, blosc2.jit(branch)): + result = s.map(func, engine=blosc2.jit) + np.testing.assert_allclose(np.asarray(result), expected) + + def test_apply_engine_accepts_configured_jit(self): + # pandas gates on hasattr(engine, "__pandas_udf__") and then uses the + # engine object as the decorator, so a configured blosc2.jit(...) call + # is a valid engine -- the only way to reach strict= through apply(). + from blosc2.dsl_kernel import DSLSyntaxError + + def branch(col): + if col >= 0: + out = col + 1.0 + else: + out = col - 1.0 + return out + + def not_dsl(col): + return np.where(col >= 0, col.mean() + 1.0, col - 1.0) + + df = pd.DataFrame({"a": [-2.0, 1.0, 3.0], "b": [4.0, -5.0, 6.0]}) + expected = df.apply(lambda col: np.where(col >= 0, col + 1.0, col - 1.0)) + + result = df.apply(branch, engine=blosc2.jit(strict=True)) + pd.testing.assert_frame_equal(result, expected) + + # strict=True refuses to silently fall back to tracing + with pytest.raises(DSLSyntaxError): + df.apply(not_dsl, engine=blosc2.jit(strict=True)) + + # strict=False forces the tracing route instead + df.apply(not_dsl, engine=blosc2.jit(strict=False)) + + def test_columns_by_keyword_unpacking(self): + # doc/guides/pandas_engine.md's row-wise pattern: a DataFrame is a + # mapping of column name to Series, so `kernel(**df)` passes each + # column as a keyword argument. Both jit routes must accept that, and + # bind by name: the columns below are in neither the parameter order + # nor alphabetical order, and the operations are asymmetric, so a + # positional binding would give a different (wrong) answer. + @blosc2.jit + def traced(a, b): + return b - a * 2.0 + + @blosc2.jit + def dsl(a, b): + if a > b: + out = a - b + else: + out = b - a * 2.0 + return out + + df = pd.DataFrame({"b": [4.0, -5.0, 6.0], "a": [-2.0, 1.0, 3.0]}) + assert list(df.columns) == ["b", "a"] + + np.testing.assert_allclose(np.asarray(traced(**df)), np.asarray(traced(df["a"], df["b"]))) + np.testing.assert_allclose(np.asarray(traced(**df)), df["b"] - df["a"] * 2.0) + np.testing.assert_allclose(np.asarray(dsl(**df)), np.asarray(dsl(df["a"], df["b"]))) + + def test_wide_frame_kwargs_error_names_the_fix(self): + # Extra columns are rejected, not dropped: a keyword that goes nowhere + # would otherwise fail silently. The message must name the subsetting fix. + @blosc2.jit + def traced(a, b): + return a + b + + @blosc2.jit + def dsl(a, b): + if a > b: + out = a - b + else: + out = b - a + return out + + df = pd.DataFrame({"a": [1.0, 2.0], "b": [4.0, 5.0], "note": [7.0, 8.0]}) + + # (`func.__name__` is the jit wrapper's; the message uses the kernel's) + for name, func in (("traced", traced), ("dsl", dsl)): + with pytest.raises(TypeError) as excinfo: + func(**df) + message = str(excinfo.value) + assert name in message + assert "'note'" in message + assert "**df[['a', 'b']]" in message + + # A missing operand is a different mistake and keeps its own message + with pytest.raises(TypeError, match="missing a required argument"): + dsl(a=df["a"])