apache/datafusion-comet PR 5612 regex PatternCache benchmarks: harness, results, sort-path repro, mechanism micro-bench

README.md

PR #5612 benchmark harness: per-expression regex PatternCache (datafusion-comet)

Note: gists have no directories, so files referenced below as harness/src/... and results/... appear here under their bare filenames.

Standalone benchmarks for apache/datafusion-comet PR #5612, comparing:

• BASE: commit 8729f6e6a (main at the time, one regex compile per batch invocation)
• HEAD: commit b4d915236 (branch perf/compile-user-regex-once, one-slot PatternCache per planned expression)

Machine and toolchain

• Apple M5, 10 cores (4 performance + 6 efficiency), macOS (Darwin 25.6.0)
• rustc 1.98.0 (88d9e12ae 2026-08-18), cargo 1.98.0
• All builds --release (opt-level 3, thin LTO, codegen-units 1)
• Dependency versions match the repo workspace pins: datafusion 54.1.0, arrow 58.4.0, regex 1.13.1 (resolved from the workspace's regex = "1.12.3"), regex-automata 0.4.16

How it builds

harness/Cargo.toml has a path dependency on the repo's native/spark-expr crate, so the identical harness source builds against whichever commit is checked out in the repo. The harness uses only symbols that exist in both commits: UDFs are constructed through create_comet_physical_fun (the serde entry point) and evaluated through ScalarUDF::invoke_with_args, with the regex pattern supplied as a scalar argument on every invoke, exactly as the engine does. Build each commit into its own target dir so the two builds never clobber each other:

cd harness
# with the repo checked out at HEAD (b4d915236):
cargo build --release --target-dir ../target-head
# with the repo checked out at BASE (git checkout 8729f6e6a):
cargo build --release --target-dir ../target-base

Adjust the absolute path in harness/Cargo.toml to your checkout location.

Binaries

1. bench5612 (src/main.rs): the main matrix. Functions {regexp_extract, regexp_extract_all, split} x sharing {shared single UDF instance across threads, one instance per worker} x workers {1,2,4,8} x batch rows {512, 8192} x pattern regime {warm, alternating}. Also measures allocations per batch (counting global allocator), cold first-invoke cost, and FNV-1a hashes of all outputs for a BASE vs HEAD equality check. Run: bench5612 <label> <out_dir>.
2. sort_repro (src/bin/sort_repro.rs): the sort-key sharing scenario. Builds the exact plan shape Comet produces (PhysicalSortExpr wrapping a scalar function expression, passed to SortExec with no fetch, single partition memory source of 128 x 8192-row batches, about 21 MB so the reservation is far past the 1 MiB sort_in_place_threshold_bytes default). DataFusion 54.1.0 ExternalSorter::in_mem_sort_stream then sorts every buffered batch in its own tokio task via spawn_buffered (datafusion-physical-plan-54.1.0/src/sorts/sort.rs lines 586-646), all tasks sharing the same UDF instance. A tracking shim around one comet UDF instance records total invokes and max concurrent in-flight invocations. Run: sort_repro <label> <function> <workers> <iters> on a multithread tokio runtime with <workers> worker threads.
3. shared_clone (src/bin/shared_clone.rs): commit-independent micro benchmark that isolates the mechanism of the shared-instance slowdown using the regex crate directly. Modes price separately: fresh compile per invoke (BASE), mutex slot + clone from a shared source (HEAD shared), mutex slot + clone from a per-thread source (HEAD per worker), plain clones from shared vs private sources, hoisted clones, direct shared use, and clone-only loops. Run: shared_clone [invokes_per_thread].

Result files

• results/{base,head}_cells.csv and {base2,head2}_cells.csv: two full replicates of the main matrix (throughput, mean/p50/p99 per-batch latency, allocations per batch)
• results/*_cold.csv: cold first-invoke cost, 30 fresh instances each
• results/*_verify.csv: output hashes; identical across all four runs (BASE and HEAD produce byte-identical results)
• results/sort_head_w1.csv, sort_head_w8.csv, sort_base_w1.csv, sort_base_w8.csv: first sort_repro runs (regexp_extract_all only, before the function column was added)
• results/sort_{base,head}_all.csv: sort_repro for both regexp_extract_all and regexp_extract at 1 and 8 workers
• results/sort_extract_w1_rep2.csv, sort_extractall_w8_rep2.csv: interleaved BASE/HEAD confirmation runs
• results/shared_clone_micro.csv: mechanism micro benchmark
• analyze.py: merges the main-matrix CSVs into markdown tables and checks the output hashes

Key findings

1. Warm patterns (the realistic case; Comet plans regexp patterns as literals, one expression instance per Spark task thread): HEAD removes a full regex compile per batch. Small batches benefit most (regexp_extract, 512 rows, 8 per-worker threads: 5.8 to 62.8 Mrows/s). Allocations per batch drop up to 92 percent.
2. Alternating patterns on one instance (one-slot cache worst case, 1 worker): within 2 percent of BASE for all three functions.
3. Concurrent same-instance evaluation does happen on the SortExec path (max in-flight 9 at 8 runtime workers, 2 at 1 worker because the merge evaluates on the driver thread concurrently with one spawned sort task).
4. On that sort path, HEAD regresses vs BASE for regexp_extract_all as the sort key (about 1.45x wall at 1 worker, about 2.2x at 8 workers), while regexp_extract as the sort key shows parity to a small HEAD win.
5. Mechanism (shared_clone_micro.csv): not the PatternCache mutex and not the per-invoke Regex clone. regex-automata 0.4.16 impl Clone for Regex (src/meta/regex.rs lines 1916-1926) shares the compiled program via Arc but creates a fresh private scratch pool, so scratch state is never shared. The contention is per row: Regex::captures_iter calls create_captures (meta/regex.rs line 657), which is Captures::all(self.group_info().clone()) (line 1571), an Arc::clone of the program-owned GroupInfo, plus one more Captures clone per match (CapturesMatches::next). When many threads run captures_iter over clones of ONE compiled program, they all hammer the refcount cache line of that single GroupInfoInner allocation. In the micro benchmark this caps every shared-program mode at about 6-15 Mrows/s regardless of thread count (2 threads drop below 1 thread total), while per-thread-program modes scale to 60-80 Mrows/s. regexp_extract is unaffected because it reuses one CaptureLocations across rows (captures_read, native/spark-expr/src/string_funcs/regexp_extract.rs lines 94-107); split is unaffected because find_iter creates no Captures. BASE never shares a program across threads (it compiles per invocation), which is why BASE shared and per-worker numbers are identical.

Fix verification

After the analysis above, regexp_extract_all was changed on the PR branch (commit f35bc97fa) to drive iteration with find_iter and resolve capture groups through captures_read_at into one CaptureLocations reused for the whole batch, in both extract_all_array and extract_one. That removes the per-row create_captures and the per-match Captures clone, and with them both the shared GroupInfo Arc refcount traffic and the per-match heap allocations.

Result files: results/sort_fix.csv (labels base3 and headfix, interleaved runs), results/{headfix,base3}_{cells,cold,verify}.csv.

Sort-key scenario (median wall seconds, regexp_extract_all as the SortExec key):

workers	BASE	HEAD before fix	HEAD with fix
1	0.171	0.251 (1.45x slower)	0.166 (3 percent faster than BASE)
8	0.130	0.287 (2.2x slower)	0.120 (8 percent faster than BASE)

Main-matrix cells that regressed (regexp_extract_all, warm, 8192 rows, Mrows/s, BASE3 / HEAD before fix / HEAD with fix):

sharing	w=1	w=2	w=4	w=8
shared	13.1 / 14.9 / 15.3	13.0 / 9.4 / 27.8	10.3 / 9.5 / 52.7	7.1 / 6.3 / 72.4
perworker	13.1 / 14.9 / 15.4	13.0 / 13.9 / 27.4	10.3 / 10.8 / 53.3	7.1 / 7.2 / 74.2

The fix does more than remove the contention: eliminating the per-match allocations lets the function scale nearly linearly to 8 workers in BOTH sharing modes (previously even the uncontended per-worker mode was allocator-bound at about 7 Mrows/s). Single thread: +17 percent vs BASE at 8192 rows (+3 percent vs unfixed HEAD); at 512 rows +89 percent vs BASE but about 12 percent below unfixed HEAD, the cost of rerunning the capture engine at each match start. Alternating-pattern worst case stays within 2 percent of BASE. Output hashes for all functions, batch sizes, and regimes remain byte-identical across BASE, both HEAD builds, and the fixed build (0 mismatches in 12 hash cells x 4 runs).

analyze.py

#!/usr/bin/env python3
"""Merge BASE/HEAD CSVs from the PR #5612 harness and emit markdown tables."""
import csv
import os
import sys

D = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")


def load(path):
    with open(path) as f:
        return list(csv.DictReader(f))


def key(r):
    return (r["function"], r["sharing"], int(r["workers"]), int(r["batch_rows"]), r["regime"])


def main():
    base = {key(r): r for r in load(f"{D}/base_cells.csv")}
    head = {key(r): r for r in load(f"{D}/head_cells.csv")}
    assert set(base) == set(head), "cell mismatch"

    out = []
    out.append("| function | sharing | workers | batch | regime | BASE Mrows/s | HEAD Mrows/s | thrpt Δ | BASE mean/p99 us | HEAD mean/p99 us | mean Δ |")
    out.append("|---|---|---:|---:|---|---:|---:|---:|---:|---:|---:|")
    for k in sorted(base):
        b, h = base[k], head[k]
        bt, ht = float(b["rows_per_s"]) / 1e6, float(h["rows_per_s"]) / 1e6
        bm, hm = float(b["mean_us"]), float(h["mean_us"])
        out.append(
            f"| {k[0]} | {k[1]} | {k[2]} | {k[3]} | {k[4]} "
            f"| {bt:.2f} | {ht:.2f} | {100*(ht-bt)/bt:+.1f}% "
            f"| {bm:.0f} / {float(b['p99_us']):.0f} | {hm:.0f} / {float(h['p99_us']):.0f} "
            f"| {100*(hm-bm)/bm:+.1f}% |"
        )
    print("\n".join(out))

    print("\n### Allocations per batch (count / MB)\n")
    print("| function | sharing | workers | batch | regime | BASE allocs | HEAD allocs | Δ | BASE MB | HEAD MB |")
    print("|---|---|---:|---:|---|---:|---:|---:|---:|---:|")
    for k in sorted(base):
        b, h = base[k], head[k]
        ba, ha = float(b["allocs_per_batch"]), float(h["allocs_per_batch"])
        print(
            f"| {k[0]} | {k[1]} | {k[2]} | {k[3]} | {k[4]} | {ba:.0f} | {ha:.0f} "
            f"| {100*(ha-ba)/ba:+.1f}% | {float(b['bytes_per_batch'])/1e6:.3f} | {float(h['bytes_per_batch'])/1e6:.3f} |"
        )

    print("\n### Cold first-invoke (8192 rows, mean of 30 fresh instances)\n")
    bc = {r["function"]: r for r in load(f"{D}/base_cold.csv")}
    hc = {r["function"]: r for r in load(f"{D}/head_cold.csv")}
    print("| function | BASE 1st us | HEAD 1st us | BASE 2nd us | HEAD 2nd us |")
    print("|---|---:|---:|---:|---:|")
    for fn in bc:
        print(
            f"| {fn} | {float(bc[fn]['first_us_mean']):.0f} | {float(hc[fn]['first_us_mean']):.0f} "
            f"| {float(bc[fn]['second_us_mean']):.0f} | {float(hc[fn]['second_us_mean']):.0f} |"
        )

    bv = {(r["function"], r["batch_rows"], r["regime"]): r["hash"] for r in load(f"{D}/base_verify.csv")}
    hv = {(r["function"], r["batch_rows"], r["regime"]): r["hash"] for r in load(f"{D}/head_verify.csv")}
    bad = [k for k in bv if bv[k] != hv.get(k)]
    print(f"\nCorrectness: {len(bv)} output-hash cells compared; mismatches: {len(bad)}")
    if bad:
        print("MISMATCHES:", bad)
        sys.exit(1)


if __name__ == "__main__":
    main()

base2_cells.csv

base2_cold.csv

base2_verify.csv

base3_cells.csv

base3_cold.csv

base3_verify.csv

base_cells.csv

base_cold.csv

base_verify.csv

Cargo.toml

# Standalone benchmark harness for datafusion-comet PR #5612 (regex pattern cache).
# Path-depends on the repo's spark-expr crate so the same harness builds against
# whichever commit is checked out in the repo.
[package]
name = "bench5612"
version = "0.1.0"
edition = "2021"

[workspace]

[dependencies]
datafusion-comet-spark-expr = { path = "/Users/dustin/code/datafusion-comet/native/spark-expr" }
# Match the repo workspace's pinned versions/features (native/Cargo.toml).
datafusion = { version = "54.1.0", default-features = false, features = [
    "unicode_expressions",
    "crypto_expressions",
    "nested_expressions",
    "parquet",
] }
arrow = { version = "58.4.0", features = ["prettyprint", "ffi", "chrono-tz"] }
regex = "1.12.3"
tokio = { version = "1", features = ["rt-multi-thread"] }

[profile.release]
opt-level = 3
debug = false
lto = "thin"
codegen-units = 1

head2_cells.csv

head2_cold.csv

head2_verify.csv

head_cells.csv

head_cold.csv

head_verify.csv

headfix_cells.csv

headfix_cold.csv

headfix_verify.csv

main.rs

// Multithreaded benchmark for datafusion-comet PR #5612 (per-expression regex PatternCache).
//
// Invokes the native spark-expr scalar UDFs exactly the way the engine does: the UDF is
// built via `create_comet_physical_fun` (the serde entry point, stable across BASE and
// HEAD) and evaluated via `ScalarUDF::invoke_with_args`, with the regex pattern supplied
// as a scalar argument on every invoke (which is why construction-time hoisting is
// impossible and the cache exists). No HEAD-only symbols are used, so the identical
// harness builds against both commits.
//
// Matrix: {regexp_extract, regexp_extract_all, split} x {shared, perworker} x
//         {1,2,4,8 workers} x {512, 8192 rows} x {warm, alternating}.
// Plus: cold first-invoke cost per function, and a correctness pass that hashes outputs
// so BASE and HEAD can be asserted identical on identical seeded input.

use std::alloc::{GlobalAlloc, Layout, System};
use std::hint::black_box;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Barrier};
use std::time::Instant;

use arrow::array::builder::StringBuilder;
use arrow::array::{Array, ArrayRef, ListArray, StringArray};
use arrow::datatypes::{DataType, Field};
use datafusion::common::ScalarValue;
use datafusion::config::ConfigOptions;
use datafusion::logical_expr::{ScalarFunctionArgs, ScalarUDF};
use datafusion::physical_plan::ColumnarValue;
use datafusion::prelude::SessionContext;
use datafusion_comet_spark_expr::create_comet_physical_fun;

// ---------------------------------------------------------------------------
// Counting global allocator (allocation count + requested bytes).
// ---------------------------------------------------------------------------

struct CountingAlloc;

static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0);
static ALLOC_BYTES: AtomicU64 = AtomicU64::new(0);

unsafe impl GlobalAlloc for CountingAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
        ALLOC_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed);
        System.alloc(layout)
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        System.dealloc(ptr, layout)
    }
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
        ALLOC_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed);
        System.alloc_zeroed(layout)
    }
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
        ALLOC_BYTES.fetch_add(new_size as u64, Ordering::Relaxed);
        System.realloc(ptr, layout, new_size)
    }
}

#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;

fn alloc_snapshot() -> (u64, u64) {
    (
        ALLOC_COUNT.load(Ordering::Relaxed),
        ALLOC_BYTES.load(Ordering::Relaxed),
    )
}

// ---------------------------------------------------------------------------
// Deterministic seeded data generation (identical for BASE and HEAD).
// ---------------------------------------------------------------------------

struct Lcg(u64);

impl Lcg {
    fn new(seed: u64) -> Self {
        Lcg(seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1))
    }
    fn next(&mut self) -> u64 {
        self.0 = self
            .0
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.0 >> 33
    }
}

const LETTERS: &[u8] = b"abcdefghijklmnopqrstuvwxyz";

fn rand_word(rng: &mut Lcg, min: usize, max: usize) -> String {
    let len = min + (rng.next() as usize) % (max - min + 1);
    (0..len)
        .map(|_| LETTERS[(rng.next() as usize) % 26] as char)
        .collect()
}

/// Input for regexp_extract / regexp_extract_all: ~10% nulls, ~30% alpha-underscore
/// strings ("foo_bar", matched by the alternating pattern), rest digit-dash strings
/// ("123-4567", matched by the warm pattern).
fn make_extract_batch(rows: usize, seed: u64) -> ArrayRef {
    let mut rng = Lcg::new(seed);
    let mut b = StringBuilder::new();
    for _ in 0..rows {
        let r = rng.next() % 10;
        if r == 0 {
            b.append_null();
        } else if r <= 3 {
            b.append_value(format!(
                "{}_{}",
                rand_word(&mut rng, 3, 8),
                rand_word(&mut rng, 3, 8)
            ));
        } else {
            b.append_value(format!("{:03}-{:04}", rng.next() % 1000, rng.next() % 10000));
        }
    }
    Arc::new(b.finish())
}

/// Input for split: ~10% nulls, rest 3..8 random words joined by mixed delimiters
/// out of ",", ";", "|" (all matched by "[,;|]+"; only ";" and "|" by the alt pattern).
fn make_split_batch(rows: usize, seed: u64) -> ArrayRef {
    let mut rng = Lcg::new(seed);
    let mut b = StringBuilder::new();
    let delims = [",", ";", "|", ";;", ",,"];
    for _ in 0..rows {
        if rng.next() % 10 == 0 {
            b.append_null();
        } else {
            let n = 3 + (rng.next() as usize) % 6;
            let mut s = String::new();
            for i in 0..n {
                if i > 0 {
                    s.push_str(delims[(rng.next() as usize) % delims.len()]);
                }
                s.push_str(&rand_word(&mut rng, 2, 9));
            }
            b.append_value(s);
        }
    }
    Arc::new(b.finish())
}

// ---------------------------------------------------------------------------
// Function descriptors.
// ---------------------------------------------------------------------------

#[derive(Clone, Copy, PartialEq)]
enum InputKind {
    Extract,
    Split,
}

struct FuncSpec {
    name: &'static str,
    input: InputKind,
    warm_pattern: &'static str,
    alt_pattern: &'static str,
    /// Third scalar argument (group index or split limit).
    extra: ScalarValue,
    return_type: DataType,
}

fn func_specs() -> Vec<FuncSpec> {
    let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
    vec![
        FuncSpec {
            name: "regexp_extract",
            input: InputKind::Extract,
            warm_pattern: r"(\d+)-(\d+)",
            alt_pattern: r"([a-z]+)_([a-z]+)",
            extra: ScalarValue::Int32(Some(1)),
            return_type: DataType::Utf8,
        },
        FuncSpec {
            name: "regexp_extract_all",
            input: InputKind::Extract,
            warm_pattern: r"(\d+)-(\d+)",
            alt_pattern: r"([a-z]+)_([a-z]+)",
            extra: ScalarValue::Int32(Some(1)),
            return_type: list_ty.clone(),
        },
        FuncSpec {
            name: "split",
            input: InputKind::Split,
            warm_pattern: r"[,;|]+",
            alt_pattern: r"[;|]+",
            extra: ScalarValue::Int32(Some(-1)),
            return_type: list_ty,
        },
    ]
}

fn make_udf(spec: &FuncSpec, ctx: &SessionContext) -> Arc<ScalarUDF> {
    create_comet_physical_fun(spec.name, spec.return_type.clone(), &ctx.state(), None)
        .expect("create_comet_physical_fun failed")
}

fn invoke(
    udf: &ScalarUDF,
    batch: &ArrayRef,
    pattern: &ScalarValue,
    extra: &ScalarValue,
    return_field: &Arc<Field>,
    config: &Arc<ConfigOptions>,
) -> ColumnarValue {
    let args = ScalarFunctionArgs {
        args: vec![
            ColumnarValue::Array(Arc::clone(batch)),
            ColumnarValue::Scalar(pattern.clone()),
            ColumnarValue::Scalar(extra.clone()),
        ],
        arg_fields: vec![],
        number_rows: batch.len(),
        return_field: Arc::clone(return_field),
        config_options: Arc::clone(config),
    };
    udf.invoke_with_args(args).expect("invoke failed")
}

// ---------------------------------------------------------------------------
// Output hashing for the BASE-vs-HEAD equality check (FNV-1a 64).
// ---------------------------------------------------------------------------

struct Fnv(u64);

impl Fnv {
    fn new() -> Self {
        Fnv(0xcbf29ce484222325)
    }
    fn bytes(&mut self, b: &[u8]) {
        for &x in b {
            self.0 ^= x as u64;
            self.0 = self.0.wrapping_mul(0x100000001b3);
        }
    }
    fn marker(&mut self, m: u8) {
        self.bytes(&[0xFE, m]);
    }
}

fn hash_string_array(h: &mut Fnv, a: &StringArray) {
    for i in 0..a.len() {
        if a.is_null(i) {
            h.marker(0);
        } else {
            h.marker(1);
            h.bytes(a.value(i).as_bytes());
        }
    }
}

fn hash_output(h: &mut Fnv, v: &ColumnarValue) {
    match v {
        ColumnarValue::Array(arr) => match arr.data_type() {
            DataType::Utf8 => {
                hash_string_array(h, arr.as_any().downcast_ref::<StringArray>().unwrap())
            }
            DataType::List(_) => {
                let la = arr.as_any().downcast_ref::<ListArray>().unwrap();
                for i in 0..la.len() {
                    if la.is_null(i) {
                        h.marker(2);
                    } else {
                        h.marker(3);
                        let inner = la.value(i);
                        hash_string_array(
                            h,
                            inner.as_any().downcast_ref::<StringArray>().unwrap(),
                        );
                    }
                }
            }
            other => panic!("unexpected output type {other:?}"),
        },
        ColumnarValue::Scalar(s) => h.bytes(format!("{s:?}").as_bytes()),
    }
}

// ---------------------------------------------------------------------------
// Benchmark cells.
// ---------------------------------------------------------------------------

const POOL_SIZE: usize = 16;
const ROWS_TARGET: usize = 4 * 1024 * 1024;

fn batch_pool(kind: InputKind, rows: usize) -> Arc<Vec<ArrayRef>> {
    let base_seed = match kind {
        InputKind::Extract => 0xE0000 + rows as u64,
        InputKind::Split => 0x50000 + rows as u64,
    };
    Arc::new(
        (0..POOL_SIZE)
            .map(|i| match kind {
                InputKind::Extract => make_extract_batch(rows, base_seed + 31 * i as u64),
                InputKind::Split => make_split_batch(rows, base_seed + 31 * i as u64),
            })
            .collect(),
    )
}

struct CellResult {
    batches: usize,
    rows_total: usize,
    wall_s: f64,
    mean_us: f64,
    p50_us: f64,
    p99_us: f64,
    allocs_per_batch: f64,
    bytes_per_batch: f64,
}

#[allow(clippy::too_many_arguments)]
fn run_cell(
    spec: &FuncSpec,
    ctx: &SessionContext,
    pool: &Arc<Vec<ArrayRef>>,
    shared: bool,
    workers: usize,
    batch_rows: usize,
    alternating: bool,
) -> CellResult {
    let nbatches = (ROWS_TARGET / batch_rows).max(64);
    let return_field = Arc::new(Field::new("out", spec.return_type.clone(), true));
    let config = Arc::new(ConfigOptions::default());
    let p0 = ScalarValue::Utf8(Some(spec.warm_pattern.to_string()));
    let p1 = ScalarValue::Utf8(Some(spec.alt_pattern.to_string()));

    // One UDF for all workers (contended cache) or one per worker (planner-typical).
    let shared_udf = if shared { Some(make_udf(spec, ctx)) } else { None };

    let barrier = Arc::new(Barrier::new(workers + 1));
    let mut handles = Vec::with_capacity(workers);
    for w in 0..workers {
        let udf = shared_udf
            .as_ref()
            .map(Arc::clone)
            .unwrap_or_else(|| make_udf(spec, ctx));
        let pool = Arc::clone(pool);
        let barrier = Arc::clone(&barrier);
        let return_field = Arc::clone(&return_field);
        let config = Arc::clone(&config);
        let (p0, p1) = (p0.clone(), p1.clone());
        let extra = spec.extra.clone();
        // Global batch indices w, w+workers, ... so alternating interleaves across
        // workers on the shared instance.
        let mine: Vec<usize> = (w..nbatches).step_by(workers).collect();
        handles.push(std::thread::spawn(move || {
            // Warmup (also puts warm-regime caches into steady state).
            for gi in mine.iter().take(16) {
                let pat = if alternating && gi % 2 == 1 { &p1 } else { &p0 };
                let out = invoke(
                    &udf,
                    &pool[gi % POOL_SIZE],
                    pat,
                    &extra,
                    &return_field,
                    &config,
                );
                black_box(&out);
            }
            let mut lat = Vec::with_capacity(mine.len());
            barrier.wait();
            for gi in &mine {
                let pat = if alternating && gi % 2 == 1 { &p1 } else { &p0 };
                let t0 = Instant::now();
                let out = invoke(
                    &udf,
                    &pool[gi % POOL_SIZE],
                    pat,
                    &extra,
                    &return_field,
                    &config,
                );
                lat.push(t0.elapsed().as_nanos() as u64);
                black_box(&out);
            }
            lat
        }));
    }

    barrier.wait();
    let (a0, b0) = alloc_snapshot();
    let t0 = Instant::now();
    let mut lats: Vec<u64> = Vec::with_capacity(nbatches);
    for h in handles {
        lats.extend(h.join().expect("worker panicked"));
    }
    let wall_s = t0.elapsed().as_secs_f64();
    let (a1, b1) = alloc_snapshot();

    lats.sort_unstable();
    let n = lats.len();
    let mean_us = lats.iter().sum::<u64>() as f64 / n as f64 / 1000.0;
    let p50_us = lats[n / 2] as f64 / 1000.0;
    let p99_us = lats[(n * 99 / 100).min(n - 1)] as f64 / 1000.0;

    CellResult {
        batches: n,
        rows_total: n * batch_rows,
        wall_s,
        mean_us,
        p50_us,
        p99_us,
        allocs_per_batch: (a1 - a0) as f64 / n as f64,
        bytes_per_batch: (b1 - b0) as f64 / n as f64,
    }
}

// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------

fn main() {
    let mut args = std::env::args().skip(1);
    let label = args.next().expect("usage: bench5612 <label> <out_dir>");
    let out_dir = args.next().expect("usage: bench5612 <label> <out_dir>");
    std::fs::create_dir_all(&out_dir).unwrap();

    let ctx = SessionContext::new();
    let specs = func_specs();

    // ---------------- correctness hashes (untimed) ----------------
    let mut verify_csv = String::from("label,function,batch_rows,regime,hash\n");
    for spec in &specs {
        for &rows in &[512usize, 8192] {
            let pool = batch_pool(spec.input, rows);
            for &alternating in &[false, true] {
                let udf = make_udf(spec, &ctx);
                let return_field = Arc::new(Field::new("out", spec.return_type.clone(), true));
                let config = Arc::new(ConfigOptions::default());
                let p0 = ScalarValue::Utf8(Some(spec.warm_pattern.to_string()));
                let p1 = ScalarValue::Utf8(Some(spec.alt_pattern.to_string()));
                let mut h = Fnv::new();
                for (gi, batch) in pool.iter().enumerate() {
                    let pat = if alternating && gi % 2 == 1 { &p1 } else { &p0 };
                    let out = invoke(&udf, batch, pat, &spec.extra, &return_field, &config);
                    hash_output(&mut h, &out);
                }
                verify_csv.push_str(&format!(
                    "{},{},{},{},{:016x}\n",
                    label,
                    spec.name,
                    rows,
                    if alternating { "alternating" } else { "warm" },
                    h.0
                ));
            }
        }
    }
    std::fs::write(format!("{out_dir}/{label}_verify.csv"), &verify_csv).unwrap();
    eprintln!("[{label}] verify hashes written");

    // ---------------- cold first-invoke cost ----------------
    let mut cold_csv = String::from("label,function,first_us_mean,second_us_mean,reps\n");
    for spec in &specs {
        let pool = batch_pool(spec.input, 8192);
        let return_field = Arc::new(Field::new("out", spec.return_type.clone(), true));
        let config = Arc::new(ConfigOptions::default());
        let p0 = ScalarValue::Utf8(Some(spec.warm_pattern.to_string()));
        let reps = 30usize;
        let (mut first_ns, mut second_ns) = (0u64, 0u64);
        for r in 0..reps {
            let udf = make_udf(spec, &ctx); // fresh planned expression => cold cache
            let batch = &pool[r % POOL_SIZE];
            let t0 = Instant::now();
            black_box(&invoke(&udf, batch, &p0, &spec.extra, &return_field, &config));
            first_ns += t0.elapsed().as_nanos() as u64;
            let t1 = Instant::now();
            black_box(&invoke(&udf, batch, &p0, &spec.extra, &return_field, &config));
            second_ns += t1.elapsed().as_nanos() as u64;
        }
        cold_csv.push_str(&format!(
            "{},{},{:.2},{:.2},{}\n",
            label,
            spec.name,
            first_ns as f64 / reps as f64 / 1000.0,
            second_ns as f64 / reps as f64 / 1000.0,
            reps
        ));
    }
    std::fs::write(format!("{out_dir}/{label}_cold.csv"), &cold_csv).unwrap();
    eprintln!("[{label}] cold costs written");

    // ---------------- main matrix ----------------
    let mut csv = String::from(
        "label,function,sharing,workers,batch_rows,regime,batches,rows_total,wall_s,\
         rows_per_s,mean_us,p50_us,p99_us,allocs_per_batch,bytes_per_batch\n",
    );
    let overall = Instant::now();
    for spec in &specs {
        for &rows in &[512usize, 8192] {
            let pool = batch_pool(spec.input, rows);
            for &shared in &[true, false] {
                for &workers in &[1usize, 2, 4, 8] {
                    for &alternating in &[false, true] {
                        let r = run_cell(spec, &ctx, &pool, shared, workers, rows, alternating);
                        let sharing = if shared { "shared" } else { "perworker" };
                        let regime = if alternating { "alternating" } else { "warm" };
                        csv.push_str(&format!(
                            "{},{},{},{},{},{},{},{},{:.4},{:.0},{:.2},{:.2},{:.2},{:.0},{:.0}\n",
                            label,
                            spec.name,
                            sharing,
                            workers,
                            rows,
                            regime,
                            r.batches,
                            r.rows_total,
                            r.wall_s,
                            r.rows_total as f64 / r.wall_s,
                            r.mean_us,
                            r.p50_us,
                            r.p99_us,
                            r.allocs_per_batch,
                            r.bytes_per_batch
                        ));
                        eprintln!(
                            "[{label}] {} {} w={} b={} {}: {:.0} rows/s (wall {:.2}s)",
                            spec.name,
                            sharing,
                            workers,
                            rows,
                            regime,
                            r.rows_total as f64 / r.wall_s,
                            r.wall_s
                        );
                    }
                }
            }
        }
    }
    std::fs::write(format!("{out_dir}/{label}_cells.csv"), &csv).unwrap();
    eprintln!(
        "[{label}] done in {:.1}s total measured wall",
        overall.elapsed().as_secs_f64()
    );
}

shared_clone.rs

// Task B: isolate the mechanism of the HEAD shared-instance slowdown seen on
// regexp_extract_all warm/8192 (HEAD-shared 6.27 vs HEAD-perworker 7.17 Mrows/s at 8
// workers; 9.38 vs 13.92 at 2 workers; BASE shared == perworker at ~6.7/13.1).
//
// Fact (regex-automata 0.4.16, src/meta/regex.rs:1916-1926): `impl Clone for Regex`
// does Arc::clone on the compiled program but Pool::new for the caches, so every
// clone has a PRIVATE fresh scratch pool. Sharing the compiled program does not share
// scratch state. So the candidate deltas in shared mode are only: the single Mutex
// cache line, the shared Arc refcount cache line touched by every clone/drop, and
// reading the one shared program's memory from all threads.
//
// Modes (each thread does the same number of "invokes"; one invoke = optionally
// obtain a Regex, then run captures_iter over one 8192-row batch, like extract_all):
//   compile_per_invoke  Regex::new every invoke                       (= BASE)
//   shared_mutex        lock shared Mutex slot, compare, clone, unlock (= HEAD shared)
//   private_mutex       same but per-thread Mutex slot                 (= HEAD perworker)
//   shared_clone        clone from one shared Arc<Regex>, no mutex
//   private_clone       clone from a per-thread Regex, no mutex
//   private_hoisted     per-thread Regex used directly, no per-invoke clone
//   shared_hoisted      one shared Regex used directly by all threads (shared pool)
//
// Usage: shared_clone [invokes_per_thread]

use regex::Regex;
use std::hint::black_box;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::time::Instant;

struct Lcg(u64);
impl Lcg {
    fn new(seed: u64) -> Self {
        Lcg(seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1))
    }
    fn next(&mut self) -> u64 {
        self.0 = self
            .0
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.0 >> 33
    }
}

const LETTERS: &[u8] = b"abcdefghijklmnopqrstuvwxyz";

fn rand_word(rng: &mut Lcg, min: usize, max: usize) -> String {
    let len = min + (rng.next() as usize) % (max - min + 1);
    (0..len)
        .map(|_| LETTERS[(rng.next() as usize) % 26] as char)
        .collect()
}

/// Same distribution as the main harness extract input (nulls become skipped rows).
fn make_rows(rows: usize, seed: u64) -> Vec<Option<String>> {
    let mut rng = Lcg::new(seed);
    (0..rows)
        .map(|_| {
            let r = rng.next() % 10;
            if r == 0 {
                None
            } else if r <= 3 {
                Some(format!(
                    "{}_{}",
                    rand_word(&mut rng, 3, 8),
                    rand_word(&mut rng, 3, 8)
                ))
            } else {
                Some(format!("{:03}-{:04}", rng.next() % 1000, rng.next() % 10000))
            }
        })
        .collect()
}

const PATTERN: &str = r"(\d+)-(\d+)";
const ROWS: usize = 8192;

/// One batch worth of extract_all-like work.
fn scan(re: &Regex, rows: &[Option<String>]) -> u64 {
    let mut acc = 0u64;
    for r in rows.iter().flatten() {
        for caps in re.captures_iter(r) {
            if let Some(m) = caps.get(1) {
                acc = acc.wrapping_add(m.len() as u64);
            }
        }
    }
    acc
}

#[derive(Clone, Copy, PartialEq)]
enum Mode {
    CompilePerInvoke,
    SharedMutex,
    PrivateMutex,
    SharedClone,
    PrivateClone,
    PrivateHoisted,
    SharedHoisted,
    /// One true Regex::clone from the shared source per THREAD (hoisted out of the
    /// loop): private scratch pool, but the compiled program is shared. Discriminates
    /// per-search program-sharing contention from per-invoke clone contention.
    SharedCloneHoisted,
    /// Just Regex::clone from the shared source in a loop, no scanning: prices the
    /// clone itself under cross-thread sharing.
    CloneOnlyShared,
    /// Same clone loop from a per-thread source.
    CloneOnlyPrivate,
}

impl Mode {
    fn name(self) -> &'static str {
        match self {
            Mode::CompilePerInvoke => "compile_per_invoke",
            Mode::SharedMutex => "shared_mutex",
            Mode::PrivateMutex => "private_mutex",
            Mode::SharedClone => "shared_clone",
            Mode::PrivateClone => "private_clone",
            Mode::PrivateHoisted => "private_hoisted",
            Mode::SharedHoisted => "shared_hoisted",
            Mode::SharedCloneHoisted => "shared_clone_hoisted",
            Mode::CloneOnlyShared => "clone_only_shared",
            Mode::CloneOnlyPrivate => "clone_only_private",
        }
    }
}

type Slot = Arc<Mutex<Option<(String, Regex)>>>;

fn get_from_slot(slot: &Slot) -> Regex {
    // Same procedure as PatternCache::get_or_compile on the hit path.
    let guard = slot.lock().unwrap();
    if let Some((p, re)) = guard.as_ref() {
        if p == PATTERN {
            return re.clone();
        }
    }
    unreachable!("slot is pre-filled");
}

fn run(mode: Mode, threads: usize, invokes: usize, rows: &Arc<Vec<Option<String>>>) -> (f64, f64) {
    let shared_src = Arc::new(Regex::new(PATTERN).unwrap());
    let shared_slot: Slot = Arc::new(Mutex::new(Some((
        PATTERN.to_string(),
        Regex::new(PATTERN).unwrap(),
    ))));
    let barrier = Arc::new(Barrier::new(threads + 1));
    let sink = Arc::new(AtomicU64::new(0));

    let mut handles = Vec::new();
    for _ in 0..threads {
        let rows = Arc::clone(rows);
        let barrier = Arc::clone(&barrier);
        let shared_src = Arc::clone(&shared_src);
        let shared_slot = Arc::clone(&shared_slot);
        let sink = Arc::clone(&sink);
        handles.push(std::thread::spawn(move || {
            // Per-thread sources, compiled on this thread.
            let private_src = Regex::new(PATTERN).unwrap();
            let private_slot: Slot = Arc::new(Mutex::new(Some((
                PATTERN.to_string(),
                Regex::new(PATTERN).unwrap(),
            ))));
            let hoisted = Regex::new(PATTERN).unwrap();
            // Warmup.
            for _ in 0..4 {
                black_box(scan(&private_src.clone(), &rows));
            }
            let shared_clone_hoisted = Regex::clone(shared_src.as_ref());
            barrier.wait();
            let mut acc = 0u64;
            for _ in 0..invokes {
                acc = acc.wrapping_add(match mode {
                    Mode::CompilePerInvoke => scan(&Regex::new(PATTERN).unwrap(), &rows),
                    Mode::SharedMutex => scan(&get_from_slot(&shared_slot), &rows),
                    Mode::PrivateMutex => scan(&get_from_slot(&private_slot), &rows),
                    // Explicit Regex::clone (an Arc<Regex>.clone() would only bump the Arc).
                    Mode::SharedClone => scan(&Regex::clone(shared_src.as_ref()), &rows),
                    Mode::PrivateClone => scan(&Regex::clone(&private_src), &rows),
                    Mode::PrivateHoisted => scan(&hoisted, &rows),
                    Mode::SharedHoisted => scan(&shared_src, &rows),
                    Mode::SharedCloneHoisted => scan(&shared_clone_hoisted, &rows),
                    Mode::CloneOnlyShared => {
                        let re = Regex::clone(shared_src.as_ref());
                        black_box(&re);
                        0
                    }
                    Mode::CloneOnlyPrivate => {
                        let re = Regex::clone(&private_src);
                        black_box(&re);
                        0
                    }
                });
            }
            sink.fetch_add(acc, Ordering::Relaxed);
        }));
    }
    barrier.wait();
    let t0 = Instant::now();
    for h in handles {
        h.join().unwrap();
    }
    let wall = t0.elapsed().as_secs_f64();
    black_box(sink.load(Ordering::Relaxed));
    let total_invokes = threads * invokes;
    let nonnull_rows = rows.iter().filter(|r| r.is_some()).count();
    let mrows_s = (total_invokes * ROWS) as f64 / wall / 1e6;
    let mean_us = wall * 1e6 / (total_invokes as f64 / threads as f64);
    let _ = nonnull_rows;
    (mrows_s, mean_us)
}

fn main() {
    let invokes: usize = std::env::args()
        .nth(1)
        .map(|s| s.parse().unwrap())
        .unwrap_or(512);
    let rows = Arc::new(make_rows(ROWS, 0xE0000 + 8192));

    println!("mode,threads,invokes_per_thread,mrows_per_s,mean_invoke_us");
    for &mode in &[
        Mode::CompilePerInvoke,
        Mode::SharedMutex,
        Mode::PrivateMutex,
        Mode::SharedClone,
        Mode::PrivateClone,
        Mode::PrivateHoisted,
        Mode::SharedHoisted,
        Mode::SharedCloneHoisted,
        Mode::CloneOnlyShared,
        Mode::CloneOnlyPrivate,
    ] {
        for &threads in &[1usize, 2, 4, 8] {
            let (mrows, mean_us) = run(mode, threads, invokes, &rows);
            println!(
                "{},{},{},{:.2},{:.1}",
                mode.name(),
                threads,
                invokes,
                mrows,
                mean_us
            );
            eprintln!(
                "{} t={} -> {:.2} Mrows/s ({:.0} us/invoke)",
                mode.name(),
                threads,
                mrows,
                mean_us
            );
        }
    }
}

shared_clone_micro.csv

sort_base_all.csv

sort_base_w1.csv

sort_base_w8.csv

sort_extract_w1_rep2.csv

sort_extractall_w8_rep2.csv

sort_fix.csv

sort_head_all.csv

sort_head_w1.csv

sort_head_w8.csv

sort_repro.rs

// Task A: sunchao's sort-key sharing scenario for PR #5612.
//
// Comet passes sort-key expressions directly into DataFusion's SortExec
// (native/core/src/execution/planner.rs: create_sort_expr builds PhysicalSortExpr
// with an arbitrary physical expression; SortExec::new(LexOrdering::new(exprs)...)).
// In DataFusion 54.1.0, ExternalSorter::in_mem_sort_stream (sorts/sort.rs:586) takes
// the spawn_buffered branch (line 633) once the memory reservation reaches
// execution.sort_in_place_threshold_bytes (default 1 MiB): every buffered batch is
// sorted in its own tokio task on the multithread runtime, and all tasks share the
// SAME cloned LexOrdering -> the SAME ScalarUDF instance evaluated concurrently.
//
// This repro builds exactly that plan shape: a single-partition memory source of
// unsorted 8192-row string batches, SortExec (no fetch) whose sort key is the comet
// regexp_extract_all UDF wrapped in a tracking shim that counts concurrent in-flight
// invocations on ONE inner UDF instance. Built via create_comet_physical_fun, so the
// identical source compiles against BASE and HEAD.
//
// Usage: sort_repro <label> <workers> <iters>

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;

use arrow::array::builder::StringBuilder;
use arrow::array::RecordBatch;
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema};
use datafusion::common::{Result as DFResult, ScalarValue};
use datafusion::config::ConfigOptions;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::execution::TaskContext;
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};
use datafusion::physical_expr::expressions::{col, lit};
use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr, ScalarFunctionExpr};
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::{collect, displayable, ExecutionPlan};
use datafusion::prelude::SessionContext;
use datafusion_comet_spark_expr::create_comet_physical_fun;

// ---------------- deterministic data (same generator as the main harness) ----------

struct Lcg(u64);
impl Lcg {
    fn new(seed: u64) -> Self {
        Lcg(seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1))
    }
    fn next(&mut self) -> u64 {
        self.0 = self
            .0
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        self.0 >> 33
    }
}

const LETTERS: &[u8] = b"abcdefghijklmnopqrstuvwxyz";

fn rand_word(rng: &mut Lcg, min: usize, max: usize) -> String {
    let len = min + (rng.next() as usize) % (max - min + 1);
    (0..len)
        .map(|_| LETTERS[(rng.next() as usize) % 26] as char)
        .collect()
}

fn make_batch(schema: &Arc<Schema>, rows: usize, seed: u64) -> RecordBatch {
    let mut rng = Lcg::new(seed);
    let mut b = StringBuilder::new();
    for _ in 0..rows {
        let r = rng.next() % 10;
        if r == 0 {
            b.append_null();
        } else if r <= 3 {
            b.append_value(format!(
                "{}_{}",
                rand_word(&mut rng, 3, 8),
                rand_word(&mut rng, 3, 8)
            ));
        } else {
            b.append_value(format!("{:03}-{:04}", rng.next() % 1000, rng.next() % 10000));
        }
    }
    RecordBatch::try_new(Arc::clone(schema), vec![Arc::new(b.finish())]).unwrap()
}

// ---------------- tracking shim around ONE comet UDF instance ----------------------

#[derive(Default)]
struct Counters {
    in_flight: AtomicU64,
    max_in_flight: AtomicU64,
    invokes: AtomicU64,
}

struct TrackingUdf {
    inner: Arc<ScalarUDF>,
    signature: Signature,
    return_type: DataType,
    counters: Arc<Counters>,
}

impl std::fmt::Debug for TrackingUdf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TrackingUdf").finish()
    }
}

impl PartialEq for TrackingUdf {
    fn eq(&self, other: &Self) -> bool {
        std::ptr::eq(self, other)
    }
}
impl Eq for TrackingUdf {}
impl std::hash::Hash for TrackingUdf {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        "tracked_regexp_extract_all".hash(state);
    }
}

impl ScalarUDFImpl for TrackingUdf {
    fn name(&self) -> &str {
        "tracked_regexp_extract_all"
    }
    fn signature(&self) -> &Signature {
        &self.signature
    }
    fn return_type(&self, _: &[DataType]) -> DFResult<DataType> {
        Ok(self.return_type.clone())
    }
    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
        let cur = self.counters.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
        self.counters.max_in_flight.fetch_max(cur, Ordering::SeqCst);
        self.counters.invokes.fetch_add(1, Ordering::SeqCst);
        let r = self.inner.invoke_with_args(args);
        self.counters.in_flight.fetch_sub(1, Ordering::SeqCst);
        r
    }
}

// ---------------- main --------------------------------------------------------------

const BATCH_ROWS: usize = 8192;
const NUM_BATCHES: usize = 128;

fn build_plan(
    func: &str,
    batches: &[RecordBatch],
    schema: &Arc<Schema>,
    ctx: &SessionContext,
) -> (Arc<dyn ExecutionPlan>, Arc<Counters>) {
    let key_ty = if func == "regexp_extract" {
        DataType::Utf8
    } else {
        DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)))
    };
    // ONE comet UDF instance (one PatternCache in HEAD), shared by every evaluation.
    let inner = create_comet_physical_fun(func, key_ty.clone(), &ctx.state(), None)
        .expect("create_comet_physical_fun");
    let counters = Arc::new(Counters::default());
    let tracked = ScalarUDF::new_from_impl(TrackingUdf {
        inner,
        signature: Signature::variadic_any(Volatility::Immutable),
        return_type: key_ty.clone(),
        counters: Arc::clone(&counters),
    });

    let sort_key = Arc::new(ScalarFunctionExpr::new(
        "tracked_regexp_extract_all",
        Arc::new(tracked),
        vec![
            col("s", schema).unwrap(),
            lit(ScalarValue::Utf8(Some(r"(\d+)-(\d+)".to_string()))),
            lit(ScalarValue::Int32(Some(1))),
        ],
        Arc::new(Field::new("key", key_ty, true)),
        Arc::new(ConfigOptions::default()),
    ));

    // Same construction as Comet's planner: PhysicalSortExpr with the expression as
    // the sort key, SortExec with no fetch, single partition.
    let source =
        MemorySourceConfig::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap();
    let sort_expr = PhysicalSortExpr {
        expr: sort_key,
        options: SortOptions::default(),
    };
    let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
        LexOrdering::new(vec![sort_expr]).unwrap(),
        source,
    ));
    (sort, counters)
}

fn main() {
    let mut args = std::env::args().skip(1);
    let label = args.next().expect("usage: sort_repro <label> <function> <workers> <iters>");
    let func = args.next().expect("function");
    let workers: usize = args.next().expect("workers").parse().unwrap();
    let iters: usize = args.next().expect("iters").parse().unwrap();

    let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
    let batches: Vec<RecordBatch> = (0..NUM_BATCHES)
        .map(|i| make_batch(&schema, BATCH_ROWS, 0xA0000 + 31 * i as u64))
        .collect();
    let input_bytes: usize = batches
        .iter()
        .map(|b| b.get_array_memory_size())
        .sum();

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(workers)
        .enable_all()
        .build()
        .unwrap();
    let ctx = SessionContext::new();

    // Print the physical plan once (workers=1 invocation is enough, but cheap to always do).
    let (plan, _c) = build_plan(&func, &batches, &schema, &ctx);
    eprintln!(
        "input: {} batches x {} rows, {} bytes buffered (sort_in_place_threshold_bytes default = {}; spawn_buffered branch requires reservation >= threshold and >1 batch)",
        NUM_BATCHES,
        BATCH_ROWS,
        input_bytes,
        ConfigOptions::default().execution.sort_in_place_threshold_bytes
    );
    eprintln!("physical plan:\n{}", displayable(plan.as_ref()).indent(true));

    println!("label,function,workers,iter,wall_s,rows,invokes,max_in_flight");
    for iter in 0..iters {
        let (plan, counters) = build_plan(&func, &batches, &schema, &ctx);
        let task_ctx = Arc::new(TaskContext::default());
        let t0 = Instant::now();
        let out = rt
            .block_on(async { collect(plan, task_ctx).await })
            .expect("sort failed");
        let wall = t0.elapsed().as_secs_f64();
        let rows: usize = out.iter().map(|b| b.num_rows()).sum();
        assert_eq!(rows, NUM_BATCHES * BATCH_ROWS);
        println!(
            "{},{},{},{},{:.4},{},{},{}",
            label,
            func,
            workers,
            iter,
            wall,
            rows,
            counters.invokes.load(Ordering::SeqCst),
            counters.max_in_flight.load(Ordering::SeqCst)
        );
    }
}