From ded5a7091d63022011a109ef51a1c457a538ce80 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:05:54 +0200 Subject: [PATCH 01/11] =?UTF-8?q?chore(backlog):=20TASK-44=20Done=20?= =?UTF-8?q?=E2=80=94=20M-cranelift=20merged=20(PR=20#27)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 --- ...degen-backend-behind-the-same-interface.md | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/backlog/tasks/task-44 - Specializer-M-cranelift-codegen-backend-behind-the-same-interface.md b/backlog/tasks/task-44 - Specializer-M-cranelift-codegen-backend-behind-the-same-interface.md index f03e131..bd6468d 100644 --- a/backlog/tasks/task-44 - Specializer-M-cranelift-codegen-backend-behind-the-same-interface.md +++ b/backlog/tasks/task-44 - Specializer-M-cranelift-codegen-backend-behind-the-same-interface.md @@ -1,10 +1,10 @@ --- id: TASK-44 title: 'Specializer M-cranelift: codegen backend behind the same interface' -status: In Progress +status: Done assignee: [] created_date: '2026-07-25 02:31' -updated_date: '2026-07-26 06:05' +updated_date: '2026-07-26 07:55' labels: [] milestone: m-7 dependencies: @@ -23,10 +23,10 @@ Cranelift-jit backend for the imperative IR (design doc §7; cranelift-jit 0.126 ## Acceptance Criteria -- [ ] #1 Every v0-subset prepared query compiles under cranelift and agrees with the interpreter backend on the corpus and on randomized inputs -- [ ] #2 Uncovered ops fall back to the interpreter backend rather than failing prepare -- [ ] #3 p50/p99 ns/call reported at n in {1, 8, 64, 1024} against the interpreter control and the existing native + codegen engines -- [ ] #4 mise gate-specializer green +- [x] #1 Every v0-subset prepared query compiles under cranelift and agrees with the interpreter backend on the corpus and on randomized inputs +- [x] #2 Uncovered ops fall back to the interpreter backend rather than failing prepare +- [x] #3 p50/p99 ns/call reported at n in {1, 8, 64, 1024} against the interpreter control and the existing native + codegen engines +- [x] #4 mise gate-specializer green ## Implementation Plan @@ -38,3 +38,15 @@ Stretch plan (recorded 2026-07-26, design doc §7 + §10): 3. Differential: gen.rs random-IR fuzz interpreter-vs-cranelift (same seeds, byte-identical outputs incl. NaN/-0.0/arena strings); corpus replay + duck_check suite through the cranelift backend; gate green. 4. Bench per §10: baseline the boundary with a no-op f first, then p50/p99 ns/call at n in {1,8,64,1024}, interpreter as control, next to the existing native + codegen engines; committed as a script, numbers into the task. + +## Implementation Notes + + +All four stretches landed in one session (commits 6fd137b/0362b37/222138c on claude/specializer-m-cranelift, PR #27). Backend shape: one JIT'd fn per program, called per row — extern "C" fn(*mut Cx) -> i64 (0 emit / 1 skip / 2 helper trap / 3+k term trap); Cx is repr(C), the generated code reads only trap_flag at offset 0 (checked after every fallible helper). Coverage-first: consts/float arith/int cmp/logic/select/itof/fabs inline in CLIF; everything nontrivial calls extern helpers delegating to the interpreter's own semantic fns (casemap, substr_window, duck_fcmp, DuckF64, arena) — backends cannot drift where they share code. Strings = (off,len) i64 pairs; IR blocks map 1:1 to CLIF blocks+params. CraneliftFn owns a compiled InterpFn (checks, statics, fallback). The 500-seed random-IR differential caught two real bugs pre-landing: load.opt and sload.opt must normalize payloads to type defaults under a false flag. DuckDBInferFn compiles cranelift-first with interpreter fallback (AC #2) + SPECIALIZER_FORCE_INTERP knob + .backend getter; the whole pytest suite (607 + corpus 53/625/0) now exercises the JIT. Bench (scripts/bench_specializer.py, §10 discipline): pydantic boundary dominates end-to-end — noop passthrough within ~20% of a full join; cranelift==interp through Python; both beat native/codegen ~1.5-2x at larger n. Raw compute isolation (backend_compute_bench, ignored test): cranelift 13.2 ns/row vs interp 37.6 ns/row = 2.8x — the JIT win is real and waiting behind the boundary (M-boundary's job). MILESTONE COMPLETE pending review; hard stop before M-boundary. + + +## Final Summary + + +M-cranelift delivered on claude/specializer-m-cranelift (PR #27): cranelift-jit 0.126 backend with FULL IR instruction coverage behind the same interface. Per-row extern "C" ABI; coverage-first op strategy (inline CLIF where trivially safe, shared-semantics extern helpers elsewhere — the two backends physically share the semantic functions). 500-seed interpreter-vs-cranelift random-IR differential green (caught two real payload-normalization bugs before landing); whole pytest suite incl. 678-case corpus replay (53 match / 625 clean-unsupported / 0 FAIL) runs on the JIT; interpreter fallback + force knob wired (AC #2). Bench per design doc §10 at n in {1,8,64,1024} vs interp control + native + codegen: boundary dominates end-to-end (JIT==interp through Python, both ~1.5-2x ahead of the existing engines); raw compute isolated: 13.2 vs 37.6 ns/row = 2.8x for cranelift. Gate green (cargo + 607 pytest, 12 xfail). Next milestone (M-boundary, generated row marshaller) is where the compute win becomes end-to-end visible. + From df5394183d9140d449f840c1362f5e9624320892 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:08:35 +0200 Subject: [PATCH 02/11] chore(backlog): TASK-45 In Progress + M-boundary stretch plan Co-Authored-By: Claude Fable 5 --- ...boundary-generated-row-marshaller-Python-API.md | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md b/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md index 20a8350..b1709ae 100644 --- a/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md +++ b/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md @@ -1,9 +1,10 @@ --- id: TASK-45 title: 'Specializer M-boundary: generated row marshaller + Python API' -status: To Do +status: In Progress assignee: [] created_date: '2026-07-25 02:32' +updated_date: '2026-07-26 08:00' labels: [] milestone: m-7 dependencies: @@ -28,3 +29,14 @@ Wire the specializer into the Python surface: SpecializedTransform (SQLTransform - [ ] #4 Steady-state hot path allocates nothing per call beyond the output objects; arena reset only - [ ] #5 mise gate-specializer green + +## Implementation Plan + + +Stretch plan (recorded 2026-07-26, design doc §3 flag 1 + §10). Measured targets from the M-cranelift bench: per-cell getattr with a fresh name string, per-call buffer allocs, and model_validate per output row — those three ARE the boundary that dominates end-to-end. +1. Marshaller core (src/duckdb, the boundary module): prepare-time interned PyString field names for input attrs and output keys; input rows accepted as dict (get_item on interned key) or pydantic model (getattr on interned name), unboxed type-directed straight into the existing Batch columns — SoA stays: both backends read Batch, the batch is L1-resident at this n, and the doc's own flag-1 argument makes layout irrelevant here (deviation from the AoS row-struct line, noted deliberately); output built model_construct-style (cached bound method + interned keys), never model_validate; RunState + input buffers owned by the fn object behind a Mutex, cleared not dropped per call. SPECIALIZER_GENERIC_BOUNDARY env knob keeps the old generic path runnable for the baseline; a .boundary getter mirrors .backend. +2. SpecializedTransform (sql_transform package): SQLTransform API minus transformer refs — ctor(sql | Template), fit(table, this_model=None) reusing desugar/inline_references/build_state_tables/rewrite_sql then preparing DuckDBInferFn (clear error on transformer refs; records output only, dense raises); infer/infer_batch pass dicts/models straight through, no SimpleNamespace hop. Window aggregates rewrite into static-table equi-joins = v0 subset, so they ride along where the rewrite output parses. +3. Bench per §10: extend scripts/bench_specializer.py — no-op f through generic vs marshalled boundary (AC #2, the marshaller's win as a measured number), end-to-end p50/p99 at n in {1,8,64,1024} vs native + codegen (AC #3); numbers into this ticket. +4. Zero-alloc steady state (AC #4): counting-global-allocator Rust test around the reused-state run path asserting no Rust-side allocation on the second call (arena reset only); marshaller buffer reuse asserted the same way. Gate green (AC #5). + + From 3201f05c822e1874b6b4b7dadab6857db4319112 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:17:59 +0200 Subject: [PATCH 03/11] =?UTF-8?q?feat(specializer):=20generated=20row=20ma?= =?UTF-8?q?rshaller=20=E2=80=94=20interned=20names,=20model=5Fconstruct=20?= =?UTF-8?q?output,=20reused=20state=20+=20span-buffer=20Str=20columns=20(T?= =?UTF-8?q?ASK-45=20stretch=201)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The boundary work knowable at prepare time now happens at prepare time: attribute names interned once, model_construct resolved once, input buffers and RunState owned by the fn and cleared (not dropped) per call. Rows cross as dicts or pydantic models. ColData::Str holds one flat buffer + spans instead of per-cell Strings, which also removes a hidden per-load String clone in the cranelift h_load_str helper. SPECIALIZER_GENERIC_BOUNDARY pins the old path as the bench baseline; .boundary getter mirrors .backend. infer_rows() is the direct hot entry. Co-Authored-By: Claude Fable 5 --- src/duckdb/mod.rs | 214 ++++++++++++++++++++++++++++-- src/specializer/exec/cranelift.rs | 11 +- src/specializer/exec/interp.rs | 2 +- src/specializer/exec/mod.rs | 92 ++++++++++++- src/specializer/exec/testutil.rs | 7 +- 5 files changed, 301 insertions(+), 25 deletions(-) diff --git a/src/duckdb/mod.rs b/src/duckdb/mod.rs index e25dbae..42038e8 100644 --- a/src/duckdb/mod.rs +++ b/src/duckdb/mod.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; use pyo3::prelude::*; -use pyo3::types::PyDict; +use pyo3::types::{PyDict, PyString}; use crate::error::InterpError; use crate::schema; @@ -200,11 +200,159 @@ impl Backend { } } +/// The generated row marshaller (design doc §3 flag 1): everything about the +/// boundary that is knowable at prepare time is done at prepare time — +/// interned attribute-name objects in fixed field order, `model_construct` +/// resolved once, input buffers and run state owned and reused (cleared, not +/// dropped, per call). The generic path stays available behind +/// `SPECIALIZER_GENERIC_BOUNDARY` as the measured baseline. +struct Marshaller { + in_names: Vec>, + out_names: Vec>, + /// `output_model.model_construct`, resolved at build. Outputs come out + /// of the typed engine already conformant; re-validating them per row + /// was the single largest boundary cost. + construct: Py, + cols: Vec, + state: RunState, +} + +impl Marshaller { + fn build( + py: Python<'_>, + in_cols: &[Col], + out_cols: &[Col], + output_model: &Py, + fun: &Backend, + ) -> PyResult { + Ok(Marshaller { + in_names: in_cols + .iter() + .map(|c| PyString::intern(py, &c.name).unbind()) + .collect(), + out_names: out_cols + .iter() + .map(|c| PyString::intern(py, &c.name).unbind()) + .collect(), + construct: output_model.bind(py).getattr("model_construct")?.unbind(), + cols: in_cols.iter().map(|c| ColData::new(c.ty.ty)).collect(), + state: fun.new_state(), + }) + } + + /// The hot path: fill reused columns from row objects (dict or model), + /// run, emit via `model_construct`. Steady-state this allocates only the + /// output objects — buffers and arena reset, never drop. + fn call( + &mut self, + py: Python<'_>, + fun: &Backend, + in_cols: &[Col], + rows: &[Py], + row_table: &str, + ) -> PyResult>> { + for col in &mut self.cols { + col.clear(); + } + for row_obj in rows { + let bound = row_obj.bind(py); + let dict = bound.cast::().ok(); + for ((c, name), col) in in_cols.iter().zip(&self.in_names).zip(&mut self.cols) { + let attr = match dict { + Some(d) => d.get_item(name.bind(py))?.ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err(format!( + "Row for table '{row_table}' is missing attribute '{}'", + c.name + )) + })?, + None => bound.getattr(name.bind(py)).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "Row for table '{row_table}' is missing attribute '{}': {e}", + c.name + )) + })?, + }; + let null = attr.is_none(); + if null && !c.ty.nullable { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "column '{}' is not nullable but a row has None", + c.name + ))); + } + match col { + ColData::I1 { valid, data } => { + valid.push(!null); + data.push(if null { false } else { attr.extract()? }); + } + ColData::I64 { valid, data } => { + valid.push(!null); + data.push(if null { 0 } else { attr.extract()? }); + } + ColData::F64 { valid, data } => { + valid.push(!null); + data.push(if null { 0.0 } else { attr.extract()? }); + } + s @ ColData::Str { .. } => { + if null { + s.push_str_cell(false, ""); + } else { + s.push_str_cell(true, attr.extract()?); + } + } + } + } + } + + // The batch borrows the reused columns for the duration of the run; + // mem::take + restore keeps `Batch` an owning type (an empty Vec + // does not allocate). + let batch = Batch { + rows: rows.len(), + cols: std::mem::take(&mut self.cols), + }; + let res = fun.run(&batch, &mut self.state); + self.cols = batch.cols; + res.map_err(|t| PyErr::from(InterpError::Eval(t.0)))?; + + let construct = self.construct.bind(py); + let mut out = Vec::with_capacity(self.state.emitted); + for r in 0..self.state.emitted { + let d = PyDict::new(py); + for (name, oc) in self.out_names.iter().zip(&self.state.out) { + let k = name.bind(py); + match oc { + OutCol::I1(v) => { + let (ok, x) = v[r]; + d.set_item(k, ok.then_some(x))?; + } + OutCol::I64(v) => { + let (ok, x) = v[r]; + d.set_item(k, ok.then_some(x))?; + } + OutCol::F64(v) => { + let (ok, x) = v[r]; + d.set_item(k, ok.then_some(x))?; + } + OutCol::Str(v) => { + let (ok, s) = v[r]; + d.set_item(k, ok.then(|| self.state.arena.get(s)))?; + } + } + } + out.push(construct.call((), Some(&d))?.unbind()); + } + Ok(out) + } +} + enum Engine { Compiled { fun: Backend, in_cols: Vec, out_cols: Vec, + /// `None` when `SPECIALIZER_GENERIC_BOUNDARY` pinned the generic + /// boundary at construction (the bench baseline). + marsh: Option, }, /// Fixed row dicts from a static-only query, re-validated through the /// output model on every `infer` call. @@ -357,11 +505,25 @@ impl DuckDBInferFn { Some(m) => m, None => synthesize_output_model(py, &prepared.program.out_cols)?, }; + // SPECIALIZER_GENERIC_BOUNDARY pins the pre-marshaller boundary — + // the bench baseline, mirroring SPECIALIZER_FORCE_INTERP. + let marsh = if std::env::var_os("SPECIALIZER_GENERIC_BOUNDARY").is_some() { + None + } else { + Some(Marshaller::build( + py, + &in_cols, + &prepared.program.out_cols, + &output_model, + &fun, + )?) + }; Ok(DuckDBInferFn { engine: Engine::Compiled { fun, in_cols, out_cols: prepared.program.out_cols.clone(), + marsh, }, row_table, output_model, @@ -377,9 +539,20 @@ impl DuckDBInferFn { } } + /// How rows cross the Python boundary: "marshaller" (generated at + /// prepare), "generic" (env-pinned baseline), or "constant". + #[getter] + fn boundary(&self) -> &'static str { + match &self.engine { + Engine::Compiled { marsh: Some(_), .. } => "marshaller", + Engine::Compiled { marsh: None, .. } => "generic", + Engine::Constant { .. } => "constant", + } + } + #[pyo3(signature = (tables=None, **kwargs))] fn infer( - &self, + &mut self, py: Python<'_>, tables: Option>>>, kwargs: Option>, @@ -397,22 +570,37 @@ impl DuckDBInferFn { ))); } let rows = merged.remove(&self.row_table).unwrap_or_default(); + self.run_rows(py, &rows) + } + + /// The direct hot entry: the row table's rows, no table-dict plumbing. + /// `SpecializedTransform.infer_batch` calls this. + fn infer_rows(&mut self, py: Python<'_>, rows: Vec>) -> PyResult>> { + self.run_rows(py, &rows) + } +} - let (fun, in_cols, out_cols) = match &self.engine { +impl DuckDBInferFn { + fn run_rows(&mut self, py: Python<'_>, rows: &[Py]) -> PyResult>> { + let (fun, in_cols, out_cols, marsh) = match &mut self.engine { Engine::Compiled { fun, in_cols, out_cols, - } => (fun, in_cols, out_cols), + marsh, + } => (&*fun, &*in_cols, &*out_cols, marsh), Engine::Constant { rows: fixed } => { let model = self.output_model.bind(py); let mut out = Vec::with_capacity(fixed.len()); - for r in fixed { - out.push(model.call_method1("model_validate", (r,))?.unbind()); + for r in fixed.iter() { + out.push(model.call_method1("model_validate", (&*r,))?.unbind()); } return Ok(out); } }; + if let Some(m) = marsh { + return m.call(py, fun, in_cols, rows, &self.row_table); + } let n = rows.len(); let mut cols: Vec = in_cols @@ -432,11 +620,12 @@ impl DuckDBInferFn { }, Ty::Str => ColData::Str { valid: Vec::with_capacity(n), - data: Vec::with_capacity(n), + buf: String::new(), + spans: Vec::with_capacity(n), }, }) .collect(); - for row_obj in &rows { + for row_obj in rows { let bound = row_obj.bind(py); for (c, col) in in_cols.iter().zip(&mut cols) { let attr = bound.getattr(c.name.as_str()).map_err(|e| { @@ -465,9 +654,12 @@ impl DuckDBInferFn { valid.push(!null); data.push(if null { 0.0 } else { attr.extract()? }); } - ColData::Str { valid, data } => { - valid.push(!null); - data.push(if null { String::new() } else { attr.extract()? }); + s @ ColData::Str { .. } => { + if null { + s.push_str_cell(false, ""); + } else { + s.push_str_cell(true, attr.extract()?); + } } } } diff --git a/src/specializer/exec/cranelift.rs b/src/specializer/exec/cranelift.rs index a7f7622..39fe05f 100644 --- a/src/specializer/exec/cranelift.rs +++ b/src/specializer/exec/cranelift.rs @@ -134,11 +134,12 @@ extern "C" fn h_load_f64(p: *mut Cx, col: i64) -> f64 { extern "C" fn h_load_str(p: *mut Cx, col: i64, len_out: *mut i64) -> i64 { let c = unsafe { cx(p) }; - let s = match &c.input().cols[col as usize] { - ColData::Str { data, .. } => data[c.row].clone(), - _ => unreachable!("load type checked by the verifier"), - }; - let r = c.arena().push_str(&s); + // SAFETY: input and arena are disjoint allocations behind separate raw + // pointers; materializing both references keeps the borrow checker out + // of a copy that never aliases (and skips the old per-load String clone). + let input = unsafe { &*c.input }; + let arena = unsafe { &mut *c.arena }; + let r = arena.push_str(input.cols[col as usize].str_at(c.row)); unsafe { *len_out = r.len as i64 }; r.off as i64 } diff --git a/src/specializer/exec/interp.rs b/src/specializer/exec/interp.rs index 23571a0..d34d8c8 100644 --- a/src/specializer/exec/interp.rs +++ b/src/specializer/exec/interp.rs @@ -1065,7 +1065,7 @@ fn load_payload(c: &ColData, row: usize, arena: &mut Arena) -> RegVal { ColData::I1 { data, .. } => RegVal::I1(data[row]), ColData::I64 { data, .. } => RegVal::I64(data[row]), ColData::F64 { data, .. } => RegVal::F64(data[row]), - ColData::Str { data, .. } => RegVal::Str(arena.push_str(&data[row])), + c @ ColData::Str { .. } => RegVal::Str(arena.push_str(c.str_at(row))), } } diff --git a/src/specializer/exec/mod.rs b/src/specializer/exec/mod.rs index 86f6e90..f2864f9 100644 --- a/src/specializer/exec/mod.rs +++ b/src/specializer/exec/mod.rs @@ -44,13 +44,95 @@ pub struct Batch { } pub enum ColData { - I1 { valid: Vec, data: Vec }, - I64 { valid: Vec, data: Vec }, - F64 { valid: Vec, data: Vec }, - Str { valid: Vec, data: Vec }, + I1 { + valid: Vec, + data: Vec, + }, + I64 { + valid: Vec, + data: Vec, + }, + F64 { + valid: Vec, + data: Vec, + }, + /// Strings live flat in one shared buffer with per-row spans — no + /// per-cell String, so a reused column refills without allocating once + /// capacity is warm (the marshaller's zero-alloc contract). + Str { + valid: Vec, + buf: String, + spans: Vec, + }, } impl ColData { + /// Empty column of type `t`, ready for `push_*` fills. + pub fn new(t: Ty) -> ColData { + match t { + Ty::I1 => ColData::I1 { + valid: Vec::new(), + data: Vec::new(), + }, + Ty::I64 => ColData::I64 { + valid: Vec::new(), + data: Vec::new(), + }, + Ty::F64 => ColData::F64 { + valid: Vec::new(), + data: Vec::new(), + }, + Ty::Str => ColData::Str { + valid: Vec::new(), + buf: String::new(), + spans: Vec::new(), + }, + } + } + + /// Capacity-preserving clear, for reuse across calls. + pub fn clear(&mut self) { + match self { + ColData::I1 { valid, data } => { + valid.clear(); + data.clear(); + } + ColData::I64 { valid, data } => { + valid.clear(); + data.clear(); + } + ColData::F64 { valid, data } => { + valid.clear(); + data.clear(); + } + ColData::Str { valid, buf, spans } => { + valid.clear(); + buf.clear(); + spans.clear(); + } + } + } + + /// Append one cell to a Str column (`""` for a NULL payload). + pub fn push_str_cell(&mut self, ok: bool, s: &str) { + let ColData::Str { valid, buf, spans } = self else { + unreachable!("push_str_cell on a non-Str column"); + }; + valid.push(ok); + let off = buf.len(); + buf.push_str(s); + spans.push(StrRef { off, len: s.len() }); + } + + /// The string payload of row `row` of a Str column. + pub fn str_at(&self, row: usize) -> &str { + let ColData::Str { buf, spans, .. } = self else { + unreachable!("str_at on a non-Str column"); + }; + let StrRef { off, len } = spans[row]; + &buf[off..off + len] + } + pub fn ty(&self) -> Ty { match self { ColData::I1 { .. } => Ty::I1, @@ -65,7 +147,7 @@ impl ColData { ColData::I1 { data, .. } => data.len(), ColData::I64 { data, .. } => data.len(), ColData::F64 { data, .. } => data.len(), - ColData::Str { data, .. } => data.len(), + ColData::Str { spans, .. } => spans.len(), } } diff --git a/src/specializer/exec/testutil.rs b/src/specializer/exec/testutil.rs index 09f26d5..13ccc8a 100644 --- a/src/specializer/exec/testutil.rs +++ b/src/specializer/exec/testutil.rs @@ -26,10 +26,11 @@ pub fn c_f64(vals: &[Option]) -> ColData { } pub fn c_str(vals: &[Option<&str>]) -> ColData { - ColData::Str { - valid: vals.iter().map(|v| v.is_some()).collect(), - data: vals.iter().map(|v| v.unwrap_or("").to_string()).collect(), + let mut col = ColData::new(super::super::ir::Ty::Str); + for v in vals { + col.push_str_cell(v.is_some(), v.unwrap_or("")); } + col } pub fn batch(rows: usize, cols: Vec) -> Batch { From 37a9cd4030e7a3c29e7d492a87b35b224da0dec1 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:19:49 +0200 Subject: [PATCH 04/11] =?UTF-8?q?feat(specializer):=20SpecializedTransform?= =?UTF-8?q?=20=E2=80=94=20the=20authoring=20surface=20served=20by=20the=20?= =?UTF-8?q?specializer=20(TASK-45=20stretch=202)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SQLTransform minus transformer refs: fit() reuses the state-extraction pipeline (window aggregates freeze into per-partition static tables, template refs inline), then prepares through DuckDBInferFn — global and partitioned window aggregates land on cranelift via the equi-join rewrite and agree with SQLTransform row for row. infer/infer_batch take dicts or pydantic models straight through the marshaller, no SimpleNamespace hop. Co-Authored-By: Claude Fable 5 --- sql_transform/__init__.py | 3 +- sql_transform/_specialized.py | 98 ++++++++++++++++++++++++++++++ sql_transform/_specialized_test.py | 62 +++++++++++++++++++ 3 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 sql_transform/_specialized.py create mode 100644 sql_transform/_specialized_test.py diff --git a/sql_transform/__init__.py b/sql_transform/__init__.py index 7f2566d..17a05bc 100644 --- a/sql_transform/__init__.py +++ b/sql_transform/__init__.py @@ -16,11 +16,12 @@ from sql_transform._interpreter import InferFn from sql_transform._rewrite import rewrite_sql from sql_transform._schema import synthesize_this_model +from sql_transform._specialized import SpecializedTransform from sql_transform._sql import find_window_aggregates, parse_and_validate from sql_transform._state import build_state_tables from sql_transform._transformer_ref import resolve_transformer_refs -__all__ = ["InferFn", "SQLTransform"] +__all__ = ["InferFn", "SQLTransform", "SpecializedTransform"] def _to_namespace(row: dict[str, Any] | BaseModel) -> SimpleNamespace: diff --git a/sql_transform/_specialized.py b/sql_transform/_specialized.py new file mode 100644 index 0000000..5a1cf90 --- /dev/null +++ b/sql_transform/_specialized.py @@ -0,0 +1,98 @@ +"""SpecializedTransform — the SQLTransform surface served by the specializer.""" + +from __future__ import annotations + +from string.templatelib import Template +from typing import Any + +import datafusion +import pyarrow as pa +from pydantic import BaseModel + +from sql_transform._compose import desugar_template, inline_references +from sql_transform._interpreter import DuckDBInferFn +from sql_transform._rewrite import rewrite_sql +from sql_transform._schema import synthesize_this_model +from sql_transform._sql import find_window_aggregates, parse_and_validate +from sql_transform._state import build_state_tables + + +class SpecializedTransform: + """SQLTransform minus transformer refs, served by the SQL specializer. + + fit() runs the same state-extraction pipeline as SQLTransform (window + aggregates freeze into per-partition static tables, template refs + inline), then hands the rewritten SQL to the specializer: frontend -> + binding-time analysis -> lowering -> native code, with the static + tables baked in as prepare-time probe structures. + + infer()/infer_batch() take dicts or pydantic models directly and cross + the boundary through the prepare-time-generated row marshaller. + """ + + def __init__(self, sql: str | Template) -> None: + if isinstance(sql, Template): + self._sql, self._refs = desugar_template(sql) + else: + self._sql, self._refs = sql, {} + if any(r.is_transformer for r in self._refs.values()): + raise ValueError( + "SpecializedTransform does not support transformer refs; " + "use SQLTransform" + ) + self._fn: DuckDBInferFn | None = None + + @classmethod + def from_file(cls, path: str) -> SpecializedTransform: + with open(path) as f: + return cls(f.read()) + + def fit( + self, + table: pa.Table, + /, + this_model: type[BaseModel] | None = None, + ) -> SpecializedTransform: + this_model = this_model or synthesize_this_model(table.schema) + tree = parse_and_validate(self._sql) + + ctx = datafusion.SessionContext() + ctx.from_arrow(table, name="__THIS__") + inline = inline_references(tree, self._refs, ctx, table) + windows = find_window_aggregates(tree) + own_state = build_state_tables( + windows, ctx, "__THIS__", join_tables=inline.scoped_state + ) + state_tables = {**inline.scoped_state, **own_state} + rewritten = rewrite_sql( + tree, windows, extra_marker_tables=tuple(inline.scoped_state) + ) + self._fn = DuckDBInferFn( + rewritten, + row_tables={"__THIS__": this_model}, + static_tables=state_tables, + ) + return self + + @property + def backend(self) -> str: + """Execution backend: "cranelift", "interpreter", or "constant".""" + return self._fn_or_raise().backend + + @property + def boundary(self) -> str: + """Boundary path: "marshaller", "generic", or "constant".""" + return self._fn_or_raise().boundary + + def infer(self, row: dict[str, Any] | BaseModel, /) -> BaseModel: + """Single-row inference; returns the typed output model instance.""" + return self._fn_or_raise().infer_rows([row])[0] + + def infer_batch(self, rows: list[dict[str, Any] | BaseModel], /) -> list[BaseModel]: + """Many-rows inference; returns typed output model instances.""" + return self._fn_or_raise().infer_rows(rows) + + def _fn_or_raise(self) -> DuckDBInferFn: + if self._fn is None: + raise RuntimeError("Must call fit() before inference") + return self._fn diff --git a/sql_transform/_specialized_test.py b/sql_transform/_specialized_test.py new file mode 100644 index 0000000..e637315 --- /dev/null +++ b/sql_transform/_specialized_test.py @@ -0,0 +1,62 @@ +"""Tests for SpecializedTransform — the specializer-served authoring surface.""" + +import pyarrow as pa +import pytest +from pydantic import BaseModel + +from sql_transform import SpecializedTransform, SQLTransform + +TRAIN = pa.table({"age": [25, 30, 35], "city": ["a", "b", "a"]}) + + +def test_infer_before_fit_raises_runtime_error(): + t = SpecializedTransform("SELECT age FROM __THIS__") + with pytest.raises(RuntimeError): + t.infer({"age": 1}) + + +def test_plain_projection_dict_and_model_rows(): + t = SpecializedTransform( + "SELECT age * 2 AS a2, upper(city) AS c FROM __THIS__" + ).fit(TRAIN) + assert t.backend == "cranelift" + assert t.boundary == "marshaller" + + class Row(BaseModel): + age: int + city: str + + got_dict = t.infer({"age": 21, "city": "xy"}) + got_model = t.infer(Row(age=21, city="xy")) + for got in (got_dict, got_model): + assert got.a2 == 42 + assert got.c == "XY" + + +@pytest.mark.parametrize( + "sql", + [ + "SELECT age / MEAN(age) OVER () AS age_norm FROM __THIS__", + "SELECT age - MEAN(age) OVER (PARTITION BY city) AS d, city FROM __THIS__", + ], +) +def test_window_aggregates_agree_with_sqltransform(sql): + spec = SpecializedTransform(sql).fit(TRAIN) + ref = SQLTransform(sql).fit(TRAIN) + assert spec.backend == "cranelift" + rows = [{"age": 40, "city": "a"}, {"age": 25, "city": "b"}] + got = spec.infer_batch(rows) + want = ref.infer_batch(rows) + assert [m.model_dump() for m in got] == [m.model_dump() for m in want] + + +def test_transformer_ref_rejected_at_construction(): + class FakeFitted: + n_features_in_ = 1 + + def transform(self, x): + return x + + scaler = FakeFitted() + with pytest.raises(ValueError, match="transformer refs"): + SpecializedTransform(t"SELECT {scaler}(age) AS s FROM __THIS__") From 95e6b7bacbd1572d0ef66ad0f892f6a0d38b768f Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:28:46 +0200 Subject: [PATCH 05/11] =?UTF-8?q?feat(specializer):=20allocation-free=20st?= =?UTF-8?q?eady=20state=20on=20both=20backends=20=E2=80=94=20AC=20#4=20(TA?= =?UTF-8?q?SK-45=20stretch=204)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Counting-allocator tests now pin zero heap allocation per warm run for the interpreter AND cranelift, over the probe/arith fixture and a string-heavy program (case, trim, substr, concat, itos/ftos, parse, compare). What had to change: substr and trim become pure sub-span arithmetic (byte ranges into the input span — no copy at all), case mapping streams char-at-a-time into the arena (Arena::case_map), number->text formats straight into the arena (Arena::push_fmt) with DuckF64 rendering to a stack buffer, string compare/parse read the arena immutably, and h_probe emits values without the per-call Vec + ScalarVal clones. All shared between backends; 500-seed differential and the full 612-test pytest suite stay green. Co-Authored-By: Claude Fable 5 --- src/specializer/exec/cranelift.rs | 154 ++++++++++++++++-------------- src/specializer/exec/interp.rs | 135 +++++++++++++++++--------- src/specializer/exec/mod.rs | 53 ++++++++++ src/specializer/exec/tests.rs | 94 ++++++++++++++++++ 4 files changed, 318 insertions(+), 118 deletions(-) diff --git a/src/specializer/exec/cranelift.rs b/src/specializer/exec/cranelift.rs index 39fe05f..9bcd7f5 100644 --- a/src/specializer/exec/cranelift.rs +++ b/src/specializer/exec/cranelift.rs @@ -36,8 +36,8 @@ use super::super::ir::{ BinOp, CmpPred, Inst, Lit, NumOp1, Program, RoundMode, StaticTy, StrOp1, Term, TrimSide, Ty, }; use super::interp::{ - self, apply_ord, col_valid, substr_range_ok, substr_window, CompileError, DuckF64, InterpFn, - PreparedStatic, + self, apply_ord, col_valid, substr_range_ok, substr_window, trim_bounds, CompileError, DuckF64, + InterpFn, PreparedStatic, }; use super::{casemap, Arena, Batch, ColData, KeyBits, OutCol, RunState, ScalarVal, StrRef, Trap}; @@ -72,16 +72,6 @@ impl Cx { fn statics(&self) -> &[PreparedStatic] { unsafe { std::slice::from_raw_parts(self.statics, self.statics_len) } } - fn get(&mut self, off: i64, len: i64) -> String { - // Owned copy: several helpers read one span and push another; a - // borrow across arena mutation is not expressible here. - self.arena() - .get(StrRef { - off: off as usize, - len: len as usize, - }) - .to_string() - } fn set_trap(&mut self, msg: String) { self.trap_flag = 1; self.trap = Some(Trap(msg)); @@ -218,8 +208,18 @@ extern "C" fn h_fcmp(a: f64, b: f64, pred: i64) -> u8 { extern "C" fn h_scmp(p: *mut Cx, ao: i64, al: i64, bo: i64, bl: i64, pred: i64) -> u8 { let c = unsafe { cx(p) }; - let (a, b) = (c.get(ao, al), c.get(bo, bl)); - apply_ord(decode_pred(pred), a.as_str().cmp(b.as_str())) as u8 + // Comparison only reads — both spans borrow the arena immutably. + let arena = unsafe { &*c.arena }; + let a = arena.get(span(ao, al)); + let b = arena.get(span(bo, bl)); + apply_ord(decode_pred(pred), a.cmp(b)) as u8 +} + +fn span(off: i64, len: i64) -> StrRef { + StrRef { + off: off as usize, + len: len as usize, + } } fn decode_pred(p: i64) -> CmpPred { @@ -257,21 +257,21 @@ extern "C" fn h_ftoi(p: *mut Cx, x: f64, round: i64) -> i64 { extern "C" fn h_itos(p: *mut Cx, v: i64, len_out: *mut i64) -> i64 { let c = unsafe { cx(p) }; - let r = c.arena().push_str(&format!("{v}")); + let r = c.arena().push_fmt(format_args!("{v}")); unsafe { *len_out = r.len as i64 }; r.off as i64 } extern "C" fn h_ftos(p: *mut Cx, v: f64, len_out: *mut i64) -> i64 { let c = unsafe { cx(p) }; - let r = c.arena().push_str(&format!("{}", DuckF64(v))); + let r = c.arena().push_fmt(format_args!("{}", DuckF64(v))); unsafe { *len_out = r.len as i64 }; r.off as i64 } extern "C" fn h_stoi(p: *mut Cx, off: i64, len: i64, valid_out: *mut u8) -> i64 { let c = unsafe { cx(p) }; - let s = c.get(off, len); + let s = unsafe { &*c.arena }.get(span(off, len)); match s.trim_ascii().parse::() { Ok(v) => { unsafe { *valid_out = 1 }; @@ -286,7 +286,7 @@ extern "C" fn h_stoi(p: *mut Cx, off: i64, len: i64, valid_out: *mut u8) -> i64 extern "C" fn h_stof(p: *mut Cx, off: i64, len: i64, valid_out: *mut u8) -> f64 { let c = unsafe { cx(p) }; - let s = c.get(off, len); + let s = unsafe { &*c.arena }.get(span(off, len)); match s.trim_ascii().parse::() { Ok(v) => { unsafe { *valid_out = 1 }; @@ -317,13 +317,12 @@ extern "C" fn h_sconcat(p: *mut Cx, ao: i64, al: i64, bo: i64, bl: i64, len_out: extern "C" fn h_scase(p: *mut Cx, off: i64, len: i64, upper: i64, len_out: *mut i64) -> i64 { let c = unsafe { cx(p) }; - let s = c.get(off, len); - let out: String = if upper != 0 { - s.chars().map(casemap::simple_upper).collect() + let map: fn(char) -> char = if upper != 0 { + casemap::simple_upper } else { - s.chars().map(casemap::simple_lower).collect() + casemap::simple_lower }; - let r = c.arena().push_str(&out); + let r = c.arena().case_map(span(off, len), map); unsafe { *len_out = r.len as i64 }; r.off as i64 } @@ -338,18 +337,16 @@ extern "C" fn h_strim( len_out: *mut i64, ) -> i64 { let c = unsafe { cx(p) }; - let (s, set_s) = (c.get(ao, al), c.get(co, cl)); - let set: Vec = set_s.chars().collect(); - let hit = |ch: char| set.contains(&ch); - let t = match side { - 0 => s.trim_matches(hit), - 1 => s.trim_start_matches(hit), - _ => s.trim_end_matches(hit), - } - .to_string(); - let r = c.arena().push_str(&t); - unsafe { *len_out = r.len as i64 }; - r.off as i64 + let arena = unsafe { &*c.arena }; + let side = match side { + 0 => TrimSide::Both, + 1 => TrimSide::Lead, + _ => TrimSide::Trail, + }; + let rng = trim_bounds(arena.get(span(ao, al)), arena.get(span(co, cl)), side); + // The trimmed value is a subview of the input span — no copy. + unsafe { *len_out = (rng.end - rng.start) as i64 }; + (ao as usize + rng.start) as i64 } extern "C" fn h_ssubstr( @@ -372,11 +369,15 @@ extern "C" fn h_ssubstr( unsafe { *len_out = 0 }; return 0; } - let s = c.get(ao, al); - let out = substr_window(&s, start, (has_len != 0).then_some(len)); - let r = c.arena().push_str(&out); - unsafe { *len_out = r.len as i64 }; - r.off as i64 + let arena = unsafe { &*c.arena }; + let rng = substr_window( + arena.get(span(ao, al)), + start, + (has_len != 0).then_some(len), + ); + // The window is a subview of the input span — no copy. + unsafe { *len_out = (rng.end - rng.start) as i64 }; + (ao as usize + rng.start) as i64 } macro_rules! store_h { @@ -449,14 +450,18 @@ extern "C" fn h_probe( let c = unsafe { cx(p) }; let desc = unsafe { &*desc }; let keys = unsafe { std::slice::from_raw_parts(keys, desc.key_tys.len()) }; - let PreparedStatic::Map { entries } = &c.statics()[desc.static_id] else { + // SAFETY: statics and arena are disjoint allocations behind separate + // raw pointers; independent references let str values append to the + // arena while the matched entry stays borrowed — no clones, no + // per-call Vec (the steady state is allocation-free). + let statics = unsafe { std::slice::from_raw_parts(c.statics, c.statics_len) }; + let arena = unsafe { &mut *c.arena }; + let PreparedStatic::Map { entries } = &statics[desc.static_id] else { unreachable!("static kind checked at compile"); }; // Mirrors interp::cmp_key: canonical f64 bits, arena strings by byte - // order. The arena is only read here, never written, so the borrow is - // fine to hold across the search. - let arena = unsafe { &*c.arena }; + // order. The search only reads the arena. let cmp = |stored: &[KeyBits]| -> std::cmp::Ordering { for (kb, cell) in stored.iter().zip(keys.iter()) { let ord = match kb { @@ -464,10 +469,7 @@ extern "C" fn h_probe( KeyBits::I64(s) => s.cmp(&(cell[0] as i64)), KeyBits::F64(s) => s.cmp(&super::canon_f64_bits(f64::from_bits(cell[0]))), KeyBits::Str(s) => { - let v = arena.get(StrRef { - off: cell[0] as usize, - len: cell[1] as usize, - }); + let v = arena.get(span(cell[0] as i64, cell[1] as i64)); s.as_str().cmp(v) } }; @@ -477,35 +479,39 @@ extern "C" fn h_probe( } std::cmp::Ordering::Equal }; - let found = entries.binary_search_by(|(k, _)| cmp(k)).ok(); - let hit = found.is_some(); - let values: Vec = match found { - Some(idx) => entries[idx].1.clone(), - None => desc - .val_tys - .iter() - .map(|ty| match ty { - Ty::I1 => ScalarVal::I1(false), - Ty::I64 => ScalarVal::I64(0), - Ty::F64 => ScalarVal::F64(0.0), - Ty::Str => ScalarVal::Str(String::new()), - }) - .collect(), - }; - for (i, v) in values.into_iter().enumerate() { - let cell = match v { - ScalarVal::I1(b) => [b as u64, 0], - ScalarVal::I64(x) => [x as u64, 0], - ScalarVal::F64(f) => [f.to_bits(), 0], - ScalarVal::Str(s) => { - let r = c.arena().push_str(&s); - [r.off as u64, r.len as u64] + + match found { + Some(idx) => { + for (i, v) in entries[idx].1.iter().enumerate() { + let cell: Cell = match v { + ScalarVal::I1(b) => [*b as u64, 0], + ScalarVal::I64(x) => [*x as u64, 0], + ScalarVal::F64(f) => [f.to_bits(), 0], + ScalarVal::Str(s) => { + let r = arena.push_str(s); + [r.off as u64, r.len as u64] + } + }; + unsafe { *outs.add(i) = cell }; } - }; - unsafe { *outs.add(i) = cell }; + 1 + } + None => { + for (i, ty) in desc.val_tys.iter().enumerate() { + let cell: Cell = match ty { + Ty::I1 | Ty::I64 => [0, 0], + Ty::F64 => [0f64.to_bits(), 0], + Ty::Str => { + let r = arena.push_str(""); + [r.off as u64, r.len as u64] + } + }; + unsafe { *outs.add(i) = cell }; + } + 0 + } } - hit as u8 } // ------------------------------------------------------------ the JIT'd fn -- diff --git a/src/specializer/exec/interp.rs b/src/specializer/exec/interp.rs index d34d8c8..471aa30 100644 --- a/src/specializer/exec/interp.rs +++ b/src/specializer/exec/interp.rs @@ -497,17 +497,6 @@ fn compile_term(p: &Program, t: &Term, slots: &HashMap) -> CTerm { } } -/// Byte sink over the arena so `write!` formats without intermediate -/// allocation. -struct ArenaWriter<'a>(&'a mut Vec); - -impl std::fmt::Write for ArenaWriter<'_> { - fn write_str(&mut self, s: &str) -> std::fmt::Result { - self.0.extend_from_slice(s.as_bytes()); - Ok(()) - } -} - /// DuckDB's substr window arithmetic (measured 1.5.5), on codepoints — NOT /// grapheme clusters (substr slices inside ZWJ emoji). 1-based virtual /// positions: negative start counts from the end (`start = n + start + 1`), @@ -523,7 +512,7 @@ impl std::fmt::Write for ArenaWriter<'_> { /// a non-negative length runs forward `[rs, rs+len)`, a NEGATIVE length /// slices BACKWARDS `[rs+len, rs)`; `len: None` is the 2-arg rest-of-string /// form. -pub(super) fn substr_window(s: &str, start: i64, len: Option) -> String { +pub(super) fn substr_window(s: &str, start: i64, len: Option) -> std::ops::Range { let n = s.chars().count() as i64; let rs = if start < 0 { (n + start + 1).max(1) @@ -537,12 +526,37 @@ pub(super) fn substr_window(s: &str, start: i64, len: Option) -> String { }; let (lo, hi) = (lo.max(1), hi.min(n + 1)); if hi <= lo { - return String::new(); + return 0..0; } - s.chars() - .skip((lo - 1) as usize) - .take((hi - lo) as usize) - .collect() + // Byte range of the char window — the output is a subview of `s`, so + // callers slice instead of copying. + let skip = (lo - 1) as usize; + let take = (hi - lo) as usize; + let b0 = s + .char_indices() + .nth(skip) + .map(|(i, _)| i) + .unwrap_or(s.len()); + let b1 = s[b0..] + .char_indices() + .nth(take) + .map(|(i, _)| b0 + i) + .unwrap_or(s.len()); + b0..b1 +} + +/// The byte range of `s` that survives trimming `set` chars from the chosen +/// ends — pure arithmetic, the output aliases the input. Membership scans +/// `set` per char (a trim set is a handful of chars; no Vec). +pub(super) fn trim_bounds(s: &str, set: &str, side: TrimSide) -> std::ops::Range { + let hit = |c: char| set.chars().any(|k| k == c); + let t = match side { + TrimSide::Both => s.trim_matches(hit), + TrimSide::Lead => s.trim_start_matches(hit), + TrimSide::Trail => s.trim_end_matches(hit), + }; + let start = t.as_ptr() as usize - s.as_ptr() as usize; + start..start + t.len() } /// The offset/length guard DuckDB applies before the window: values outside @@ -561,9 +575,16 @@ impl std::fmt::Display for DuckF64 { if self.0.is_nan() { return f.write_str("nan"); } - let s = format!("{:?}", self.0); + // Stack-render the shortest round-trip form (≤ 24 bytes for any + // f64) so the hot path never builds a temp String. + let mut buf = StackStr::<32>::default(); + { + use std::fmt::Write; + write!(buf, "{:?}", self.0).expect("f64 debug fits 32 bytes"); + } + let s = buf.as_str(); match s.find('e') { - None => f.write_str(&s), + None => f.write_str(s), Some(pos) => { let exp: i64 = s[pos + 1..].parse().expect("float exponent"); write!( @@ -578,12 +599,36 @@ impl std::fmt::Display for DuckF64 { } } -fn fmt_into_arena(arena: &mut Arena, args: std::fmt::Arguments<'_>) -> StrRef { - let off = arena.0.len(); - let _ = ArenaWriter(&mut arena.0).write_fmt(args); - StrRef { - off, - len: arena.0.len() - off, +/// Fixed-capacity ASCII scratch for `write!` — errors instead of growing. +struct StackStr { + buf: [u8; N], + len: usize, +} + +impl Default for StackStr { + fn default() -> Self { + StackStr { + buf: [0; N], + len: 0, + } + } +} + +impl StackStr { + fn as_str(&self) -> &str { + std::str::from_utf8(&self.buf[..self.len]).expect("writes were valid UTF-8") + } +} + +impl std::fmt::Write for StackStr { + fn write_str(&mut self, s: &str) -> std::fmt::Result { + let b = s.as_bytes(); + if self.len + b.len() > N { + return Err(std::fmt::Error); + } + self.buf[self.len..self.len + b.len()].copy_from_slice(b); + self.len += b.len(); + Ok(()) } } @@ -750,7 +795,7 @@ fn compile_inst(p: &Program, inst: &Inst, slots: &HashMap) -> InstFn { let (dst, a) = (sl(slots, dst), sl(slots, a)); Box::new(move |ctx| { let v = as_i64(ctx.regs[a]); - ctx.regs[dst] = RegVal::Str(fmt_into_arena(ctx.arena, format_args!("{v}"))); + ctx.regs[dst] = RegVal::Str(ctx.arena.push_fmt(format_args!("{v}"))); Ok(()) }) } @@ -758,8 +803,7 @@ fn compile_inst(p: &Program, inst: &Inst, slots: &HashMap) -> InstFn { let (dst, a) = (sl(slots, dst), sl(slots, a)); Box::new(move |ctx| { let v = as_f64(ctx.regs[a]); - ctx.regs[dst] = - RegVal::Str(fmt_into_arena(ctx.arena, format_args!("{}", DuckF64(v)))); + ctx.regs[dst] = RegVal::Str(ctx.arena.push_fmt(format_args!("{}", DuckF64(v)))); Ok(()) }) } @@ -814,9 +858,7 @@ fn compile_inst(p: &Program, inst: &Inst, slots: &HashMap) -> InstFn { StrOp1::Lower => super::casemap::simple_lower, }; Box::new(move |ctx| { - let s = ctx.arena.get(as_str(ctx.regs[a])); - let out: String = s.chars().map(map).collect(); - ctx.regs[dst] = RegVal::Str(ctx.arena.push_str(&out)); + ctx.regs[dst] = RegVal::Str(ctx.arena.case_map(as_str(ctx.regs[a]), map)); Ok(()) }) } @@ -828,16 +870,17 @@ fn compile_inst(p: &Program, inst: &Inst, slots: &HashMap) -> InstFn { } => { let (dst, a, chars) = (sl(slots, dst), sl(slots, a), sl(slots, chars)); Box::new(move |ctx| { - let set: Vec = ctx.arena.get(as_str(ctx.regs[chars])).chars().collect(); - let s = ctx.arena.get(as_str(ctx.regs[a])); - let hit = |c: char| set.contains(&c); - let t = match side { - TrimSide::Both => s.trim_matches(hit), - TrimSide::Lead => s.trim_start_matches(hit), - TrimSide::Trail => s.trim_end_matches(hit), - } - .to_owned(); - ctx.regs[dst] = RegVal::Str(ctx.arena.push_str(&t)); + let sref = as_str(ctx.regs[a]); + let rng = trim_bounds( + ctx.arena.get(sref), + ctx.arena.get(as_str(ctx.regs[chars])), + side, + ); + // The trimmed value is a subview of the input span — no copy. + ctx.regs[dst] = RegVal::Str(StrRef { + off: sref.off + rng.start, + len: rng.end - rng.start, + }); Ok(()) }) } @@ -864,9 +907,13 @@ fn compile_inst(p: &Program, inst: &Inst, slots: &HashMap) -> InstFn { } None => None, }; - let s = ctx.arena.get(as_str(ctx.regs[a])); - let out = substr_window(s, st, ln); - ctx.regs[dst] = RegVal::Str(ctx.arena.push_str(&out)); + let sref = as_str(ctx.regs[a]); + let rng = substr_window(ctx.arena.get(sref), st, ln); + // The window is a subview of the input span — no copy. + ctx.regs[dst] = RegVal::Str(StrRef { + off: sref.off + rng.start, + len: rng.end - rng.start, + }); Ok(()) }) } diff --git a/src/specializer/exec/mod.rs b/src/specializer/exec/mod.rs index f2864f9..73a3850 100644 --- a/src/specializer/exec/mod.rs +++ b/src/specializer/exec/mod.rs @@ -201,6 +201,59 @@ impl Arena { std::str::from_utf8(&self.0[r.off..r.off + r.len]) .expect("arena spans are always whole UTF-8 strings") } + + /// Format directly into the arena — no intermediate String. + pub fn push_fmt(&mut self, args: std::fmt::Arguments<'_>) -> StrRef { + use std::fmt::Write; + let off = self.0.len(); + let _ = ArenaWriter(&mut self.0).write_fmt(args); + StrRef { + off, + len: self.0.len() - off, + } + } + + /// Char-by-char 1:1 case map of `r` into a fresh span. Decodes one char + /// at a time (width from the leading byte, O(1) validation) so no temp + /// String is needed while the arena grows under the read span — offsets + /// stay valid across reallocation where borrows would not. + pub fn case_map(&mut self, r: StrRef, map: fn(char) -> char) -> StrRef { + let off = self.0.len(); + let mut pos = r.off; + let end = r.off + r.len; + let mut buf = [0u8; 4]; + while pos < end { + let w = match self.0[pos] { + 0x00..=0x7f => 1, + 0xc0..=0xdf => 2, + 0xe0..=0xef => 3, + _ => 4, + }; + let ch = std::str::from_utf8(&self.0[pos..pos + w]) + .expect("arena spans are always whole UTF-8 strings") + .chars() + .next() + .expect("non-empty UTF-8 sequence"); + pos += w; + self.0 + .extend_from_slice(map(ch).encode_utf8(&mut buf).as_bytes()); + } + StrRef { + off, + len: self.0.len() - off, + } + } +} + +/// Byte sink over the arena so `write!` formats without intermediate +/// allocation. +struct ArenaWriter<'a>(&'a mut Vec); + +impl std::fmt::Write for ArenaWriter<'_> { + fn write_str(&mut self, s: &str) -> std::fmt::Result { + self.0.extend_from_slice(s.as_bytes()); + Ok(()) + } } /// An owned scalar, used by static structures and test expectations. diff --git a/src/specializer/exec/tests.rs b/src/specializer/exec/tests.rs index 04353eb..4cbe6b8 100644 --- a/src/specializer/exec/tests.rs +++ b/src/specializer/exec/tests.rs @@ -285,6 +285,100 @@ fn steady_state_run_allocates_nothing() { } let delta = alloc_count() - before; assert_eq!(delta, 0, "steady-state run heap-allocated {delta} time(s)"); + + // The cranelift backend honors the same contract — its helpers share + // the interpreter's semantic functions, so they must also share the + // no-alloc property (TASK-45 AC #4). + let statics = vec![StaticData::Map(vec![( + vec![KeyBits::Str("a".into())], + vec![ScalarVal::F64(10.0)], + )])]; + let cf = super::cranelift::compile(&p, statics).unwrap(); + let mut cst = cf.new_state(); + cf.run(&input, &mut cst).unwrap(); + cf.run(&input, &mut cst).unwrap(); + + let before = alloc_count(); + for _ in 0..5 { + cf.run(&input, &mut cst).unwrap(); + } + let delta = alloc_count() - before; + assert_eq!( + delta, 0, + "cranelift steady state heap-allocated {delta} time(s)" + ); +} + +/// The string surface — case map, trim, substr, concat, int/float text, +/// parse, compare — is arena-only in steady state on BOTH backends. These +/// ops all used to build temp Strings per row (TASK-45 AC #4). +#[test] +fn steady_state_string_ops_allocate_nothing() { + let p = built( + r#" +fn stringy(in: batch{s: str, t: str, n: i64}, out: batch{a: str, b: str}) { +entry: + %s = load in.s + %t = load in.t + %up = supper %s + %lo = slower %up + %set = const.str " x" + %tr = strim.both %lo, %set + %one = const.i64 1 + %three = const.i64 3 + %sub = ssubstr %tr, %one, %three + %cat = sconcat %sub, %up + %n = load in.n + %ns = itos %n + %cat2 = sconcat %cat, %ns + store out.a, %cat2 + %ok, %iv = stoi.opt %ns + %f = itof %iv + %fs = ftos %f + %eq = scmp.eq %s, %t + %sel = select %eq, %fs, %tr + %sel2 = select %ok, %sel, %up + store out.b, %sel2 + emit +} +"#, + ); + let input = batch( + 3, + vec![ + c_str(&[Some(" héLLo x"), Some("wörld"), Some("")]), + c_str(&[Some("wörld"), Some("wörld"), Some("a")]), + c_i64(&[Some(42), Some(-7), Some(0)]), + ], + ); + + let f = compile(&p, vec![]).unwrap(); + let mut st = f.new_state(); + f.run(&input, &mut st).unwrap(); + f.run(&input, &mut st).unwrap(); + let before = alloc_count(); + for _ in 0..5 { + f.run(&input, &mut st).unwrap(); + } + let delta = alloc_count() - before; + assert_eq!( + delta, 0, + "interp string steady state allocated {delta} time(s)" + ); + + let cf = super::cranelift::compile(&p, vec![]).unwrap(); + let mut cst = cf.new_state(); + cf.run(&input, &mut cst).unwrap(); + cf.run(&input, &mut cst).unwrap(); + let before = alloc_count(); + for _ in 0..5 { + cf.run(&input, &mut cst).unwrap(); + } + let delta = alloc_count() - before; + assert_eq!( + delta, 0, + "cranelift string steady state allocated {delta} time(s)" + ); } // ----------------------------------------- adversarial-pass regressions -- From f15829836938041fa5b9bc7c4f48df95340a6859 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:36:13 +0200 Subject: [PATCH 06/11] feat(specializer): slot-fill output + boundary bench with generic baseline (TASK-45 stretch 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measurement killed an assumption: pydantic v2's model_construct is pure-Python and SLOWER than model_validate (1432ns vs 882ns per row, pydantic 2.13). The marshaller now builds output rows the way the design doc meant — object.__new__ + direct slot writes (__dict__, __pydantic_fields_set__, extra, private), 491ns measured — and the bench gained a 'generic' engine (SPECIALIZER_GENERIC_BOUNDARY) as the AC #2 baseline. Results: marshaller beats the generic boundary 1.9-2.5x on the no-op f; the specializer beats the shipping native engine 3.7-4.2x end-to-end at n=1024 and 1.5-2.8x at n=1; the JIT-vs- interp compute gap is finally visible through Python. Co-Authored-By: Claude Fable 5 --- scripts/bench_specializer.py | 28 +++++++++++++++------ src/duckdb/mod.rs | 47 ++++++++++++++++++++++++++++++------ 2 files changed, 59 insertions(+), 16 deletions(-) diff --git a/scripts/bench_specializer.py b/scripts/bench_specializer.py index 9f7836f..350eeb9 100644 --- a/scripts/bench_specializer.py +++ b/scripts/bench_specializer.py @@ -1,16 +1,20 @@ """Specializer ns/call bench — the design doc §10 measurement discipline. Reports p50/p99 ns per call at n in {1, 8, 64, 1024} for: - boundary — a no-op passthrough (SELECT a FROM __THIS__) through the - generic pydantic path: the floor every engine pays today - cranelift — the specializer's JIT backend (default) + cranelift — the specializer's JIT backend through the generated row + marshaller (the default production path) interp — the same programs on the interpreter backend (the control; - SPECIALIZER_FORCE_INTERP=1 in a subprocess) + SPECIALIZER_FORCE_INTERP=1 in a subprocess), also marshalled + generic — the JIT backend behind the PRE-marshaller generic boundary + (SPECIALIZER_GENERIC_BOUNDARY=1): per-cell getattr with + fresh name strings, per-call buffers, model_validate per + output row. generic vs cranelift on the noop case IS the + marshaller's win (TASK-45 AC #2). native — the existing DataFusion-semantics InferFn codegen — the existing Python codegen engine Run: `uv run python scripts/bench_specializer.py [--json out.json]` -The interp control re-execs this script in a subprocess so the env knob is +The env-knob engines re-exec this script in a subprocess so the knob is set before the module builds its functions. """ @@ -79,7 +83,11 @@ def main(): dim = pa.table({"id": list(range(60)), "name": [f"n{i}" for i in range(60)]}) engine = os.environ.get("BENCH_ENGINE", "cranelift") - builders = {"cranelift": build_specializer, "interp": build_specializer} + builders = { + "cranelift": build_specializer, + "interp": build_specializer, + "generic": build_specializer, + } if engine in ("native", "codegen"): builders = {"native": build_native, "codegen": build_codegen} @@ -90,9 +98,11 @@ def main(): except Exception as e: # noqa: BLE001 -- engines differ in coverage results[case] = {"error": str(e)[:120]} continue - if engine in ("cranelift", "interp"): + if engine in ("cranelift", "interp", "generic"): want = "interpreter" if engine == "interp" else "cranelift" assert fn.backend == want, f"{case}: backend is {fn.backend}, wanted {want}" + wantb = "generic" if engine == "generic" else "marshaller" + assert fn.boundary == wantb, f"{case}: boundary {fn.boundary} != {wantb}" per_n = {} for n in NS: objs = [model(**r) for r in rows_of(n)] @@ -105,11 +115,13 @@ def main(): def orchestrate(): """Run every engine (interp in a subprocess for the env knob), merge.""" merged = {} - for engine in ("cranelift", "interp", "native", "codegen"): + for engine in ("cranelift", "interp", "generic", "native", "codegen"): env = os.environ.copy() env["BENCH_ENGINE"] = engine if engine == "interp": env["SPECIALIZER_FORCE_INTERP"] = "1" + if engine == "generic": + env["SPECIALIZER_GENERIC_BOUNDARY"] = "1" out = subprocess.run( # noqa: S603 -- fixed argv [sys.executable, __file__, "--engine-run"], env=env, diff --git a/src/duckdb/mod.rs b/src/duckdb/mod.rs index 42038e8..a20dd6a 100644 --- a/src/duckdb/mod.rs +++ b/src/duckdb/mod.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyString}; +use pyo3::types::{PyDict, PySet, PyString}; use crate::error::InterpError; use crate::schema; @@ -209,10 +209,21 @@ impl Backend { struct Marshaller { in_names: Vec>, out_names: Vec>, - /// `output_model.model_construct`, resolved at build. Outputs come out - /// of the typed engine already conformant; re-validating them per row - /// was the single largest boundary cost. - construct: Py, + /// Output rows are built by filling pydantic v2's instance slots + /// directly (`object.__new__` + `object.__setattr__` of `__dict__`, + /// `__pydantic_fields_set__`, `__pydantic_extra__`, + /// `__pydantic_private__`) — what `model_construct` does, minus its + /// pure-Python per-field loop (measured 2026-07-26 on pydantic 2.13: + /// construct 1432ns > validate 882ns > slot fill 491ns per row). + /// Assumes a plain model: no `extra="allow"`, no private attrs — + /// true of synthesized output models; supplied ones are v0-trusted. + model: Py, + object_new: Py, + object_setattr: Py, + s_dict: Py, + s_fields_set: Py, + s_extra: Py, + s_private: Py, cols: Vec, state: RunState, } @@ -225,6 +236,7 @@ impl Marshaller { output_model: &Py, fun: &Backend, ) -> PyResult { + let object = PyModule::import(py, "builtins")?.getattr("object")?; Ok(Marshaller { in_names: in_cols .iter() @@ -234,7 +246,13 @@ impl Marshaller { .iter() .map(|c| PyString::intern(py, &c.name).unbind()) .collect(), - construct: output_model.bind(py).getattr("model_construct")?.unbind(), + model: output_model.clone_ref(py), + object_new: object.getattr("__new__")?.unbind(), + object_setattr: object.getattr("__setattr__")?.unbind(), + s_dict: PyString::intern(py, "__dict__").unbind(), + s_fields_set: PyString::intern(py, "__pydantic_fields_set__").unbind(), + s_extra: PyString::intern(py, "__pydantic_extra__").unbind(), + s_private: PyString::intern(py, "__pydantic_private__").unbind(), cols: in_cols.iter().map(|c| ColData::new(c.ty.ty)).collect(), state: fun.new_state(), }) @@ -314,7 +332,14 @@ impl Marshaller { self.cols = batch.cols; res.map_err(|t| PyErr::from(InterpError::Eval(t.0)))?; - let construct = self.construct.bind(py); + let model = self.model.bind(py); + let object_new = self.object_new.bind(py); + let object_setattr = self.object_setattr.bind(py); + let s_dict = self.s_dict.bind(py); + let s_fields_set = self.s_fields_set.bind(py); + let s_extra = self.s_extra.bind(py); + let s_private = self.s_private.bind(py); + let none = py.None(); let mut out = Vec::with_capacity(self.state.emitted); for r in 0..self.state.emitted { let d = PyDict::new(py); @@ -339,7 +364,13 @@ impl Marshaller { } } } - out.push(construct.call((), Some(&d))?.unbind()); + let fields_set = PySet::new(py, self.out_names.iter().map(|n| n.bind(py)))?; + let inst = object_new.call1((model,))?; + object_setattr.call1((&inst, s_dict, &d))?; + object_setattr.call1((&inst, s_fields_set, fields_set))?; + object_setattr.call1((&inst, s_extra, &none))?; + object_setattr.call1((&inst, s_private, &none))?; + out.push(inst.unbind()); } Ok(out) } From 240a56c4ec3b6c0fb5d7329735f61196762fb7d3 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 09:38:30 +0200 Subject: [PATCH 07/11] chore(backlog): TASK-45 all acceptance criteria checked + implementation notes Co-Authored-By: Claude Fable 5 --- ...dary-generated-row-marshaller-Python-API.md | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md b/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md index b1709ae..83005ce 100644 --- a/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md +++ b/backlog/tasks/task-45 - Specializer-M-boundary-generated-row-marshaller-Python-API.md @@ -4,7 +4,7 @@ title: 'Specializer M-boundary: generated row marshaller + Python API' status: In Progress assignee: [] created_date: '2026-07-25 02:32' -updated_date: '2026-07-26 08:00' +updated_date: '2026-07-26 09:10' labels: [] milestone: m-7 dependencies: @@ -23,11 +23,11 @@ Wire the specializer into the Python surface: SpecializedTransform (SQLTransform ## Acceptance Criteria -- [ ] #1 SpecializedTransform fit/infer/infer_batch works end-to-end on the v0 subset with dict and pydantic-model rows -- [ ] #2 No-op-f boundary baseline reported: generic pydantic path vs generated marshaller, p50/p99 at n in {1, 8, 64, 1024} -- [ ] #3 End-to-end p50/p99 vs the current native and codegen engines reported -- [ ] #4 Steady-state hot path allocates nothing per call beyond the output objects; arena reset only -- [ ] #5 mise gate-specializer green +- [x] #1 SpecializedTransform fit/infer/infer_batch works end-to-end on the v0 subset with dict and pydantic-model rows +- [x] #2 No-op-f boundary baseline reported: generic pydantic path vs generated marshaller, p50/p99 at n in {1, 8, 64, 1024} +- [x] #3 End-to-end p50/p99 vs the current native and codegen engines reported +- [x] #4 Steady-state hot path allocates nothing per call beyond the output objects; arena reset only +- [x] #5 mise gate-specializer green ## Implementation Plan @@ -40,3 +40,9 @@ Stretch plan (recorded 2026-07-26, design doc §3 flag 1 + §10). Measured targe 4. Zero-alloc steady state (AC #4): counting-global-allocator Rust test around the reused-state run path asserting no Rust-side allocation on the second call (arena reset only); marshaller buffer reuse asserted the same way. Gate green (AC #5). +## Implementation Notes + + +All four stretches landed on claude/specializer-m-boundary. The marshaller (src/duckdb/mod.rs) does at prepare time everything knowable at prepare time: interned attribute-name PyStrings in fixed field order, input buffers + RunState owned and cleared-not-dropped per call, dict rows via get_item and model rows via getattr on the interned names, output rows by direct pydantic-v2 slot fill. Two assumptions died by measurement: (1) pydantic's literal model_construct API is pure-Python and SLOWER than model_validate (1432 vs 882 ns/row on 2.13) — the shipped path is object.__new__ + object.__setattr__ of __dict__/__pydantic_fields_set__/__pydantic_extra__/__pydantic_private__ at 491 ns, semantically equal (eq, fields_set, assignment all verified); (2) the design doc's AoS row structs buy nothing at L1-resident n — the marshaller fills the existing SoA Batch directly (deliberate deviation, noted in the plan). AC #4 forced real work: ColData::Str became one flat buffer + spans (killing per-cell Strings AND a hidden per-load clone in the JIT's h_load_str), substr/trim became pure sub-span arithmetic, case mapping streams into the arena (Arena::case_map), number→text formats via Arena::push_fmt with DuckF64 on a stack buffer, and h_probe emits without its per-call Vec + ScalarVal clones — all shared between backends, pinned by counting-allocator tests over a probe/arith fixture and a string-heavy program on BOTH backends (the remaining per-call allocs are the output objects themselves plus pyo3's input-list Vec, i.e. the AC's "beyond the output objects"). SpecializedTransform (sql_transform/_specialized.py) reuses the whole fit pipeline minus transformer refs (ValueError at ctor), so window aggregates ride the equi-join rewrite onto cranelift — parity with SQLTransform asserted row-for-row; WHERE stays rejected at the authoring surface (parse_and_validate), while raw DuckDBInferFn keeps it. Bench (scripts/bench_specializer.py, +generic engine via SPECIALIZER_GENERIC_BOUNDARY): noop p50 marshaller vs generic = 1.1/2.1µs at n=1, 468/1192µs at n=1024 (1.9-2.5x, AC #2); vs shipping engines at n=1024 the specializer is 3.7-4.2x faster than native and 2.4-3.8x than codegen (AC #3); cranelift-vs-interp is now visible end-to-end (arith 339 vs 391µs at n=1024). SPECIALIZER_GENERIC_BOUNDARY + .boundary getter mirror the FORCE_INTERP pattern; infer_rows() is the direct hot entry SpecializedTransform uses. + + From 1a6360bbcc12da754f849f33b37ef09c3afb01e8 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 10:04:00 +0200 Subject: [PATCH 08/11] =?UTF-8?q?fix(specializer):=20adversarial-review=20?= =?UTF-8?q?fixes=20=E2=80=94=20supplied-model=20validate=20semantics,=20di?= =?UTF-8?q?ct=20rows=20on=20the=20generic=20baseline,=20reentrancy=20fallb?= =?UTF-8?q?ack=20(TASK-45)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A 17-agent adversarial fleet confirmed 9 findings reducing to three root causes, all fixed with regression tests: 1. The slot-fill output silently broke user-SUPPLIED output models (validators skipped, defaults dropped to AttributeError, extra=allow and private attrs crashed). Slot fill now applies only to the synthesized plain model; a supplied model goes through model_validate per row — master semantics exactly. 2. SPECIALIZER_GENERIC_BOUNDARY dropped dict-row support, making the baseline reject documented inputs with a misleading missing- attribute error. The generic fill now takes the same dict path. 3. infer's move to &mut self regressed same-object reentrancy (a row property calling infer mid-marshal raised 'Already borrowed'). infer is &self again with the marshaller in a RefCell; a reentrant call finds it borrowed and completes via the per-call generic path. Bench spot-check after the fixes: unchanged (noop 1.1µs at n=1, ~2.6x vs generic at n=1024). Co-Authored-By: Claude Fable 5 --- src/duckdb/mod.rs | 91 +++++++++++++++++++++++--------- src/specializer/exec/interp.rs | 1 - tests/test_duckdb_interpreter.py | 63 ++++++++++++++++++++++ 3 files changed, 129 insertions(+), 26 deletions(-) diff --git a/src/duckdb/mod.rs b/src/duckdb/mod.rs index a20dd6a..86c1a94 100644 --- a/src/duckdb/mod.rs +++ b/src/duckdb/mod.rs @@ -7,6 +7,7 @@ //! is only the Python boundary — schema extraction on the way in, map //! materialization for the join probes, output-model rows on the way out. +use std::cell::RefCell; use std::collections::HashMap; use pyo3::prelude::*; @@ -209,15 +210,19 @@ impl Backend { struct Marshaller { in_names: Vec>, out_names: Vec>, - /// Output rows are built by filling pydantic v2's instance slots - /// directly (`object.__new__` + `object.__setattr__` of `__dict__`, - /// `__pydantic_fields_set__`, `__pydantic_extra__`, + /// Output rows for the SYNTHESIZED model are built by filling pydantic + /// v2's instance slots directly (`object.__new__` + `object.__setattr__` + /// of `__dict__`, `__pydantic_fields_set__`, `__pydantic_extra__`, /// `__pydantic_private__`) — what `model_construct` does, minus its /// pure-Python per-field loop (measured 2026-07-26 on pydantic 2.13: - /// construct 1432ns > validate 882ns > slot fill 491ns per row). - /// Assumes a plain model: no `extra="allow"`, no private attrs — - /// true of synthesized output models; supplied ones are v0-trusted. + /// construct 1432ns > validate 882ns > slot fill 491ns per row). The + /// slot fill is sound only for that plain model (required scalar + /// fields, no validators/defaults/extra/private) — so a user-SUPPLIED + /// output model goes through `validate` below instead and keeps full + /// pydantic semantics (adversarial-review finding, 2026-07-26). model: Py, + /// `output_model.model_validate`, present iff the model was supplied. + validate: Option>, object_new: Py, object_setattr: Py, s_dict: Py, @@ -234,6 +239,7 @@ impl Marshaller { in_cols: &[Col], out_cols: &[Col], output_model: &Py, + supplied: bool, fun: &Backend, ) -> PyResult { let object = PyModule::import(py, "builtins")?.getattr("object")?; @@ -247,6 +253,11 @@ impl Marshaller { .map(|c| PyString::intern(py, &c.name).unbind()) .collect(), model: output_model.clone_ref(py), + validate: if supplied { + Some(output_model.bind(py).getattr("model_validate")?.unbind()) + } else { + None + }, object_new: object.getattr("__new__")?.unbind(), object_setattr: object.getattr("__setattr__")?.unbind(), s_dict: PyString::intern(py, "__dict__").unbind(), @@ -364,6 +375,12 @@ impl Marshaller { } } } + if let Some(v) = &self.validate { + // Supplied model: full pydantic semantics (validators, + // defaults, coercion, extra/private) — master parity. + out.push(v.bind(py).call1((&d,))?.unbind()); + continue; + } let fields_set = PySet::new(py, self.out_names.iter().map(|n| n.bind(py)))?; let inst = object_new.call1((model,))?; object_setattr.call1((&inst, s_dict, &d))?; @@ -382,8 +399,13 @@ enum Engine { in_cols: Vec, out_cols: Vec, /// `None` when `SPECIALIZER_GENERIC_BOUNDARY` pinned the generic - /// boundary at construction (the bench baseline). - marsh: Option, + /// boundary at construction (the bench baseline). RefCell so infer + /// stays `&self`: a reentrant call (a row property calling infer on + /// the same object mid-marshal) finds the cell borrowed and falls + /// through to the per-call generic path instead of erroring — + /// master behavior (adversarial-review finding, 2026-07-26). The + /// pyclass is unsendable, so single-threaded RefCell suffices. + marsh: Option>, }, /// Fixed row dicts from a static-only query, re-validated through the /// output model on every `infer` call. @@ -531,8 +553,10 @@ impl DuckDBInferFn { } }, }; + let supplied = output_model.is_some(); let output_model = match output_model { - // Supplied models are trusted as-is in v0 (no shape validation). + // Supplied models are trusted as-is in v0 (no shape validation) + // and keep full model_validate semantics per row. Some(m) => m, None => synthesize_output_model(py, &prepared.program.out_cols)?, }; @@ -541,13 +565,14 @@ impl DuckDBInferFn { let marsh = if std::env::var_os("SPECIALIZER_GENERIC_BOUNDARY").is_some() { None } else { - Some(Marshaller::build( + Some(RefCell::new(Marshaller::build( py, &in_cols, &prepared.program.out_cols, &output_model, + supplied, &fun, - )?) + )?)) }; Ok(DuckDBInferFn { engine: Engine::Compiled { @@ -583,7 +608,7 @@ impl DuckDBInferFn { #[pyo3(signature = (tables=None, **kwargs))] fn infer( - &mut self, + &self, py: Python<'_>, tables: Option>>>, kwargs: Option>, @@ -606,31 +631,35 @@ impl DuckDBInferFn { /// The direct hot entry: the row table's rows, no table-dict plumbing. /// `SpecializedTransform.infer_batch` calls this. - fn infer_rows(&mut self, py: Python<'_>, rows: Vec>) -> PyResult>> { + fn infer_rows(&self, py: Python<'_>, rows: Vec>) -> PyResult>> { self.run_rows(py, &rows) } } impl DuckDBInferFn { - fn run_rows(&mut self, py: Python<'_>, rows: &[Py]) -> PyResult>> { - let (fun, in_cols, out_cols, marsh) = match &mut self.engine { + fn run_rows(&self, py: Python<'_>, rows: &[Py]) -> PyResult>> { + let (fun, in_cols, out_cols, marsh) = match &self.engine { Engine::Compiled { fun, in_cols, out_cols, marsh, - } => (&*fun, &*in_cols, &*out_cols, marsh), + } => (fun, in_cols, out_cols, marsh), Engine::Constant { rows: fixed } => { let model = self.output_model.bind(py); let mut out = Vec::with_capacity(fixed.len()); for r in fixed.iter() { - out.push(model.call_method1("model_validate", (&*r,))?.unbind()); + out.push(model.call_method1("model_validate", (r,))?.unbind()); } return Ok(out); } }; - if let Some(m) = marsh { - return m.call(py, fun, in_cols, rows, &self.row_table); + if let Some(cell) = marsh { + // A reentrant call (row property re-entering infer mid-marshal) + // finds the cell borrowed and takes the generic path below. + if let Ok(mut m) = cell.try_borrow_mut() { + return m.call(py, fun, in_cols, rows, &self.row_table); + } } let n = rows.len(); @@ -658,13 +687,25 @@ impl DuckDBInferFn { .collect(); for row_obj in rows { let bound = row_obj.bind(py); + // Dict rows are part of the API surface; the baseline path must + // accept the same inputs as the marshaller, differing only in + // cost (adversarial-review finding, 2026-07-26). + let dict = bound.cast::().ok(); for (c, col) in in_cols.iter().zip(&mut cols) { - let attr = bound.getattr(c.name.as_str()).map_err(|e| { - pyo3::exceptions::PyValueError::new_err(format!( - "Row for table '{}' is missing attribute '{}': {e}", - self.row_table, c.name - )) - })?; + let attr = match dict { + Some(d) => d.get_item(c.name.as_str())?.ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err(format!( + "Row for table '{}' is missing attribute '{}'", + self.row_table, c.name + )) + })?, + None => bound.getattr(c.name.as_str()).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "Row for table '{}' is missing attribute '{}': {e}", + self.row_table, c.name + )) + })?, + }; let null = attr.is_none(); if null && !c.ty.nullable { return Err(pyo3::exceptions::PyValueError::new_err(format!( diff --git a/src/specializer/exec/interp.rs b/src/specializer/exec/interp.rs index 471aa30..2a699c6 100644 --- a/src/specializer/exec/interp.rs +++ b/src/specializer/exec/interp.rs @@ -24,7 +24,6 @@ //! normalizes even if the input batch carries garbage in invalid slots. use std::collections::HashMap; -use std::fmt::Write as _; use super::super::ir::verify::{verify, VerifyError}; use super::super::ir::{ diff --git a/tests/test_duckdb_interpreter.py b/tests/test_duckdb_interpreter.py index 443be3e..eacc85d 100644 --- a/tests/test_duckdb_interpreter.py +++ b/tests/test_duckdb_interpreter.py @@ -449,3 +449,66 @@ def test_static_only_backend_is_constant(): static_tables={"dim": static({"v": "int"}, [{"v": 1}, {"v": 2}])}, ) assert fn.backend == "constant" + + +# --------------------------------------------------------- M-boundary: +# the generated row marshaller. These pin the adversarial-review fixes +# (2026-07-26): supplied output models keep full pydantic semantics, the +# generic baseline accepts the same inputs, and reentrancy degrades to the +# generic path instead of erroring. + + +def test_supplied_output_model_keeps_validate_semantics(): + from pydantic import BaseModel, field_validator + + class Out(BaseModel): + x: float # engine emits int; validate coerces + note: str = "default" # not in the projection; validate fills + + @field_validator("x") + @classmethod + def clamp(cls, v): + return min(v, 10.0) + + fn = DuckDBInferFn( + "SELECT k + 1 AS x FROM __THIS__", + row_tables={"__THIS__": _row_model({"k": "int"})}, + static_tables={}, + output_model=Out, + ) + assert fn.boundary == "marshaller" + (m,) = fn.infer_rows([{"k": 41}]) + assert m.x == 10.0 # validator ran AND coerced to float + assert m.note == "default" # default applied + assert m.model_dump() == {"x": 10.0, "note": "default"} + + +def test_generic_boundary_accepts_dict_rows(monkeypatch): + monkeypatch.setenv("SPECIALIZER_GENERIC_BOUNDARY", "1") + fn = DuckDBInferFn( + "SELECT k * 2 AS d FROM __THIS__", + row_tables={"__THIS__": _row_model({"k": "int"})}, + static_tables={}, + ) + assert fn.boundary == "generic" + assert fn.infer_rows([{"k": 21}])[0].d == 42 + + +def test_reentrant_infer_falls_back_instead_of_erroring(): + fn = DuckDBInferFn( + "SELECT k * 2 AS d FROM __THIS__", + row_tables={"__THIS__": _row_model({"k": "int"})}, + static_tables={}, + ) + inner: list = [] + + class Row: + @property + def k(self): + if not inner: + inner.append(fn.infer_rows([{"k": 5}])[0].d) + return 7 + + (m,) = fn.infer_rows([Row()]) + assert m.d == 14 + assert inner == [10] # the nested call completed via the generic path From 588f6f1f0b5514eaa6b641e431f6bc58940fa93c Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 11:09:05 +0200 Subject: [PATCH 09/11] =?UTF-8?q?feat(bench):=20realistic=20serving-path?= =?UTF-8?q?=20bench=20=E2=80=94=20scenario=20package,=207-engine=20harness?= =?UTF-8?q?,=20three-way=20parity=20gate=20(titanic=20+=20fraud=20so=20far?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wide-table Kaggle-style inference paths as repo assets: titanic (10->20, string/CASE-heavy) and IEEE-CIS-style fraud (32->38, five join tables, NULL-indicator flags) — each with a handcrafted-Python twin and exact three-way parity (specializer == DuckDB multiset == handcrafted) gated in pytest before any timing. Engines: spec / interp / generic / native / codegen / duckdb-per-call / python(+dict floor). First numbers: spec beats the handcrafted twin 1.8-1.9x at n=1024; duckdb-per-call is ms- scale; native and codegen cannot build either scenario (IS NULL projections unsupported there). house_prices + store_sales to follow. Co-Authored-By: Claude Fable 5 --- benchmarks/__init__.py | 0 benchmarks/bench_serving.py | 190 +++++++++++ benchmarks/serving_scenarios/__init__.py | 154 +++++++++ benchmarks/serving_scenarios/fraud_txn.py | 381 ++++++++++++++++++++++ benchmarks/serving_scenarios/titanic.py | 293 +++++++++++++++++ pyproject.toml | 3 + tests/test_serving_scenarios.py | 33 ++ 7 files changed, 1054 insertions(+) create mode 100644 benchmarks/__init__.py create mode 100644 benchmarks/bench_serving.py create mode 100644 benchmarks/serving_scenarios/__init__.py create mode 100644 benchmarks/serving_scenarios/fraud_txn.py create mode 100644 benchmarks/serving_scenarios/titanic.py create mode 100644 tests/test_serving_scenarios.py diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/benchmarks/bench_serving.py b/benchmarks/bench_serving.py new file mode 100644 index 0000000..a93ecde --- /dev/null +++ b/benchmarks/bench_serving.py @@ -0,0 +1,190 @@ +"""Realistic serving-path bench — wide tables, involved feature engineering. + +The scenarios (benchmarks/serving_scenarios/) reproduce the inference paths +of famous tabular-ML problems. Engines compared, p50/p99 ns per call at +n in {1, 8, 64, 1024}: + + spec — the specializer: cranelift + generated marshaller (ours) + interp — interpreter backend, same marshaller (backend control; + SPECIALIZER_FORCE_INTERP=1) + generic — cranelift behind the pre-marshaller boundary (previous + boundary architecture; SPECIALIZER_GENERIC_BOUNDARY=1) + native — the previous DataFusion-semantics InferFn + codegen — the previous Python-codegen engine + duckdb — DuckDB itself per call (statics pre-materialized as native + tables; per call: Arrow batch from row dicts -> register -> + execute -> fetch) + python — the handcrafted twin: what an engineer would hand-write for a + microservice, returning the same typed output models + python_dict — same, returning plain dicts (the absolute floor; JSON only) + +A three-way parity gate (specializer == DuckDB == handcrafted) runs before +any timing; a scenario that disagrees aborts the bench. + +Run: uv run python -m benchmarks.bench_serving [--json out.json] +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import time + +from benchmarks import serving_scenarios as sc + +NS = [1, 8, 64, 1024] + +ENGINE_ENV = { + "spec": {}, + "interp": {"SPECIALIZER_FORCE_INTERP": "1"}, + "generic": {"SPECIALIZER_GENERIC_BOUNDARY": "1"}, +} + + +def scenario_names(): + """All scenarios, or a comma-separated BENCH_SCENARIOS subset.""" + sub = os.environ.get("BENCH_SCENARIOS") + return sub.split(",") if sub else sc.NAMES + + +def quantiles(samples): + s = sorted(samples) + return s[len(s) // 2], s[min(len(s) - 1, int(len(s) * 0.99))] + + +def bench_call(f, n): + iters = max(30, min(2000, 2_000_000 // max(n, 1))) + f() # warm + out = [] + for _ in range(iters): + t0 = time.perf_counter_ns() + f() + out.append(time.perf_counter_ns() - t0) + p50, p99 = quantiles(out) + return {"p50_ns": p50, "p99_ns": p99} + + +def build_callers(mod, engines): + """-> {engine: callable(rows) or None if the engine can't serve it}.""" + statics = mod.make_statics(sc.SEED) + model = sc.row_model(mod.ROW_SCHEMA) + callers = {} + + if any(e in engines for e in ("spec", "interp", "generic")): + fn = sc.build_spec_fn(mod, statics) + for e in ("spec", "interp", "generic"): + if e in engines: + callers[e] = fn.infer_rows + + for eng, cls_path in ( + ("native", "sql_transform._interpreter.InferFn"), + ("codegen", "sql_transform._codegen.CodegenFn"), + ): + if eng not in engines: + continue + mod_path, cls_name = cls_path.rsplit(".", 1) + cls = getattr(__import__(mod_path, fromlist=[cls_name]), cls_name) + try: + f = cls(mod.SQL, row_tables={"__THIS__": model}, static_tables=statics) + callers[eng] = lambda rows, f=f: f.infer({"__THIS__": rows}) + except Exception as e: # noqa: BLE001 -- engines differ in coverage + callers[eng] = str(e)[:140] + + if "duckdb" in engines: + duck = sc.duckdb_server(mod, statics) + callers["duckdb"] = duck + + if "python" in engines or "python_dict" in engines: + hand = mod.handcrafted(statics) + out_model = sc.build_spec_fn(mod, statics).output_model + if "python" in engines: + callers["python"] = lambda rows: [out_model(**hand(r)) for r in rows] + if "python_dict" in engines: + callers["python_dict"] = lambda rows: [hand(r) for r in rows] + + return callers + + +def engine_run(engines): + results = {} + for mod in map(sc.load, scenario_names()): + callers = build_callers(mod, engines) + per = {} + for eng, call in callers.items(): + if isinstance(call, str): + per[eng] = {"error": call} + continue + per[eng] = {} + for n in NS: + rows = mod.make_rows(sc.SEED + 2, n) + # Engines take model objects on their classic surface; + # spec-family and duckdb/python take dicts natively. Feed + # each what its real caller would: dict rows for + # spec/duckdb/python, model objects for native/codegen + # (their only supported input shape). + if eng in ("native", "codegen"): + model = sc.row_model(mod.ROW_SCHEMA) + rows = [model(**r) for r in rows] + per[eng][n] = bench_call(lambda c=call, r=rows: c(r), n) + results[mod.NAME] = per + print(json.dumps(results)) + + +def orchestrate(): + print("parity gate:", flush=True) + for mod in map(sc.load, scenario_names()): + problems = sc.verify_parity(mod) + tag = "ok" if not problems else "FAIL" + print(f" {mod.NAME:<14} {tag}") + if problems: + print("\n".join(" " + p for p in problems)) + sys.exit(1) + + groups = { + "main": ["spec", "native", "codegen", "duckdb", "python", "python_dict"], + "interp": ["interp"], + "generic": ["generic"], + } + merged: dict = {} + for group, engines in groups.items(): + env = os.environ.copy() + env.update(ENGINE_ENV.get(group if group != "main" else "spec", {})) + env["BENCH_ENGINES"] = ",".join(engines) + out = subprocess.run( # noqa: S603 -- fixed argv + [sys.executable, "-m", "benchmarks.bench_serving", "--engine-run"], + env=env, + capture_output=True, + text=True, + check=True, + ) + for scenario, per in json.loads(out.stdout.strip().splitlines()[-1]).items(): + merged.setdefault(scenario, {}).update(per) + + order = ["spec", "interp", "generic", "native", "codegen", "duckdb", "python"] + hdr = f"{'scenario':<14} {'engine':<9}" + "".join(f"{f'n={n}':>15}" for n in NS) + print("\n" + hdr) + print("-" * len(hdr)) + for scenario, per in merged.items(): + for eng in order: + r = per.get(eng) + if r is None: + continue + if "error" in r: + print(f"{scenario:<14} {eng:<9} unsupported: {r['error'][:70]}") + continue + cells = "".join(f"{r[str(n)]['p50_ns']:>13,}ns" for n in NS) + print(f"{scenario:<14} {eng:<9}{cells}") + if "--json" in sys.argv: + path = sys.argv[sys.argv.index("--json") + 1] + with open(path, "w", encoding="utf-8") as f: + json.dump(merged, f, indent=1) + print(f"wrote {path}") + + +if __name__ == "__main__": + if "--engine-run" in sys.argv: + engine_run(os.environ["BENCH_ENGINES"].split(",")) + else: + orchestrate() diff --git a/benchmarks/serving_scenarios/__init__.py b/benchmarks/serving_scenarios/__init__.py new file mode 100644 index 0000000..ad93beb --- /dev/null +++ b/benchmarks/serving_scenarios/__init__.py @@ -0,0 +1,154 @@ +"""Serving-path benchmark scenarios — famous tabular-ML inference paths. + +Each module reproduces the SERVING shape of a well-known Kaggle-style +solution: a wide input row -> scalar feature expressions + LEFT JOINs to +fitted static tables (lookup dims and pre-materialized target/frequency +encodings — that is what a mean encoding IS at serve time). + +Module contract (see any scenario module): + NAME, KAGGLE, N_INPUT_COLS, N_OUTPUT_COLS, ROW_SCHEMA, SQL, + make_statics(seed) -> dict[str, pa.Table], + make_rows(seed, n) -> list[dict], + handcrafted(statics) -> Callable[[dict], dict] + +The standing trust anchor is three-way parity, enforced by +tests/test_serving_scenarios.py and re-checked by the bench harness: +specializer output == DuckDB itself == the handcrafted Python twin. +""" + +from __future__ import annotations + +import importlib +from collections.abc import Callable +from typing import Any + +import duckdb +import pyarrow as pa +from pydantic import create_model + +NAMES = ["titanic", "house_prices", "fraud_txn", "store_sales"] + +SEED = 20260726 +_PY = {"int": int, "float": float, "str": str, "bool": bool} +_ARROW = { + "int": pa.int64(), + "float": pa.float64(), + "str": pa.string(), + "bool": pa.bool_(), +} + + +def load(name: str): + return importlib.import_module(f"benchmarks.serving_scenarios.{name}") + + +def all_scenarios(): + return [load(n) for n in NAMES] + + +def row_model(schema: dict[str, str]): + fields: dict[str, Any] = {} + for name, spec in schema.items(): + if spec.endswith("?"): + fields[name] = (_PY[spec[:-1]] | None, None) + else: + fields[name] = (_PY[spec], ...) + return create_model("Row", **fields) + + +def arrow_schema(schema: dict[str, str]) -> pa.Schema: + return pa.schema( + pa.field(n, _ARROW[s.rstrip("?")], nullable=s.endswith("?")) + for n, s in schema.items() + ) + + +def rows_table(mod, rows: list[dict]) -> pa.Table: + return pa.Table.from_pylist(rows, schema=arrow_schema(mod.ROW_SCHEMA)) + + +def build_spec_fn(mod, statics: dict[str, pa.Table]): + """The specializer serving fn (env knobs select backend/boundary).""" + from sql_transform._interpreter import DuckDBInferFn + + return DuckDBInferFn( + mod.SQL, + row_tables={"__THIS__": row_model(mod.ROW_SCHEMA)}, + static_tables=statics, + ) + + +def duckdb_server( + mod, statics: dict[str, pa.Table] +) -> Callable[[list[dict]], list[dict]]: + """The 'just run DuckDB' baseline: statics materialized ONCE as native + tables (DuckDB gets the same prepare-time courtesy as everyone else); + each call pays what a DuckDB-backed service would — build an Arrow + batch from the row dicts, register, execute, fetch.""" + con = duckdb.connect() + for name, table in statics.items(): + con.register(f"__arrow_{name}", table) + # S608: table names come from our own scenario modules, not input. + ddl = f'CREATE TABLE "{name}" AS SELECT * FROM "__arrow_{name}"' # noqa: S608 + con.execute(ddl) + con.unregister(f"__arrow_{name}") + schema = arrow_schema(mod.ROW_SCHEMA) + + def call(rows: list[dict]) -> list[dict]: + con.register("__THIS__", pa.Table.from_pylist(rows, schema=schema)) + out = con.execute(mod.SQL).to_arrow_table().to_pylist() + con.unregister("__THIS__") + return out + + return call + + +def verify_parity(mod, n: int = 300) -> list[str]: + """specializer == DuckDB == handcrafted, exact (repr-level floats). + Returns human-readable mismatches; empty list = trusted.""" + statics = mod.make_statics(SEED) + rows = mod.make_rows(SEED + 1, n) + + fn = build_spec_fn(mod, statics) + got_spec = [m.model_dump() for m in fn.infer_rows(rows)] + got_duck = duckdb_server(mod, statics)(rows) + hand = mod.handcrafted(statics) + got_hand = [hand(r) for r in rows] + + problems: list[str] = [] + if fn.backend != "cranelift": + problems.append(f"{mod.NAME}: backend is {fn.backend}, not cranelift") + + def norm(row: dict) -> tuple: + return tuple((k, repr(v)) for k, v in row.items()) + + # DuckDB reorders rows after hash joins (the specializer preserves input + # order — a contract, not an accident): multiset equality, like the + # corpus replay. The handcrafted twin is positional. + if sorted(map(norm, got_spec)) != sorted(map(norm, got_duck)): + spec_only = sorted(map(norm, got_spec)) + duck_only = sorted(map(norm, got_duck)) + first = next( + ( + i + for i, (a, b) in enumerate(zip(spec_only, duck_only, strict=False)) + if a != b + ), + min(len(spec_only), len(duck_only)), + ) + problems.append( + f"{mod.NAME} vs duckdb: multiset mismatch at sorted index {first}: " + f"{spec_only[first] if first < len(spec_only) else ''} != " + f"{duck_only[first] if first < len(duck_only) else ''}" + ) + if len(got_hand) != len(got_spec): + problems.append( + f"{mod.NAME} vs handcrafted: {len(got_spec)} vs {len(got_hand)} rows" + ) + else: + for i, (a, b) in enumerate(zip(got_spec, got_hand, strict=True)): + if norm(a) != norm(b): + problems.append(f"{mod.NAME} vs handcrafted row {i}: {a!r} != {b!r}") + if len(problems) > 5: + return problems + return problems diff --git a/benchmarks/serving_scenarios/fraud_txn.py b/benchmarks/serving_scenarios/fraud_txn.py new file mode 100644 index 0000000..676c797 --- /dev/null +++ b/benchmarks/serving_scenarios/fraud_txn.py @@ -0,0 +1,381 @@ +"""fraud_txn — IEEE-CIS-style transaction-fraud serving benchmark scenario. + +Serving shape of the famous IEEE-CIS Fraud Detection winning solutions: +uid/frequency/target encodings materialized at prepare time become plain +LEFT-JOIN lookup tables at serve time; everything else is scalar math over +the raw transaction row. +""" + +import math +import random +from collections.abc import Callable +from typing import Any + +import pyarrow as pa + +NAME = "fraud_txn" +KAGGLE = ( + "IEEE-CIS Fraud Detection (Kaggle 2019, Deotte/Yakovlev winning-solution " + "serving shape): card1/addr1/email-domain frequency+target encodings as " + "join tables, amt ratio/delta vs card1 mean, cents feature, D-column NULL " + "flags, hour/day-of-week from unix ts, amount buckets, ProductCD one-hots." +) + +N_INPUT_COLS = 32 +N_OUTPUT_COLS = 38 + +ROW_SCHEMA: dict[str, str] = { + "txn_id": "int", + "transaction_amt": "float", + "product_cd": "str", + "card1": "int", + "card2": "int?", + "card3": "int?", + "card4": "str?", + "card5": "int?", + "card6": "str?", + "addr1": "int?", + "addr2": "int?", + "dist1": "float?", + "dist2": "float?", + "p_email_domain": "str?", + "r_email_domain": "str?", + "c1": "int", + "c2": "int", + "c3": "int", + "c4": "int", + "c5": "int", + "c6": "int", + "d1": "int?", + "d2": "int?", + "d3": "int?", + "d4": "int?", + "transaction_dt": "int", + "m1": "bool?", + "m2": "bool?", + "m3": "bool?", + "device_type": "str?", + "v1": "float?", + "v2": "float?", +} + +_DOMAINS = [ + "gmail.com", + "yahoo.com", + "hotmail.com", + "outlook.com", + "aol.com", + "anonymous.com", + "protonmail.com", + "icloud.com", + "comcast.net", + "live.com", + "msn.com", + "att.net", + "verizon.net", + "ymail.com", + "mail.com", + "optonline.net", + "cox.net", + "charter.net", + "earthlink.net", + "juno.com", +] +# serving traffic also carries domains never seen at prepare time -> join miss +_SERVE_DOMAINS = _DOMAINS + ["qq.com", "rocketmail.com", "protonmail.ch"] +_PRODUCTS = ["W", "C", "R", "H", "S"] + + +def make_statics(seed: int) -> dict[str, pa.Table]: + rng = random.Random(seed * 7919 + 1) + + card_ids = sorted(rng.sample(range(1000, 1500), 400)) + card1_stats = pa.table( + { + "card1_id": card_ids, + "card1_amt_mean": [round(rng.uniform(8.0, 420.0), 2) for _ in card_ids], + "card1_txn_cnt": [rng.randint(1, 5000) for _ in card_ids], + "card1_fraud_rate": [round(rng.uniform(0.001, 0.35), 6) for _ in card_ids], + } + ) + + addr_ids = sorted(rng.sample(range(100, 540), 150)) + addr1_stats = pa.table( + { + "addr1_id": addr_ids, + "addr1_fraud_rate": [round(rng.uniform(0.001, 0.25), 6) for _ in addr_ids], + "addr1_txn_cnt": [rng.randint(1, 20000) for _ in addr_ids], + } + ) + + p_email_stats = pa.table( + { + "p_domain": list(_DOMAINS), + "p_email_fraud_rate": [round(rng.uniform(0.002, 0.4), 6) for _ in _DOMAINS], + "p_email_freq": [rng.randint(50, 200000) for _ in _DOMAINS], + } + ) + + r_domains = _DOMAINS[:14] # recipient encoding covers fewer domains + r_email_stats = pa.table( + { + "r_domain": list(r_domains), + "r_email_fraud_rate": [ + round(rng.uniform(0.002, 0.4), 6) for _ in r_domains + ], + } + ) + + pcd_stats = pa.table( + { + "pcd": list(_PRODUCTS), + "pcd_fraud_rate": [round(rng.uniform(0.01, 0.12), 6) for _ in _PRODUCTS], + } + ) + + return { + "card1_stats": card1_stats, + "addr1_stats": addr1_stats, + "p_email_stats": p_email_stats, + "r_email_stats": r_email_stats, + "pcd_stats": pcd_stats, + } + + +def make_rows(seed: int, n: int) -> list[dict]: + rng = random.Random(seed * 104729 + 3) + base_ts = 1_700_000_000 + rows: list[dict[str, Any]] = [] + for i in range(n): + if rng.random() < 0.2: # whole-dollar transactions (the cents trick) + amt = float(rng.choice([20, 25, 30, 50, 75, 100, 150, 200, 300, 500])) + else: + amt = round(math.exp(rng.gauss(3.8, 0.9)), 2) + p_dom = None if rng.random() < 0.16 else rng.choice(_SERVE_DOMAINS) + if p_dom is not None and rng.random() < 0.3: + r_dom: str | None = p_dom # purchaser == recipient is common + else: + r_dom = None if rng.random() < 0.45 else rng.choice(_SERVE_DOMAINS) + rows.append( + { + "txn_id": 3_000_000 + i, + "transaction_amt": amt, + "product_cd": rng.choices(_PRODUCTS, weights=[65, 10, 8, 9, 8])[0], + "card1": rng.randint(1000, 1599), # tail misses the stats table + "card2": None if rng.random() < 0.15 else rng.randint(100, 600), + "card3": None if rng.random() < 0.1 else rng.choice([150, 185]), + "card4": None + if rng.random() < 0.1 + else rng.choice( + [ + "visa", + "Visa ", + "mastercard", + "MasterCard", + "discover", + "american express", + ] + ), + "card5": None if rng.random() < 0.2 else rng.randint(100, 240), + "card6": None + if rng.random() < 0.08 + else rng.choices( + ["debit", "credit", "charge card"], weights=[70, 27, 3] + )[0], + "addr1": None if rng.random() < 0.3 else rng.randint(100, 560), + "addr2": None if rng.random() < 0.3 else rng.choice([87, 60, 96]), + "dist1": None if rng.random() < 0.6 else float(rng.randint(0, 3000)), + "dist2": None if rng.random() < 0.93 else float(rng.randint(0, 8000)), + "p_email_domain": p_dom, + "r_email_domain": r_dom, + "c1": rng.randint(0, 20) + if rng.random() < 0.85 + else rng.randint(20, 400), + "c2": rng.randint(0, 15) + if rng.random() < 0.85 + else rng.randint(15, 300), + "c3": rng.randint(0, 2), + "c4": rng.randint(0, 8), + "c5": rng.randint(0, 30), + "c6": rng.randint(0, 12), + "d1": None if rng.random() < 0.45 else rng.randint(0, 640), + "d2": None if rng.random() < 0.55 else rng.randint(0, 640), + "d3": None if rng.random() < 0.6 else rng.randint(0, 500), + "d4": None if rng.random() < 0.65 else rng.randint(0, 500), + "transaction_dt": base_ts + rng.randint(0, 60 * 86400), + "m1": None if rng.random() < 0.3 else rng.random() < 0.6, + "m2": None if rng.random() < 0.35 else rng.random() < 0.5, + "m3": None if rng.random() < 0.4 else rng.random() < 0.5, + "device_type": None + if rng.random() < 0.25 + else rng.choice(["desktop", "mobile"]), + "v1": None if rng.random() < 0.5 else round(rng.uniform(0.0, 10.0), 4), + "v2": None if rng.random() < 0.5 else round(rng.uniform(0.0, 10.0), 4), + } + ) + return rows + + +SQL = """ +SELECT + txn_id, + transaction_amt AS amt, + card1_stats.card1_amt_mean AS card1_amt_mean, + card1_stats.card1_txn_cnt AS card1_txn_cnt, + card1_stats.card1_fraud_rate AS card1_fraud_rate, + transaction_amt / card1_stats.card1_amt_mean AS amt_to_card1_mean, + transaction_amt - card1_stats.card1_amt_mean AS amt_minus_card1_mean, + addr1_stats.addr1_fraud_rate AS addr1_fraud_rate, + addr1_stats.addr1_txn_cnt AS addr1_txn_cnt, + p_email_stats.p_email_fraud_rate AS p_email_fraud_rate, + p_email_stats.p_email_freq AS p_email_freq, + r_email_stats.r_email_fraud_rate AS r_email_fraud_rate, + pcd_stats.pcd_fraud_rate AS pcd_fraud_rate, + transaction_amt * pcd_stats.pcd_fraud_rate AS amt_x_pcd_rate, + CAST(((transaction_dt % 86400) - (transaction_dt % 3600)) / 3600 AS INTEGER) AS txn_hour, + (CAST((transaction_dt - (transaction_dt % 86400)) / 86400 AS INTEGER) + 4) % 7 AS txn_dow, + CASE WHEN (transaction_dt % 86400) < 21600 THEN 1 ELSE 0 END AS is_night, + CASE WHEN transaction_amt < 20.0 THEN 0 + WHEN transaction_amt < 50.0 THEN 1 + WHEN transaction_amt < 100.0 THEN 2 + WHEN transaction_amt < 300.0 THEN 3 + ELSE 4 END AS amt_bucket, + transaction_amt % 1.0 AS amt_cents, + CASE WHEN (transaction_amt % 1.0) = 0.0 THEN 1 ELSE 0 END AS is_whole_amt, + CASE WHEN d1 IS NULL THEN 1 ELSE 0 END AS d1_null, + CASE WHEN d2 IS NULL THEN 1 ELSE 0 END AS d2_null, + CASE WHEN d3 IS NULL THEN 1 ELSE 0 END AS d3_null, + CASE WHEN d4 IS NULL THEN 1 ELSE 0 END AS d4_null, + coalesce(d1, -1) AS d1_filled, + CASE WHEN product_cd = 'W' THEN 1 ELSE 0 END AS pcd_w, + CASE WHEN product_cd = 'C' THEN 1 ELSE 0 END AS pcd_c, + CASE WHEN product_cd = 'R' THEN 1 ELSE 0 END AS pcd_r, + CASE WHEN product_cd = 'H' THEN 1 ELSE 0 END AS pcd_h, + CASE WHEN product_cd = 'S' THEN 1 ELSE 0 END AS pcd_s, + coalesce(upper(trim(card4)), 'UNK') AS card4_norm, + CASE WHEN card6 = 'debit' THEN 1 ELSE 0 END AS is_debit, + CASE WHEN p_email_domain = r_email_domain THEN 1 + WHEN p_email_domain IS NULL OR r_email_domain IS NULL THEN -1 + ELSE 0 END AS email_match, + c1 + c2 + c3 + c4 + c5 + c6 AS c_sum, + c1 / (c1 + c2 + c3 + c4 + c5 + c6 + 1) AS c1_share, + coalesce(dist1, 0.0) AS dist1_filled, + CASE WHEN m1 IS NULL THEN -1 WHEN m1 THEN 1 ELSE 0 END AS m1_flag, + CASE WHEN addr1 IS NULL THEN 0 ELSE 1 END AS addr_known +FROM __THIS__ +LEFT JOIN card1_stats ON card1 = card1_stats.card1_id +LEFT JOIN addr1_stats ON addr1 = addr1_stats.addr1_id +LEFT JOIN p_email_stats ON p_email_domain = p_email_stats.p_domain +LEFT JOIN r_email_stats ON r_email_domain = r_email_stats.r_domain +LEFT JOIN pcd_stats ON product_cd = pcd_stats.pcd +""" + + +def handcrafted(statics: dict[str, pa.Table]) -> Callable[[dict], dict]: + """What a competent engineer hand-writes for the same features: hydrate the + encoding tables into plain dicts once, then a per-row closure.""" + card1_map = { + r["card1_id"]: (r["card1_amt_mean"], r["card1_txn_cnt"], r["card1_fraud_rate"]) + for r in statics["card1_stats"].to_pylist() + } + addr1_map = { + r["addr1_id"]: (r["addr1_fraud_rate"], r["addr1_txn_cnt"]) + for r in statics["addr1_stats"].to_pylist() + } + p_email_map = { + r["p_domain"]: (r["p_email_fraud_rate"], r["p_email_freq"]) + for r in statics["p_email_stats"].to_pylist() + } + r_email_map = { + r["r_domain"]: r["r_email_fraud_rate"] + for r in statics["r_email_stats"].to_pylist() + } + pcd_map = {r["pcd"]: r["pcd_fraud_rate"] for r in statics["pcd_stats"].to_pylist()} + + def fn(row: dict) -> dict: + amt = row["transaction_amt"] + ts = row["transaction_dt"] + + c1s = card1_map.get(row["card1"]) + card1_amt_mean, card1_txn_cnt, card1_fraud_rate = ( + c1s if c1s else (None, None, None) + ) + a1s = addr1_map.get(row["addr1"]) + addr1_fraud_rate, addr1_txn_cnt = a1s if a1s else (None, None) + pes = p_email_map.get(row["p_email_domain"]) + p_email_fraud_rate, p_email_freq = pes if pes else (None, None) + r_email_fraud_rate = r_email_map.get(row["r_email_domain"]) + pcd_fraud_rate = pcd_map.get(row["product_cd"]) + + if amt < 20.0: + amt_bucket = 0 + elif amt < 50.0: + amt_bucket = 1 + elif amt < 100.0: + amt_bucket = 2 + elif amt < 300.0: + amt_bucket = 3 + else: + amt_bucket = 4 + + amt_cents = math.fmod(amt, 1.0) # SQL float %: fmod, amounts positive + + p_dom, r_dom = row["p_email_domain"], row["r_email_domain"] + if p_dom is not None and r_dom is not None and p_dom == r_dom: + email_match = 1 + elif p_dom is None or r_dom is None: + email_match = -1 + else: + email_match = 0 + + c_sum = row["c1"] + row["c2"] + row["c3"] + row["c4"] + row["c5"] + row["c6"] + card4, m1 = row["card4"], row["m1"] + + return { + "txn_id": row["txn_id"], + "amt": amt, + "card1_amt_mean": card1_amt_mean, + "card1_txn_cnt": card1_txn_cnt, + "card1_fraud_rate": card1_fraud_rate, + "amt_to_card1_mean": amt / card1_amt_mean + if card1_amt_mean is not None + else None, + "amt_minus_card1_mean": amt - card1_amt_mean + if card1_amt_mean is not None + else None, + "addr1_fraud_rate": addr1_fraud_rate, + "addr1_txn_cnt": addr1_txn_cnt, + "p_email_fraud_rate": p_email_fraud_rate, + "p_email_freq": p_email_freq, + "r_email_fraud_rate": r_email_fraud_rate, + "pcd_fraud_rate": pcd_fraud_rate, + "amt_x_pcd_rate": amt * pcd_fraud_rate + if pcd_fraud_rate is not None + else None, + "txn_hour": (ts % 86400) // 3600, + "txn_dow": (ts // 86400 + 4) % 7, + "is_night": 1 if ts % 86400 < 21600 else 0, + "amt_bucket": amt_bucket, + "amt_cents": amt_cents, + "is_whole_amt": 1 if amt_cents == 0.0 else 0, + "d1_null": 1 if row["d1"] is None else 0, + "d2_null": 1 if row["d2"] is None else 0, + "d3_null": 1 if row["d3"] is None else 0, + "d4_null": 1 if row["d4"] is None else 0, + "d1_filled": row["d1"] if row["d1"] is not None else -1, + "pcd_w": 1 if row["product_cd"] == "W" else 0, + "pcd_c": 1 if row["product_cd"] == "C" else 0, + "pcd_r": 1 if row["product_cd"] == "R" else 0, + "pcd_h": 1 if row["product_cd"] == "H" else 0, + "pcd_s": 1 if row["product_cd"] == "S" else 0, + "card4_norm": card4.strip(" ").upper() if card4 is not None else "UNK", + "is_debit": 1 if row["card6"] == "debit" else 0, + "email_match": email_match, + "c_sum": c_sum, + "c1_share": row["c1"] / (c_sum + 1), + "dist1_filled": row["dist1"] if row["dist1"] is not None else 0.0, + "m1_flag": -1 if m1 is None else (1 if m1 else 0), + "addr_known": 0 if row["addr1"] is None else 1, + } + + return fn diff --git a/benchmarks/serving_scenarios/titanic.py b/benchmarks/serving_scenarios/titanic.py new file mode 100644 index 0000000..e5dd957 --- /dev/null +++ b/benchmarks/serving_scenarios/titanic.py @@ -0,0 +1,293 @@ +"""Titanic survival serving scenario for the SQL specializer benchmark. + +Reproduces the canonical Kaggle-Titanic public-kernel feature pipeline as it +looks at SERVE time: scalar expressions over the raw passenger row plus LEFT +JOINs to prepare-time fitted tables (per-pclass fare medians, per-title age +medians and group mean fares, sex-x-pclass target-mean encoding, embarked +target encoding). +""" + +import random +from collections.abc import Callable + +import pyarrow as pa + +NAME = "titanic_survival_features" +KAGGLE = ( + "Kaggle Titanic (survival prediction) — canonical public-kernel tricks: " + "family_size/is_alone, fare_per_person with per-pclass median imputation, " + "cabin deck letter + has_cabin, title-based age imputation with missing " + "flag, age*pclass interaction, age bins, embarked one-hots, sex-x-pclass " + "target-mean encoding and title-GROUP mean fare as fitted join tables." +) + +N_INPUT_COLS = 10 +N_OUTPUT_COLS = 20 + +ROW_SCHEMA = { + "passenger_id": "int", + "pclass": "int", + "sex": "str", + "age": "float?", + "sibsp": "int", + "parch": "int", + "fare": "float?", + "cabin": "str?", + "embarked": "str?", + "title": "str", +} + +# Fitted fallback constants a real pipeline bakes into its generated SQL. +_GLOBAL_MEDIAN_AGE = 29.7 +_GLOBAL_TITLE_FARE = 32.2 +_MODE_PORT_RATE = 0.339 + +SQL = """ +SELECT + passenger_id, + sibsp + parch + 1 AS family_size, + CASE WHEN sibsp + parch = 0 THEN 1 ELSE 0 END AS is_alone, + coalesce(fare, pclass_dim.median_fare) AS fare_filled, + coalesce(fare, pclass_dim.median_fare) / (sibsp + parch + 1) AS fare_per_person, + coalesce(upper(substr(trim(cabin), 1, 1)), 'U') AS deck, + CASE WHEN cabin IS NULL THEN 0 ELSE 1 END AS has_cabin, + CASE WHEN age IS NULL THEN 1 ELSE 0 END AS age_missing, + coalesce(age, title_stats.median_age, 29.7) AS age_filled, + coalesce(age, title_stats.median_age, 29.7) * pclass AS age_class, + CASE + WHEN coalesce(age, title_stats.median_age, 29.7) < 13.0 THEN 'child' + WHEN coalesce(age, title_stats.median_age, 29.7) < 20.0 THEN 'teen' + WHEN coalesce(age, title_stats.median_age, 29.7) < 60.0 THEN 'adult' + ELSE 'senior' + END AS age_bin, + CASE WHEN coalesce(embarked, 'S') = 'C' THEN 1 ELSE 0 END AS embarked_c, + CASE WHEN coalesce(embarked, 'S') = 'Q' THEN 1 ELSE 0 END AS embarked_q, + CASE WHEN coalesce(embarked, 'S') = 'S' THEN 1 ELSE 0 END AS embarked_s, + CASE WHEN sex = 'female' THEN 1 ELSE 0 END AS sex_female, + sex || '-' || CAST(pclass AS VARCHAR) AS sex_pclass, + sex_pclass_enc.survival_rate AS sex_pclass_rate, + coalesce(title_stats.group_mean_fare, 32.2) AS title_fare_mean, + pclass_dim.pclass_survival_rate AS pclass_rate, + coalesce(embarked_enc.port_survival_rate, 0.339) AS port_rate +FROM __THIS__ +LEFT JOIN pclass_dim ON pclass = pclass_dim.pc +LEFT JOIN title_stats ON title = title_stats.t_title +LEFT JOIN sex_pclass_enc + ON sex = sex_pclass_enc.sp_sex AND pclass = sex_pclass_enc.sp_pclass +LEFT JOIN embarked_enc ON embarked = embarked_enc.port +""" + +# Title -> (group median age, denormalized title-GROUP mean fare). The +# grouping {Mr, Mrs, Miss, Master, Rare} happened at fit time; the fitted +# table is keyed on the raw title. Capt/Jonkheer/Dona are deliberately NOT +# here: unseen-at-fit titles cause serve-time LEFT JOIN misses. +_TITLES = { + "Mr": (30.0, 24.4), + "Mrs": (35.0, 45.0), + "Miss": (21.0, 43.8), + "Master": (3.5, 37.0), + "Dr": (46.5, 40.9), + "Rev": (46.5, 40.9), + "Col": (46.5, 40.9), + "Major": (46.5, 40.9), + "Mlle": (21.0, 43.8), + "Ms": (21.0, 43.8), + "Mme": (35.0, 45.0), + "Lady": (46.5, 40.9), + "Sir": (46.5, 40.9), + "Countess": (46.5, 40.9), +} + + +def make_statics(seed: int) -> dict[str, pa.Table]: + r = random.Random(seed) + + def jit(x: float) -> float: + return round(x + r.uniform(-0.015, 0.015), 6) + + pclass_dim = pa.table( + { + "pc": [1, 2, 3], + "median_fare": [round(jit(60.2875), 4), 14.25, 8.05], + "pclass_survival_rate": [jit(0.6296), jit(0.4728), jit(0.2424)], + } + ) + titles = sorted(_TITLES) + title_stats = pa.table( + { + "t_title": titles, + "median_age": [_TITLES[t][0] for t in titles], + "group_mean_fare": [jit(_TITLES[t][1]) for t in titles], + } + ) + sex_pclass_enc = pa.table( + { + "sp_sex": ["female", "female", "female", "male", "male", "male"], + "sp_pclass": [1, 2, 3, 1, 2, 3], + "survival_rate": [ + jit(0.9681), + jit(0.9211), + jit(0.5000), + jit(0.3689), + jit(0.1574), + jit(0.1354), + ], + } + ) + embarked_enc = pa.table( + { + "port": ["S", "C", "Q"], + "port_survival_rate": [jit(0.3370), jit(0.5539), jit(0.3896)], + } + ) + return { + "pclass_dim": pclass_dim, + "title_stats": title_stats, + "sex_pclass_enc": sex_pclass_enc, + "embarked_enc": embarked_enc, + } + + +def make_rows(seed: int, n: int) -> list[dict]: + r = random.Random(seed) + male_titles = ["Mr"] * 90 + ["Master"] * 6 + ["Dr", "Rev", "Capt", "Jonkheer"] + female_titles = ( + ["Miss"] * 44 + + ["Mrs"] * 44 + + ["Mlle", "Ms", "Mme", "Lady", "Countess", "Dona"] * 2 + ) + decks = {1: "ABCDE", 2: "DEF", 3: "EFG"} + fare_base = {1: 84.15, 2: 20.66, 3: 13.68} + rows = [] + for i in range(n): + pclass = r.choices([1, 2, 3], weights=[24, 21, 55])[0] + sex = "male" if r.random() < 0.65 else "female" + title = r.choice(male_titles if sex == "male" else female_titles) + if r.random() < 0.20: + age = None + elif title == "Master": + age = round(r.uniform(0.5, 12.0) * 2) / 2 + else: + age = min(80.0, max(14.0, round(r.gauss(30.0, 13.0) * 2) / 2)) + sibsp = r.choices([0, 1, 2, 3, 4, 8], weights=[68, 23, 4, 3, 1, 1])[0] + parch = r.choices([0, 1, 2, 5], weights=[76, 13, 9, 2])[0] + if r.random() < 0.01: + fare = None + elif r.random() < 0.015: + fare = 0.0 + else: + fare = round(fare_base[pclass] * (0.35 + r.random() * 2.2), 4) + if r.random() < 0.77: + cabin = None + else: + letter = r.choice(decks[pclass]) + if r.random() < 0.15: + letter = letter.lower() + cabin = f"{letter}{r.randint(1, 130)}" + if r.random() < 0.10: + cabin = f" {cabin} " + embarked = ( + None + if r.random() < 0.02 + else r.choices(["S", "C", "Q"], weights=[72, 19, 9])[0] + ) + rows.append( + { + "passenger_id": 1000 + i, + "pclass": pclass, + "sex": sex, + "age": age, + "sibsp": sibsp, + "parch": parch, + "fare": fare, + "cabin": cabin, + "embarked": embarked, + "title": title, + } + ) + return rows + + +def handcrafted(statics: dict[str, pa.Table]) -> Callable[[dict], dict]: + def cols(t: pa.Table) -> dict[str, list]: + return {name: t.column(name).to_pylist() for name in t.column_names} + + p = cols(statics["pclass_dim"]) + pclass_lut = { + k: (mf, sr) + for k, mf, sr in zip( + p["pc"], p["median_fare"], p["pclass_survival_rate"], strict=True + ) + } + t = cols(statics["title_stats"]) + title_lut = { + k: (ma, gf) + for k, ma, gf in zip( + t["t_title"], t["median_age"], t["group_mean_fare"], strict=True + ) + } + s = cols(statics["sex_pclass_enc"]) + sp_lut = { + (sx, pc): sr + for sx, pc, sr in zip( + s["sp_sex"], s["sp_pclass"], s["survival_rate"], strict=True + ) + } + e = cols(statics["embarked_enc"]) + port_lut = dict(zip(e["port"], e["port_survival_rate"], strict=True)) + + def fn(row: dict) -> dict: + sibsp, parch, pclass = row["sibsp"], row["parch"], row["pclass"] + family_size = sibsp + parch + 1 + median_fare, pclass_rate = pclass_lut.get(pclass, (None, None)) + fare = row["fare"] + fare_filled = fare if fare is not None else median_fare + fare_per_person = None if fare_filled is None else fare_filled / family_size + cabin = row["cabin"] + deck = "U" if cabin is None else cabin.strip()[:1].upper() + median_age, group_fare = title_lut.get(row["title"], (None, None)) + age = row["age"] + if age is not None: + age_filled = age + elif median_age is not None: + age_filled = median_age + else: + age_filled = _GLOBAL_MEDIAN_AGE + if age_filled < 13.0: + age_bin = "child" + elif age_filled < 20.0: + age_bin = "teen" + elif age_filled < 60.0: + age_bin = "adult" + else: + age_bin = "senior" + embarked = row["embarked"] + emb = embarked if embarked is not None else "S" + port_rate = port_lut.get(embarked) if embarked is not None else None + sex = row["sex"] + return { + "passenger_id": row["passenger_id"], + "family_size": family_size, + "is_alone": 1 if sibsp + parch == 0 else 0, + "fare_filled": fare_filled, + "fare_per_person": fare_per_person, + "deck": deck, + "has_cabin": 0 if cabin is None else 1, + "age_missing": 1 if age is None else 0, + "age_filled": age_filled, + "age_class": age_filled * pclass, + "age_bin": age_bin, + "embarked_c": 1 if emb == "C" else 0, + "embarked_q": 1 if emb == "Q" else 0, + "embarked_s": 1 if emb == "S" else 0, + "sex_female": 1 if sex == "female" else 0, + "sex_pclass": f"{sex}-{pclass}", + "sex_pclass_rate": sp_lut.get((sex, pclass)), + "title_fare_mean": group_fare + if group_fare is not None + else _GLOBAL_TITLE_FARE, + "pclass_rate": pclass_rate, + "port_rate": port_rate if port_rate is not None else _MODE_PORT_RATE, + } + + return fn diff --git a/pyproject.toml b/pyproject.toml index 26894ce..55dd67c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,9 @@ known-third-party = ["datafusion", "duckdb"] "_state.py" = ["S608"] # Doc generator: emits long markdown lines verbatim; wrapping would corrupt output. "scripts/gen_datafusion_catalogue.py" = ["E501"] +# Bench fixtures: seeded pseudo-random data generators (S311 is about crypto) +# and dense SQL/data-gen literals where wrapping hurts readability. +"benchmarks/serving_scenarios/*.py" = ["S311", "E501"] [tool.pytest.ini_options] python_files = ["*_test.py", "test_*.py"] diff --git a/tests/test_serving_scenarios.py b/tests/test_serving_scenarios.py new file mode 100644 index 0000000..aa011cf --- /dev/null +++ b/tests/test_serving_scenarios.py @@ -0,0 +1,33 @@ +"""Standing parity gate for the serving-bench scenarios. + +Each scenario (benchmarks/serving_scenarios/) is a realistic wide-table +feature-engineering inference path. The specializer's output must equal +DuckDB itself AND the handcrafted Python twin, exactly, on 300 seeded rows. +""" + +from __future__ import annotations + +import pytest + +from benchmarks import serving_scenarios as sc + + +@pytest.mark.parametrize("name", sc.NAMES) +def test_scenario_three_way_parity(name): + mod = sc.load(name) + problems = sc.verify_parity(mod, n=300) + assert not problems, "\n".join(problems) + + +@pytest.mark.parametrize("name", sc.NAMES) +def test_scenario_is_wide_and_deterministic(name): + mod = sc.load(name) + assert mod.N_INPUT_COLS == len(mod.ROW_SCHEMA) + rows_a = mod.make_rows(sc.SEED, 50) + rows_b = mod.make_rows(sc.SEED, 50) + assert rows_a == rows_b, "make_rows must be deterministic" + statics_a = mod.make_statics(sc.SEED) + statics_b = mod.make_statics(sc.SEED) + assert {k: t.to_pylist() for k, t in statics_a.items()} == { + k: t.to_pylist() for k, t in statics_b.items() + }, "make_statics must be deterministic" From 82ab6ca08b10bb476207d0bcd344a7b4252f5263 Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 11:13:13 +0200 Subject: [PATCH 10/11] =?UTF-8?q?feat(bench):=20all=20four=20serving=20sce?= =?UTF-8?q?narios=20=E2=80=94=20house=5Fprices,=20store=5Fsales=20+=20harn?= =?UTF-8?q?ess-aware=20titanic/fraud=20rewrites?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four famous-problem inference paths, all passing exact three-way parity (specializer == DuckDB multiset == handcrafted twin) on 300 seeded rows: titanic 10->20 (strings/CASE), house_prices ~40->30+ (wide arithmetic, quality ordinals, neighborhood encoding), fraud_txn 32->38 (five probe tables, NULL flags, time arithmetic), store_sales ~20+dim (store-dim join, competition/promo2 arithmetic, seasonal factors). Each module records which real winning-solution features the v0 surface could NOT express — the workload-driven support ladder (log/log1p in all four; instr/position; true floor; IN-list; fractional pow; trig). Co-Authored-By: Claude Fable 5 --- benchmarks/bench_serving.py | 5 + benchmarks/serving_scenarios/fraud_txn.py | 65 ++- benchmarks/serving_scenarios/house_prices.py | 470 ++++++++++++++++++ benchmarks/serving_scenarios/store_sales.py | 496 +++++++++++++++++++ benchmarks/serving_scenarios/titanic.py | 474 ++++++++++-------- pyproject.toml | 7 +- 6 files changed, 1305 insertions(+), 212 deletions(-) create mode 100644 benchmarks/serving_scenarios/house_prices.py create mode 100644 benchmarks/serving_scenarios/store_sales.py diff --git a/benchmarks/bench_serving.py b/benchmarks/bench_serving.py index a93ecde..180f68e 100644 --- a/benchmarks/bench_serving.py +++ b/benchmarks/bench_serving.py @@ -58,10 +58,15 @@ def bench_call(f, n): iters = max(30, min(2000, 2_000_000 // max(n, 1))) f() # warm out = [] + # 3s budget per cell (min 30 samples) — keeps the ms-scale engines + # (duckdb per call) from dominating wall time. + deadline = time.perf_counter_ns() + 3_000_000_000 for _ in range(iters): t0 = time.perf_counter_ns() f() out.append(time.perf_counter_ns() - t0) + if len(out) >= 30 and time.perf_counter_ns() > deadline: + break p50, p99 = quantiles(out) return {"p50_ns": p50, "p99_ns": p99} diff --git a/benchmarks/serving_scenarios/fraud_txn.py b/benchmarks/serving_scenarios/fraud_txn.py index 676c797..ca6bbac 100644 --- a/benchmarks/serving_scenarios/fraud_txn.py +++ b/benchmarks/serving_scenarios/fraud_txn.py @@ -16,13 +16,14 @@ NAME = "fraud_txn" KAGGLE = ( "IEEE-CIS Fraud Detection (Kaggle 2019, Deotte/Yakovlev winning-solution " - "serving shape): card1/addr1/email-domain frequency+target encodings as " - "join tables, amt ratio/delta vs card1 mean, cents feature, D-column NULL " - "flags, hour/day-of-week from unix ts, amount buckets, ProductCD one-hots." + "serving shape): composite uid (card1 x addr1) + card1/addr1/email-domain " + "frequency+target encodings as join tables, amt ratio/delta vs card1 and " + "uid means, cents feature, D-column NULL flags, hour/day-of-week from " + "unix ts, amount buckets, ProductCD one-hots." ) N_INPUT_COLS = 32 -N_OUTPUT_COLS = 38 +N_OUTPUT_COLS = 41 ROW_SCHEMA: dict[str, str] = { "txn_id": "int", @@ -85,6 +86,27 @@ _SERVE_DOMAINS = _DOMAINS + ["qq.com", "rocketmail.com", "protonmail.ch"] _PRODUCTS = ["W", "C", "R", "H", "S"] +# The fitted uid (card1 x addr1) population is a property of the training +# data, not of the value-fit seed: fixed here so make_rows can send +# repeat-customer traffic that hits the uid encoding at a realistic rate. +_UID_SEED = 990721 +_uid_cache: list[tuple[int, int]] | None = None + + +def _uid_pairs() -> list[tuple[int, int]]: + global _uid_cache + if _uid_cache is None: + rng = random.Random(_UID_SEED) + pairs: list[tuple[int, int]] = [] + seen: set[tuple[int, int]] = set() + while len(pairs) < 2500: + p = (rng.randint(1000, 1599), rng.randint(100, 559)) + if p not in seen: + seen.add(p) + pairs.append(p) + _uid_cache = pairs + return _uid_cache + def make_statics(seed: int) -> dict[str, pa.Table]: rng = random.Random(seed * 7919 + 1) @@ -133,20 +155,37 @@ def make_statics(seed: int) -> dict[str, pa.Table]: } ) + uid_pairs = _uid_pairs() + uid_stats = pa.table( + { + "uid_card1": [c for c, _ in uid_pairs], + "uid_addr1": [a for _, a in uid_pairs], + "uid_amt_mean": [round(rng.uniform(8.0, 420.0), 2) for _ in uid_pairs], + "uid_txn_cnt": [rng.randint(1, 800) for _ in uid_pairs], + } + ) + return { "card1_stats": card1_stats, "addr1_stats": addr1_stats, "p_email_stats": p_email_stats, "r_email_stats": r_email_stats, "pcd_stats": pcd_stats, + "uid_stats": uid_stats, } def make_rows(seed: int, n: int) -> list[dict]: rng = random.Random(seed * 104729 + 3) + uid_pairs = _uid_pairs() base_ts = 1_700_000_000 rows: list[dict[str, Any]] = [] for i in range(n): + if rng.random() < 0.55: # repeat customer: hits a fitted uid pair + card1, addr1 = uid_pairs[rng.randrange(len(uid_pairs))] + else: + card1 = rng.randint(1000, 1599) # tail misses the stats table + addr1 = None if rng.random() < 0.3 else rng.randint(100, 560) if rng.random() < 0.2: # whole-dollar transactions (the cents trick) amt = float(rng.choice([20, 25, 30, 50, 75, 100, 150, 200, 300, 500])) else: @@ -161,7 +200,7 @@ def make_rows(seed: int, n: int) -> list[dict]: "txn_id": 3_000_000 + i, "transaction_amt": amt, "product_cd": rng.choices(_PRODUCTS, weights=[65, 10, 8, 9, 8])[0], - "card1": rng.randint(1000, 1599), # tail misses the stats table + "card1": card1, "card2": None if rng.random() < 0.15 else rng.randint(100, 600), "card3": None if rng.random() < 0.1 else rng.choice([150, 185]), "card4": None @@ -182,7 +221,7 @@ def make_rows(seed: int, n: int) -> list[dict]: else rng.choices( ["debit", "credit", "charge card"], weights=[70, 27, 3] )[0], - "addr1": None if rng.random() < 0.3 else rng.randint(100, 560), + "addr1": addr1, "addr2": None if rng.random() < 0.3 else rng.choice([87, 60, 96]), "dist1": None if rng.random() < 0.6 else float(rng.randint(0, 3000)), "dist2": None if rng.random() < 0.93 else float(rng.randint(0, 8000)), @@ -227,6 +266,9 @@ def make_rows(seed: int, n: int) -> list[dict]: transaction_amt - card1_stats.card1_amt_mean AS amt_minus_card1_mean, addr1_stats.addr1_fraud_rate AS addr1_fraud_rate, addr1_stats.addr1_txn_cnt AS addr1_txn_cnt, + uid_stats.uid_amt_mean AS uid_amt_mean, + uid_stats.uid_txn_cnt AS uid_txn_cnt, + transaction_amt / uid_stats.uid_amt_mean AS amt_to_uid_mean, p_email_stats.p_email_fraud_rate AS p_email_fraud_rate, p_email_stats.p_email_freq AS p_email_freq, r_email_stats.r_email_fraud_rate AS r_email_fraud_rate, @@ -265,6 +307,7 @@ def make_rows(seed: int, n: int) -> list[dict]: FROM __THIS__ LEFT JOIN card1_stats ON card1 = card1_stats.card1_id LEFT JOIN addr1_stats ON addr1 = addr1_stats.addr1_id +LEFT JOIN uid_stats ON card1 = uid_stats.uid_card1 AND addr1 = uid_stats.uid_addr1 LEFT JOIN p_email_stats ON p_email_domain = p_email_stats.p_domain LEFT JOIN r_email_stats ON r_email_domain = r_email_stats.r_domain LEFT JOIN pcd_stats ON product_cd = pcd_stats.pcd @@ -291,6 +334,10 @@ def handcrafted(statics: dict[str, pa.Table]) -> Callable[[dict], dict]: for r in statics["r_email_stats"].to_pylist() } pcd_map = {r["pcd"]: r["pcd_fraud_rate"] for r in statics["pcd_stats"].to_pylist()} + uid_map = { + (r["uid_card1"], r["uid_addr1"]): (r["uid_amt_mean"], r["uid_txn_cnt"]) + for r in statics["uid_stats"].to_pylist() + } def fn(row: dict) -> dict: amt = row["transaction_amt"] @@ -302,6 +349,9 @@ def fn(row: dict) -> dict: ) a1s = addr1_map.get(row["addr1"]) addr1_fraud_rate, addr1_txn_cnt = a1s if a1s else (None, None) + # NULL addr1 can never equal a fitted key, same as the SQL join + us = uid_map.get((row["card1"], row["addr1"])) + uid_amt_mean, uid_txn_cnt = us if us else (None, None) pes = p_email_map.get(row["p_email_domain"]) p_email_fraud_rate, p_email_freq = pes if pes else (None, None) r_email_fraud_rate = r_email_map.get(row["r_email_domain"]) @@ -345,6 +395,9 @@ def fn(row: dict) -> dict: else None, "addr1_fraud_rate": addr1_fraud_rate, "addr1_txn_cnt": addr1_txn_cnt, + "uid_amt_mean": uid_amt_mean, + "uid_txn_cnt": uid_txn_cnt, + "amt_to_uid_mean": amt / uid_amt_mean if uid_amt_mean is not None else None, "p_email_fraud_rate": p_email_fraud_rate, "p_email_freq": p_email_freq, "r_email_fraud_rate": r_email_fraud_rate, diff --git a/benchmarks/serving_scenarios/house_prices.py b/benchmarks/serving_scenarios/house_prices.py new file mode 100644 index 0000000..faaeb6b --- /dev/null +++ b/benchmarks/serving_scenarios/house_prices.py @@ -0,0 +1,470 @@ +"""house_prices — Ames wide-arithmetic serving scenario for the SQL specializer. + +The canonical Kaggle House Prices feature set, expressed as the serve-time +query a fitted pipeline reduces to: scalar expressions over the row plus +LEFT JOINs to prepare-time encoding tables. This is the wide-arith stress +case: 43 input columns, 42 output features, almost all of them arithmetic +and CASE over the row. + +Realistic serving gotchas baked into the row distributions: +- MSSubClass 150 and the GrnHill/Landmrk neighborhoods exist only outside + the Kaggle train set -> LEFT JOIN misses at serve time (COALESCE to the + train-global mean price, the standard target-encoding fallback). +- The test set has NA garage/basement numerics the train set never had + (famous rows 2121/2189/2577) -> None on nullable columns, propagating + through arithmetic exactly like SQL NULL. +- LotFrontage is missing ~17% of the time -> imputed from the fitted + per-neighborhood median via the join table (the classic trick). +""" + +import random +from collections.abc import Callable + +import pyarrow as pa + +NAME = "house_prices" +KAGGLE = ( + "House Prices: Advanced Regression Techniques (Ames) — the canonical " + "public-kernel feature set: TotalSF/TotalBath sums, age features, " + "Ex..Po ordinal quality maps, porch total, has_* flags, nullif-guarded " + "lot ratios, qual*cond crosses, neighborhood-median LotFrontage " + "imputation, and Neighborhood/MSSubClass target+frequency encodings " + "served as join tables." +) + +N_INPUT_COLS = 43 +N_OUTPUT_COLS = 42 + +ROW_SCHEMA = { + "id": "int", + "ms_sub_class": "int", + "lot_frontage": "float?", + "lot_area": "int", + "neighborhood": "str", + "overall_qual": "int", + "overall_cond": "int", + "year_built": "int", + "year_remod_add": "int", + "exter_qual": "str", + "exter_cond": "str", + "mas_vnr_area": "float?", + "bsmt_qual": "str?", + "bsmt_fin_sf1": "float?", + "bsmt_unf_sf": "float?", + "total_bsmt_sf": "float?", + "heating_qc": "str", + "central_air": "str", + "first_flr_sf": "int", + "second_flr_sf": "int", + "gr_liv_area": "int", + "bsmt_full_bath": "int?", + "bsmt_half_bath": "int?", + "full_bath": "int", + "half_bath": "int", + "bedroom_abv_gr": "int", + "kitchen_abv_gr": "int", + "kitchen_qual": "str?", + "tot_rms_abv_grd": "int", + "fireplaces": "int", + "fireplace_qu": "str?", + "garage_yr_blt": "float?", + "garage_cars": "int?", + "garage_area": "float?", + "garage_qual": "str?", + "wood_deck_sf": "int", + "open_porch_sf": "int", + "enclosed_porch": "int", + "three_ssn_porch": "int", + "screen_porch": "int", + "pool_area": "int", + "mo_sold": "int", + "yr_sold": "int", +} + +# (name, train mean price, train count, median lot frontage) — the 25 Kaggle +# train neighborhoods. GrnHill/Landmrk exist in full Ames but not in train: +# the unseen-category LEFT JOIN miss at serve time. +_NBHDS = [ + ("NAmes", 145847, 225, 73.0), + ("CollgCr", 197966, 150, 70.0), + ("OldTown", 128225, 113, 60.0), + ("Edwards", 128220, 100, 66.0), + ("Somerst", 225380, 86, 73.5), + ("Gilbert", 192854, 79, 64.0), + ("NridgHt", 316271, 77, 92.0), + ("Sawyer", 136793, 74, 71.0), + ("NWAmes", 189050, 73, 80.0), + ("SawyerW", 186556, 59, 66.5), + ("BrkSide", 124834, 58, 52.0), + ("Crawfor", 210625, 51, 70.0), + ("Mitchel", 156270, 49, 73.0), + ("NoRidge", 335295, 41, 91.0), + ("Timber", 242247, 38, 85.0), + ("IDOTRR", 100124, 37, 60.0), + ("ClearCr", 212565, 28, 80.0), + ("StoneBr", 310499, 25, 61.5), + ("SWISU", 142591, 25, 60.0), + ("Blmngtn", 194871, 17, 43.0), + ("MeadowV", 98576, 17, 21.0), + ("BrDale", 104494, 16, 21.0), + ("Veenker", 238772, 11, 68.0), + ("NPkVill", 142694, 9, 24.0), + ("Blueste", 137500, 2, 24.0), +] +_UNSEEN_NBHDS = ["GrnHill", "Landmrk"] + +# (MSSubClass, train mean price, train count). 150 is the famous class that +# appears only in the test set — the unseen join miss. +_SUBCLASSES = [ + (20, 185224, 536), + (60, 240403, 299), + (50, 143302, 144), + (120, 200779, 87), + (30, 95829, 69), + (160, 138647, 63), + (70, 166772, 60), + (80, 169736, 58), + (90, 133541, 52), + (190, 129613, 30), + (85, 147810, 20), + (75, 192437, 16), + (45, 108591, 12), + (180, 102300, 10), + (40, 156125, 4), +] +_UNSEEN_SUBCLASS = 150 + +_QUAL_LEVELS = ["Ex", "Gd", "TA", "Fa", "Po"] + + +def make_statics(seed: int) -> dict[str, pa.Table]: + rng = random.Random(seed) + names, means, freqs, medfr = [], [], [], [] + for name, base_mean, count, med in _NBHDS: + names.append(name) + means.append(round(base_mean * rng.uniform(0.97, 1.03), 2)) + freqs.append(round(count / 1460.0, 6)) + medfr.append(med) + nbhd = pa.table( + { + "nbhd": pa.array(names, type=pa.string()), + "mean_price": pa.array(means, type=pa.float64()), + "freq": pa.array(freqs, type=pa.float64()), + "median_frontage": pa.array(medfr, type=pa.float64()), + } + ) + classes, sc_means, sc_freqs = [], [], [] + for cls, base_mean, count in _SUBCLASSES: + classes.append(cls) + sc_means.append(round(base_mean * rng.uniform(0.97, 1.03), 2)) + sc_freqs.append(round(count / 1460.0, 6)) + sub = pa.table( + { + "sub_class": pa.array(classes, type=pa.int64()), + "mean_price": pa.array(sc_means, type=pa.float64()), + "freq": pa.array(sc_freqs, type=pa.float64()), + } + ) + return {"nbhd_price_enc": nbhd, "subclass_enc": sub} + + +def make_rows(seed: int, n: int) -> list[dict]: + rng = random.Random(seed) + nb_names = [t[0] for t in _NBHDS] + nb_wts = [t[2] for t in _NBHDS] + sc_names = [t[0] for t in _SUBCLASSES] + sc_wts = [t[2] for t in _SUBCLASSES] + rows = [] + for i in range(n): + yr_sold = rng.randint(2006, 2010) + year_built = yr_sold if rng.random() < 0.06 else rng.randint(1872, yr_sold) + if rng.random() < 0.45 and year_built < yr_sold: + year_remod_add = rng.randint(max(1950, year_built), yr_sold) + else: + year_remod_add = max(1950, year_built) + + first_flr_sf = rng.randint(334, 2600) + second_flr_sf = 0 if rng.random() < 0.55 else rng.randint(300, 1200) + gr_liv_area = first_flr_sf + second_flr_sf + + u = rng.random() + if u < 0.01: # the test-set row with every basement field missing + bsmt_qual = bsmt_fin_sf1 = bsmt_unf_sf = total_bsmt_sf = None + bsmt_full_bath = bsmt_half_bath = None + elif u < 0.05: # no basement + bsmt_qual = None + bsmt_fin_sf1 = bsmt_unf_sf = total_bsmt_sf = 0.0 + bsmt_full_bath = bsmt_half_bath = 0 + else: + total_bsmt_sf = float(rng.randint(300, first_flr_sf + 200)) + bsmt_fin_sf1 = float(rng.randint(0, int(total_bsmt_sf))) + bsmt_unf_sf = total_bsmt_sf - bsmt_fin_sf1 + bsmt_qual = rng.choices(_QUAL_LEVELS, weights=[8, 40, 44, 6, 2])[0] + bsmt_full_bath = rng.choices([0, 1, 2], weights=[58, 39, 3])[0] + bsmt_half_bath = rng.choices([0, 1], weights=[94, 6])[0] + + u = rng.random() + if u < 0.055: # no garage + garage_yr_blt = garage_qual = None + garage_cars, garage_area = 0, 0.0 + elif u < 0.065: # the famous test row 2577: garage fields missing + garage_yr_blt = garage_qual = garage_cars = garage_area = None + else: + garage_yr_blt = float(rng.randint(year_built, yr_sold)) + garage_cars = rng.choices([1, 2, 3, 4], weights=[25, 55, 18, 2])[0] + garage_area = float(rng.randint(200, 350) * garage_cars) + garage_qual = rng.choices(_QUAL_LEVELS, weights=[1, 2, 90, 5, 2])[0] + + fireplaces = rng.choices([0, 1, 2, 3], weights=[47, 43, 9, 1])[0] + bedroom_abv_gr = rng.choices( + [0, 1, 2, 3, 4, 5, 6], weights=[1, 4, 24, 50, 17, 3, 1] + )[0] + kitchen_abv_gr = rng.choices([1, 2], weights=[95, 5])[0] + full_bath = rng.choices([1, 2, 3], weights=[45, 50, 5])[0] + half_bath = rng.choices([0, 1, 2], weights=[62, 36, 2])[0] + + rows.append( + { + "id": 1461 + i, # Kaggle test-set ids start at 1461 + "ms_sub_class": _UNSEEN_SUBCLASS + if rng.random() < 0.01 + else rng.choices(sc_names, weights=sc_wts)[0], + "lot_frontage": None + if rng.random() < 0.17 + else float(rng.randint(21, 150)), + "lot_area": min(215245, max(1300, int(rng.lognormvariate(9.2, 0.45)))), + "neighborhood": rng.choice(_UNSEEN_NBHDS) + if rng.random() < 0.02 + else rng.choices(nb_names, weights=nb_wts)[0], + "overall_qual": min(10, max(1, round(rng.gauss(6.1, 1.4)))), + "overall_cond": min(10, max(1, round(rng.gauss(5.6, 1.1)))), + "year_built": year_built, + "year_remod_add": year_remod_add, + "exter_qual": rng.choices(_QUAL_LEVELS, weights=[4, 33, 60, 2, 1])[0], + "exter_cond": rng.choices(_QUAL_LEVELS, weights=[1, 10, 86, 2, 1])[0], + "mas_vnr_area": None + if rng.random() < 0.005 + else (0.0 if rng.random() < 0.6 else float(rng.randint(50, 1600))), + "bsmt_qual": bsmt_qual, + "bsmt_fin_sf1": bsmt_fin_sf1, + "bsmt_unf_sf": bsmt_unf_sf, + "total_bsmt_sf": total_bsmt_sf, + "heating_qc": rng.choices(_QUAL_LEVELS, weights=[51, 16, 29, 3, 1])[0], + "central_air": "Y" if rng.random() < 0.93 else "N", + "first_flr_sf": first_flr_sf, + "second_flr_sf": second_flr_sf, + "gr_liv_area": gr_liv_area, + "bsmt_full_bath": bsmt_full_bath, + "bsmt_half_bath": bsmt_half_bath, + "full_bath": full_bath, + "half_bath": half_bath, + "bedroom_abv_gr": bedroom_abv_gr, + "kitchen_abv_gr": kitchen_abv_gr, + "kitchen_qual": None + if rng.random() < 0.01 + else rng.choices(_QUAL_LEVELS, weights=[7, 40, 50, 2, 1])[0], + "tot_rms_abv_grd": bedroom_abv_gr + kitchen_abv_gr + rng.randint(2, 5), + "fireplaces": fireplaces, + "fireplace_qu": None + if fireplaces == 0 + else rng.choices(_QUAL_LEVELS, weights=[3, 49, 41, 4, 3])[0], + "garage_yr_blt": garage_yr_blt, + "garage_cars": garage_cars, + "garage_area": garage_area, + "garage_qual": garage_qual, + "wood_deck_sf": 0 if rng.random() < 0.48 else rng.randint(20, 800), + "open_porch_sf": 0 if rng.random() < 0.45 else rng.randint(10, 500), + "enclosed_porch": 0 if rng.random() < 0.86 else rng.randint(20, 400), + "three_ssn_porch": 0 if rng.random() < 0.98 else rng.randint(100, 300), + "screen_porch": 0 if rng.random() < 0.92 else rng.randint(80, 400), + "pool_area": 0 if rng.random() < 0.995 else rng.randint(400, 800), + "mo_sold": rng.choices( + list(range(1, 13)), + weights=[4, 4, 7, 9, 14, 17, 15, 8, 6, 6, 6, 4], + )[0], + "yr_sold": yr_sold, + } + ) + return rows + + +def _ord_case(col: str) -> str: + """Ex..Po -> 5..1, anything else (incl. NULL: no branch matches) -> 0.""" + branches = " ".join( + f"WHEN {col} = '{lvl}' THEN {5 - i}" for i, lvl in enumerate(_QUAL_LEVELS) + ) + return f"CASE {branches} ELSE 0 END" + + +SQL = f"""SELECT + id, + coalesce(total_bsmt_sf, 0.0) + first_flr_sf + second_flr_sf AS total_sf, + full_bath + CAST(0.5 AS DOUBLE) * half_bath + coalesce(bsmt_full_bath, 0) + + CAST(0.5 AS DOUBLE) * coalesce(bsmt_half_bath, 0) AS total_bath, + yr_sold - year_built AS house_age, + yr_sold - year_remod_add AS remod_age, + yr_sold - garage_yr_blt AS garage_age, + CASE WHEN yr_sold = year_built THEN 1 ELSE 0 END AS is_new, + CASE WHEN year_remod_add > year_built THEN 1 ELSE 0 END AS is_remodeled, + {_ord_case("exter_qual")} AS exter_qual_ord, + {_ord_case("exter_cond")} AS exter_cond_ord, + CASE WHEN kitchen_qual IS NULL THEN 3 WHEN kitchen_qual = 'Ex' THEN 5 + WHEN kitchen_qual = 'Gd' THEN 4 WHEN kitchen_qual = 'TA' THEN 3 + WHEN kitchen_qual = 'Fa' THEN 2 WHEN kitchen_qual = 'Po' THEN 1 + ELSE 0 END AS kitchen_qual_ord, + {_ord_case("heating_qc")} AS heating_qc_ord, + {_ord_case("bsmt_qual")} AS bsmt_qual_ord, + {_ord_case("fireplace_qu")} AS fireplace_qu_ord, + {_ord_case("garage_qual")} AS garage_qual_ord, + wood_deck_sf + open_porch_sf + enclosed_porch + three_ssn_porch + + screen_porch AS porch_total, + CASE WHEN pool_area > 0 THEN 1 ELSE 0 END AS has_pool, + CASE WHEN coalesce(garage_area, 0.0) > 0.0 THEN 1 ELSE 0 END AS has_garage, + CASE WHEN coalesce(total_bsmt_sf, 0.0) > 0.0 THEN 1 ELSE 0 END AS has_bsmt, + CASE WHEN fireplaces > 0 THEN 1 ELSE 0 END AS has_fireplace, + CASE WHEN second_flr_sf > 0 THEN 1 ELSE 0 END AS has_2nd_floor, + CASE WHEN coalesce(mas_vnr_area, 0.0) > 0.0 THEN 1 ELSE 0 END AS has_mas_vnr, + CASE WHEN central_air = 'Y' THEN 1 ELSE 0 END AS central_air_flag, + CASE WHEN kitchen_abv_gr > 1 THEN 1 ELSE 0 END AS multi_kitchen, + coalesce(lot_frontage, nbhd_price_enc.median_frontage, 69.0) AS lot_frontage_filled, + lot_frontage / nullif(lot_area, 0) AS lot_frontage_ratio, + CAST(gr_liv_area AS DOUBLE) / nullif(lot_area, 0) AS liv_lot_ratio, + CAST(gr_liv_area AS DOUBLE) / nullif(tot_rms_abv_grd, 0) AS sf_per_room, + CAST(bedroom_abv_gr AS DOUBLE) / nullif(full_bath + half_bath, 0) AS bed_bath_ratio, + bsmt_fin_sf1 / nullif(total_bsmt_sf, 0.0) AS bsmt_fin_ratio, + bsmt_unf_sf / nullif(total_bsmt_sf, 0.0) AS bsmt_unf_ratio, + garage_area / nullif(garage_cars, 0) AS garage_area_per_car, + overall_qual * overall_cond AS qual_cond_cross, + overall_qual * overall_qual AS overall_qual_sq, + overall_qual * gr_liv_area AS qual_sf_cross, + coalesce(nbhd_price_enc.mean_price, 180921.0) AS nbhd_price, + nbhd_price_enc.mean_price AS nbhd_price_raw, + coalesce(nbhd_price_enc.freq, 0.0) AS nbhd_freq, + coalesce(subclass_enc.mean_price, 180921.0) AS subclass_price, + coalesce(subclass_enc.freq, 0.0) AS subclass_freq, + CASE WHEN mo_sold >= 3 AND mo_sold <= 5 THEN 1 + WHEN mo_sold >= 6 AND mo_sold <= 8 THEN 2 + WHEN mo_sold >= 9 AND mo_sold <= 11 THEN 3 + ELSE 0 END AS season_sold, + CASE WHEN mo_sold >= 5 AND mo_sold <= 7 THEN 1 ELSE 0 END AS is_peak_season +FROM __THIS__ +LEFT JOIN nbhd_price_enc ON neighborhood = nbhd_price_enc.nbhd +LEFT JOIN subclass_enc ON ms_sub_class = subclass_enc.sub_class""" + + +def handcrafted(statics: dict[str, pa.Table]) -> Callable[[dict], dict]: + t = statics["nbhd_price_enc"] + nb = { + k: (mp, fr, mf) + for k, mp, fr, mf in zip( + t.column("nbhd").to_pylist(), + t.column("mean_price").to_pylist(), + t.column("freq").to_pylist(), + t.column("median_frontage").to_pylist(), + strict=True, + ) + } + t = statics["subclass_enc"] + sc = { + k: (mp, fr) + for k, mp, fr in zip( + t.column("sub_class").to_pylist(), + t.column("mean_price").to_pylist(), + t.column("freq").to_pylist(), + strict=True, + ) + } + q_ord = {"Ex": 5, "Gd": 4, "TA": 3, "Fa": 2, "Po": 1} + + def fe(r: dict) -> dict: + tb = r["total_bsmt_sf"] + fin = r["bsmt_fin_sf1"] + unf = r["bsmt_unf_sf"] + ga = r["garage_area"] + gc = r["garage_cars"] + gy = r["garage_yr_blt"] + lf = r["lot_frontage"] + la = r["lot_area"] + mv = r["mas_vnr_area"] + kq = r["kitchen_qual"] + bf = r["bsmt_full_bath"] + bh = r["bsmt_half_bath"] + oq = r["overall_qual"] + gla = r["gr_liv_area"] + bath_den = r["full_bath"] + r["half_bath"] + nbe = nb.get(r["neighborhood"]) + sce = sc.get(r["ms_sub_class"]) + lff = lf + if lff is None: + lff = nbe[2] if nbe is not None else None + if lff is None: + lff = 69.0 + return { + "id": r["id"], + "total_sf": (tb if tb is not None else 0.0) + + r["first_flr_sf"] + + r["second_flr_sf"], + "total_bath": r["full_bath"] + + 0.5 * r["half_bath"] + + (bf if bf is not None else 0) + + 0.5 * (bh if bh is not None else 0), + "house_age": r["yr_sold"] - r["year_built"], + "remod_age": r["yr_sold"] - r["year_remod_add"], + "garage_age": None if gy is None else r["yr_sold"] - gy, + "is_new": 1 if r["yr_sold"] == r["year_built"] else 0, + "is_remodeled": 1 if r["year_remod_add"] > r["year_built"] else 0, + "exter_qual_ord": q_ord.get(r["exter_qual"], 0), + "exter_cond_ord": q_ord.get(r["exter_cond"], 0), + "kitchen_qual_ord": 3 if kq is None else q_ord.get(kq, 0), + "heating_qc_ord": q_ord.get(r["heating_qc"], 0), + "bsmt_qual_ord": q_ord.get(r["bsmt_qual"], 0), + "fireplace_qu_ord": q_ord.get(r["fireplace_qu"], 0), + "garage_qual_ord": q_ord.get(r["garage_qual"], 0), + "porch_total": r["wood_deck_sf"] + + r["open_porch_sf"] + + r["enclosed_porch"] + + r["three_ssn_porch"] + + r["screen_porch"], + "has_pool": 1 if r["pool_area"] > 0 else 0, + "has_garage": 1 if (ga if ga is not None else 0.0) > 0.0 else 0, + "has_bsmt": 1 if (tb if tb is not None else 0.0) > 0.0 else 0, + "has_fireplace": 1 if r["fireplaces"] > 0 else 0, + "has_2nd_floor": 1 if r["second_flr_sf"] > 0 else 0, + "has_mas_vnr": 1 if (mv if mv is not None else 0.0) > 0.0 else 0, + "central_air_flag": 1 if r["central_air"] == "Y" else 0, + "multi_kitchen": 1 if r["kitchen_abv_gr"] > 1 else 0, + "lot_frontage_filled": lff, + "lot_frontage_ratio": None if lf is None or la == 0 else lf / la, + "liv_lot_ratio": None if la == 0 else gla / la, + "sf_per_room": None + if r["tot_rms_abv_grd"] == 0 + else gla / r["tot_rms_abv_grd"], + "bed_bath_ratio": None if bath_den == 0 else r["bedroom_abv_gr"] / bath_den, + "bsmt_fin_ratio": None + if fin is None or tb is None or tb == 0.0 + else fin / tb, + "bsmt_unf_ratio": None + if unf is None or tb is None or tb == 0.0 + else unf / tb, + "garage_area_per_car": None + if ga is None or gc is None or gc == 0 + else ga / gc, + "qual_cond_cross": oq * r["overall_cond"], + "overall_qual_sq": oq * oq, + "qual_sf_cross": oq * gla, + "nbhd_price": nbe[0] if nbe is not None else 180921.0, + "nbhd_price_raw": nbe[0] if nbe is not None else None, + "nbhd_freq": nbe[1] if nbe is not None else 0.0, + "subclass_price": sce[0] if sce is not None else 180921.0, + "subclass_freq": sce[1] if sce is not None else 0.0, + "season_sold": 1 + if 3 <= r["mo_sold"] <= 5 + else 2 + if 6 <= r["mo_sold"] <= 8 + else 3 + if 9 <= r["mo_sold"] <= 11 + else 0, + "is_peak_season": 1 if 5 <= r["mo_sold"] <= 7 else 0, + } + + return fe diff --git a/benchmarks/serving_scenarios/store_sales.py b/benchmarks/serving_scenarios/store_sales.py new file mode 100644 index 0000000..493dfa4 --- /dev/null +++ b/benchmarks/serving_scenarios/store_sales.py @@ -0,0 +1,496 @@ +"""store_sales — Rossmann/Walmart-style store-sales feature serving. + +Serving shape reproduced here: + * the Rossmann ``store.csv`` dim (store_type, assortment, competition_*, + promo2_since_*) LEFT-JOINed on store_id — with ~10% of serving store_ids + absent from the dim (new stores), so join-miss NULL semantics are real; + * per-store mean-sales / mean-customers / sales-per-customer target + encodings (what a fitted mean encoding IS at serve time), with partial + coverage and literal global priors as fallback; + * day-of-week / month / region seasonal-factor encodings; + * competition-open-months and promo2-active-weeks arithmetic with + NULL-guarded clamping, promo interactions, day-of-week one-hots, + state-holiday CASEs, a store_type x promo cross, ratios vs store means, + and Walmart-style econ covariates (fuel/CPI/unemployment/markdowns). +""" + +from __future__ import annotations + +import math +import random +from collections.abc import Callable + +import pyarrow as pa + +NAME = "store_sales" +KAGGLE = ( + "Rossmann Store Sales (+ Walmart Recruiting) serving path: store.csv dim join, " + "competition-open-months / promo2-weeks arithmetic, per-store mean-target encodings " + "with global priors, dow/month seasonal factors, store_type x promo cross, " + "holiday CASEs, Walmart econ covariates." +) + +N_INPUT_COLS = 21 +N_OUTPUT_COLS = 44 + +ROW_SCHEMA: dict[str, str] = { + "record_id": "int", # the Rossmann test.csv `Id`: unique serving-row key + "store_id": "int", + "day_of_week": "int", # 1=Mon .. 7=Sun + "day_of_month": "int", + "month": "int", + "year": "int", + "week_of_year": "int", + "is_open": "bool", + "promo": "bool", + "promo2": "bool", # store enrolled in Promo2, known to the caller + "school_holiday": "bool", + "state_holiday": "str", # '0', 'a' public, 'b' easter, 'c' christmas + "customers_expected": "float?", # upstream forecast, sometimes missing + "temperature": "float?", + "fuel_price": "float?", + "cpi": "float?", + "unemployment": "float?", + "markdown_total": "float?", # Walmart-style markdown spend, often absent + "days_since_prev_promo": "int?", + "region": "str", + "competitor_promo": "bool", +} + +# Fitted constants a real pipeline would bake into the generated SQL at fit +# time: global-prior fallbacks for encoding misses and imputation fills. +PRIOR_MEAN_SALES = 5773.0 +PRIOR_MEAN_CUSTOMERS = 633.0 +PRIOR_SPC = 9.12 +COMP_DIST_FILL = 75000.0 +PRIOR_FUEL = 3.45 +PRIOR_UNEMP = 8.0 +DAYS_SINCE_PROMO_FILL = 30 + +N_STORES = 60 # serving traffic hits 1..60 +DIM_COVERED = 54 # store.csv knows 1..54 (55..60 are new stores -> join miss) +STATS_COVERED = 50 # encodings fitted for 1..50 (51..60 lack history) + +_COMP_MONTHS = ( + "(__THIS__.year - store_dim.competition_open_since_year) * 12" + " + (__THIS__.month - store_dim.competition_open_since_month)" +) +_PROMO2_WEEKS = ( + "(__THIS__.year - store_dim.promo2_since_year) * 52" + " + (__THIS__.week_of_year - store_dim.promo2_since_week)" +) + +SQL = f""" +SELECT + __THIS__.record_id AS record_id, + __THIS__.store_id AS store_id, + CASE WHEN __THIS__.is_open THEN 1 ELSE 0 END AS is_open_flag, + CASE WHEN __THIS__.promo THEN 1 ELSE 0 END AS promo_flag, + CASE WHEN __THIS__.promo2 THEN 1 ELSE 0 END AS promo2_flag, + CASE WHEN __THIS__.school_holiday THEN 1 ELSE 0 END AS school_holiday_flag, + CASE WHEN __THIS__.state_holiday = 'a' THEN 1 ELSE 0 END AS state_hol_public, + CASE WHEN __THIS__.state_holiday = 'b' THEN 1 ELSE 0 END AS state_hol_easter, + CASE WHEN __THIS__.state_holiday = 'c' THEN 1 ELSE 0 END AS state_hol_christmas, + CASE WHEN __THIS__.state_holiday = '0' THEN 0 ELSE 1 END AS any_state_holiday, + CASE WHEN __THIS__.state_holiday <> '0' AND __THIS__.school_holiday + THEN 1 ELSE 0 END AS holiday_x_school, + CASE WHEN __THIS__.day_of_week = 1 THEN 1 ELSE 0 END AS dow_mon, + CASE WHEN __THIS__.day_of_week = 2 THEN 1 ELSE 0 END AS dow_tue, + CASE WHEN __THIS__.day_of_week = 3 THEN 1 ELSE 0 END AS dow_wed, + CASE WHEN __THIS__.day_of_week = 4 THEN 1 ELSE 0 END AS dow_thu, + CASE WHEN __THIS__.day_of_week = 5 THEN 1 ELSE 0 END AS dow_fri, + CASE WHEN __THIS__.day_of_week = 6 THEN 1 ELSE 0 END AS dow_sat, + CASE WHEN __THIS__.day_of_week = 7 THEN 1 ELSE 0 END AS dow_sun, + CASE WHEN __THIS__.day_of_week >= 6 THEN 1 ELSE 0 END AS is_weekend, + CASE WHEN __THIS__.promo AND __THIS__.day_of_week <= 5 + THEN 1 ELSE 0 END AS promo_weekday, + CASE WHEN __THIS__.competitor_promo AND NOT __THIS__.promo + THEN 1 ELSE 0 END AS comp_promo_pressure, + CASE WHEN store_dim.store_type = 'a' THEN 1 + WHEN store_dim.store_type = 'b' THEN 2 + WHEN store_dim.store_type = 'c' THEN 3 + WHEN store_dim.store_type = 'd' THEN 4 + ELSE 0 END AS store_type_ord, + CASE WHEN store_dim.assortment = 'a' THEN 1 + WHEN store_dim.assortment = 'b' THEN 2 + WHEN store_dim.assortment = 'c' THEN 3 + ELSE 0 END AS assortment_ord, + CASE WHEN NOT __THIS__.promo THEN 0 + WHEN store_dim.store_type = 'a' THEN 1 + WHEN store_dim.store_type = 'b' THEN 2 + WHEN store_dim.store_type = 'c' THEN 3 + WHEN store_dim.store_type = 'd' THEN 4 + ELSE 0 END AS promo_x_store_type, + COALESCE(store_dim.competition_distance, {COMP_DIST_FILL}) AS comp_distance, + 1.0 / (1.0 + COALESCE(store_dim.competition_distance, {COMP_DIST_FILL})) + AS comp_distance_inv, + CASE WHEN {_COMP_MONTHS} IS NULL THEN 0 + WHEN {_COMP_MONTHS} < 0 THEN 0 + ELSE {_COMP_MONTHS} END AS comp_open_months, + CASE WHEN NOT __THIS__.promo2 THEN 0 + WHEN {_PROMO2_WEEKS} IS NULL THEN 0 + WHEN {_PROMO2_WEEKS} < 0 THEN 0 + ELSE {_PROMO2_WEEKS} END AS promo2_active_weeks, + COALESCE(store_stats.mean_sales, {PRIOR_MEAN_SALES}) AS store_mean_sales, + COALESCE(store_stats.mean_customers, {PRIOR_MEAN_CUSTOMERS}) + AS store_mean_customers, + COALESCE(store_stats.sales_per_customer, {PRIOR_SPC}) AS store_spc, + COALESCE(region_stats.region_factor, 1.0) AS region_factor, + COALESCE(store_stats.mean_sales, {PRIOR_MEAN_SALES}) + * COALESCE(dow_stats.sales_factor, 1.0) + * COALESCE(month_stats.sales_factor, 1.0) + * COALESCE(region_stats.region_factor, 1.0) + * CASE WHEN __THIS__.promo THEN 1.22 ELSE 1.0 END AS expected_sales, + COALESCE(store_stats.mean_customers, {PRIOR_MEAN_CUSTOMERS}) + * COALESCE(dow_stats.customers_factor, 1.0) + * COALESCE(month_stats.customers_factor, 1.0) AS expected_customers, + __THIS__.customers_expected + / NULLIF(COALESCE(store_stats.mean_customers, {PRIOR_MEAN_CUSTOMERS}), 0.0) + AS customers_ratio, + __THIS__.customers_expected + * COALESCE(store_stats.sales_per_customer, {PRIOR_SPC}) AS forecast_sales, + COALESCE(__THIS__.markdown_total, 0.0) AS markdown_filled, + CASE WHEN __THIS__.markdown_total IS NULL THEN 0 ELSE 1 END AS has_markdown, + COALESCE(__THIS__.markdown_total, 0.0) + / NULLIF(__THIS__.customers_expected, 0.0) AS markdown_per_customer, + (__THIS__.temperature - 15.0) / 10.0 AS temp_norm, + __THIS__.cpi / NULLIF(__THIS__.unemployment, 0.0) AS cpi_unemployment_ratio, + COALESCE(__THIS__.fuel_price, {PRIOR_FUEL}) + * COALESCE(__THIS__.unemployment, {PRIOR_UNEMP}) AS econ_pressure, + COALESCE(__THIS__.days_since_prev_promo, {DAYS_SINCE_PROMO_FILL}) + AS days_since_promo_filled, + CASE WHEN __THIS__.promo + AND COALESCE(__THIS__.days_since_prev_promo, {DAYS_SINCE_PROMO_FILL}) < 7 + THEN 1 ELSE 0 END AS promo_fatigue +FROM __THIS__ +LEFT JOIN store_dim ON __THIS__.store_id = store_dim.store_id +LEFT JOIN store_stats ON __THIS__.store_id = store_stats.store_id +LEFT JOIN dow_stats ON __THIS__.day_of_week = dow_stats.day_of_week +LEFT JOIN month_stats ON __THIS__.month = month_stats.month +LEFT JOIN region_stats ON __THIS__.region = region_stats.region +""" + +_REGIONS = ["north", "south", "east", "west"] + + +def make_statics(seed: int) -> dict[str, pa.Table]: + rng = random.Random(seed) + + dim_rows = [] + for sid in range(1, DIM_COVERED + 1): + enrolled = rng.random() < 0.5 + dim_rows.append( + { + "store_id": sid, + "store_type": rng.choices("abcd", weights=[54, 2, 13, 31])[0], + "assortment": rng.choices("abc", weights=[53, 1, 46])[0], + # prepare-time imputation keeps the static NULL-free, like the + # classic fill of missing CompetitionDistance with a far value + "competition_distance": round( + math.exp(rng.uniform(math.log(30.0), math.log(75000.0))), 1 + ), + "competition_open_since_month": rng.randint(1, 12), + "competition_open_since_year": rng.randint(2000, 2015), + # sentinel far-future start for stores not enrolled in Promo2 + # (clamps to 0 active weeks) + "promo2_since_week": rng.randint(1, 52) if enrolled else 1, + "promo2_since_year": rng.randint(2009, 2015) if enrolled else 2099, + } + ) + + stats_rows = [] + for sid in range(1, STATS_COVERED + 1): + mean_customers = round(rng.uniform(300.0, 1400.0), 2) + spc = round(rng.uniform(6.5, 11.5), 4) + mean_sales = round(mean_customers * spc, 2) + stats_rows.append( + { + "store_id": sid, + "mean_sales": mean_sales, + "mean_customers": mean_customers, + "sales_per_customer": round(mean_sales / mean_customers, 4), + } + ) + + dow_base = [1.15, 0.98, 0.95, 0.96, 1.02, 1.08, 0.45] + dow_rows = [ + { + "day_of_week": d + 1, + "sales_factor": round(dow_base[d] + rng.uniform(-0.02, 0.02), 4), + "customers_factor": round(dow_base[d] + rng.uniform(-0.03, 0.03), 4), + } + for d in range(7) + ] + + month_base = [0.96, 0.94, 0.99, 1.0, 1.01, 0.98, 1.02, 0.99, 0.97, 1.0, 1.05, 1.35] + month_rows = [ + { + "month": m + 1, + "sales_factor": round(month_base[m] + rng.uniform(-0.02, 0.02), 4), + "customers_factor": round(month_base[m] + rng.uniform(-0.03, 0.03), 4), + } + for m in range(12) + ] + + region_base = {"north": 1.04, "south": 0.97, "east": 0.92, "west": 1.07} + region_rows = [ + { + "region": r, + "region_factor": round(region_base[r] + rng.uniform(-0.02, 0.02), 4), + } + for r in _REGIONS + ] + + return { + "store_dim": pa.Table.from_pylist( + dim_rows, + schema=pa.schema( + [ + ("store_id", pa.int64()), + ("store_type", pa.string()), + ("assortment", pa.string()), + ("competition_distance", pa.float64()), + ("competition_open_since_month", pa.int64()), + ("competition_open_since_year", pa.int64()), + ("promo2_since_week", pa.int64()), + ("promo2_since_year", pa.int64()), + ] + ), + ), + "store_stats": pa.Table.from_pylist( + stats_rows, + schema=pa.schema( + [ + ("store_id", pa.int64()), + ("mean_sales", pa.float64()), + ("mean_customers", pa.float64()), + ("sales_per_customer", pa.float64()), + ] + ), + ), + "dow_stats": pa.Table.from_pylist( + dow_rows, + schema=pa.schema( + [ + ("day_of_week", pa.int64()), + ("sales_factor", pa.float64()), + ("customers_factor", pa.float64()), + ] + ), + ), + "month_stats": pa.Table.from_pylist( + month_rows, + schema=pa.schema( + [ + ("month", pa.int64()), + ("sales_factor", pa.float64()), + ("customers_factor", pa.float64()), + ] + ), + ), + "region_stats": pa.Table.from_pylist( + region_rows, + schema=pa.schema( + [("region", pa.string()), ("region_factor", pa.float64())] + ), + ), + } + + +def make_rows(seed: int, n: int) -> list[dict]: + rng = random.Random(seed) + rows = [] + for i in range(n): + dow = rng.randint(1, 7) + month = rng.randint(1, 12) + state_holiday = rng.choices(["0", "a", "b", "c"], weights=[91, 5, 2, 2])[0] + open_p = 0.1 if state_holiday != "0" else (0.6 if dow == 7 else 0.97) + rows.append( + { + "record_id": 9_000_000 + i, + "store_id": rng.randint(1, N_STORES), + "day_of_week": dow, + "day_of_month": rng.randint(1, 28), + "month": month, + "year": rng.choice([2014, 2015]), + "week_of_year": (month - 1) * 4 + rng.randint(1, 4), + "is_open": rng.random() < open_p, + "promo": dow <= 5 and rng.random() < 0.45, + "promo2": rng.random() < 0.5, + "school_holiday": rng.random() < 0.18, + "state_holiday": state_holiday, + "customers_expected": ( + None + if rng.random() < 0.12 + else round(rng.uniform(200.0, 1600.0), 1) + ), + "temperature": ( + None if rng.random() < 0.05 else round(rng.uniform(-5.0, 35.0), 1) + ), + "fuel_price": ( + None if rng.random() < 0.05 else round(rng.uniform(2.4, 4.5), 3) + ), + "cpi": None + if rng.random() < 0.08 + else round(rng.uniform(126.0, 228.0), 2), + "unemployment": ( + None if rng.random() < 0.08 else round(rng.uniform(3.8, 14.3), 3) + ), + "markdown_total": ( + None + if rng.random() < 0.55 + else round(rng.uniform(50.0, 20000.0), 2) + ), + "days_since_prev_promo": ( + None if rng.random() < 0.10 else rng.randint(0, 60) + ), + "region": rng.choice(_REGIONS), + "competitor_promo": rng.random() < 0.3, + } + ) + return rows + + +def handcrafted(statics: dict[str, pa.Table]) -> Callable[[dict], dict]: + """What a competent engineer hand-writes for a Python microservice: + plain dict lookups prepared once, then a per-row closure computing the + identical features with SQL NULL semantics (join miss => None + propagation through arithmetic, CASE falls through on unknown).""" + dim = {r["store_id"]: r for r in statics["store_dim"].to_pylist()} + stats = {r["store_id"]: r for r in statics["store_stats"].to_pylist()} + dow_f = {r["day_of_week"]: r for r in statics["dow_stats"].to_pylist()} + month_f = {r["month"]: r for r in statics["month_stats"].to_pylist()} + region_f = {r["region"]: r for r in statics["region_stats"].to_pylist()} + + def fn(row: dict) -> dict: + d = dim.get(row["store_id"]) + s = stats.get(row["store_id"]) + w = dow_f[row["day_of_week"]] # full coverage by construction + m = month_f[row["month"]] + g = region_f[row["region"]] + + dow = row["day_of_week"] + promo = row["promo"] + hol = row["state_holiday"] + + st = d["store_type"] if d is not None else None + sort = d["assortment"] if d is not None else None + comp_dist = d["competition_distance"] if d is not None else COMP_DIST_FILL + + if d is None: + comp_open_months = 0 + else: + cm = (row["year"] - d["competition_open_since_year"]) * 12 + ( + row["month"] - d["competition_open_since_month"] + ) + comp_open_months = 0 if cm < 0 else cm + + if not row["promo2"] or d is None: + promo2_active_weeks = 0 + else: + pw = (row["year"] - d["promo2_since_year"]) * 52 + ( + row["week_of_year"] - d["promo2_since_week"] + ) + promo2_active_weeks = 0 if pw < 0 else pw + + mean_sales = s["mean_sales"] if s is not None else PRIOR_MEAN_SALES + mean_customers = s["mean_customers"] if s is not None else PRIOR_MEAN_CUSTOMERS + spc = s["sales_per_customer"] if s is not None else PRIOR_SPC + + ce = row["customers_expected"] + md = row["markdown_total"] + md_filled = md if md is not None else 0.0 + t = row["temperature"] + cpi = row["cpi"] + unemp = row["unemployment"] + fuel = row["fuel_price"] + ds = row["days_since_prev_promo"] + ds_filled = ds if ds is not None else DAYS_SINCE_PROMO_FILL + + return { + "record_id": row["record_id"], + "store_id": row["store_id"], + "is_open_flag": 1 if row["is_open"] else 0, + "promo_flag": 1 if promo else 0, + "promo2_flag": 1 if row["promo2"] else 0, + "school_holiday_flag": 1 if row["school_holiday"] else 0, + "state_hol_public": 1 if hol == "a" else 0, + "state_hol_easter": 1 if hol == "b" else 0, + "state_hol_christmas": 1 if hol == "c" else 0, + "any_state_holiday": 0 if hol == "0" else 1, + "holiday_x_school": 1 if hol != "0" and row["school_holiday"] else 0, + "dow_mon": 1 if dow == 1 else 0, + "dow_tue": 1 if dow == 2 else 0, + "dow_wed": 1 if dow == 3 else 0, + "dow_thu": 1 if dow == 4 else 0, + "dow_fri": 1 if dow == 5 else 0, + "dow_sat": 1 if dow == 6 else 0, + "dow_sun": 1 if dow == 7 else 0, + "is_weekend": 1 if dow >= 6 else 0, + "promo_weekday": 1 if promo and dow <= 5 else 0, + "comp_promo_pressure": 1 if row["competitor_promo"] and not promo else 0, + # None == 'x' is False in Python, matching CASE falling through + # NULL comparisons to ELSE + "store_type_ord": ( + 1 + if st == "a" + else 2 + if st == "b" + else 3 + if st == "c" + else 4 + if st == "d" + else 0 + ), + "assortment_ord": ( + 1 if sort == "a" else 2 if sort == "b" else 3 if sort == "c" else 0 + ), + "promo_x_store_type": ( + 0 + if not promo + else 1 + if st == "a" + else 2 + if st == "b" + else 3 + if st == "c" + else 4 + if st == "d" + else 0 + ), + "comp_distance": comp_dist, + "comp_distance_inv": 1.0 / (1.0 + comp_dist), + "comp_open_months": comp_open_months, + "promo2_active_weeks": promo2_active_weeks, + "store_mean_sales": mean_sales, + "store_mean_customers": mean_customers, + "store_spc": spc, + "region_factor": g["region_factor"], + "expected_sales": mean_sales + * w["sales_factor"] + * m["sales_factor"] + * g["region_factor"] + * (1.22 if promo else 1.0), + "expected_customers": mean_customers + * w["customers_factor"] + * m["customers_factor"], + "customers_ratio": ( + None if ce is None or mean_customers == 0.0 else ce / mean_customers + ), + "forecast_sales": None if ce is None else ce * spc, + "markdown_filled": md_filled, + "has_markdown": 0 if md is None else 1, + "markdown_per_customer": ( + None if ce is None or ce == 0.0 else md_filled / ce + ), + "temp_norm": None if t is None else (t - 15.0) / 10.0, + "cpi_unemployment_ratio": ( + None if cpi is None or unemp is None or unemp == 0.0 else cpi / unemp + ), + "econ_pressure": (fuel if fuel is not None else PRIOR_FUEL) + * (unemp if unemp is not None else PRIOR_UNEMP), + "days_since_promo_filled": ds_filled, + "promo_fatigue": 1 if promo and ds_filled < 7 else 0, + } + + return fn diff --git a/benchmarks/serving_scenarios/titanic.py b/benchmarks/serving_scenarios/titanic.py index e5dd957..7fcb2f3 100644 --- a/benchmarks/serving_scenarios/titanic.py +++ b/benchmarks/serving_scenarios/titanic.py @@ -1,293 +1,361 @@ -"""Titanic survival serving scenario for the SQL specializer benchmark. +"""Titanic survival — THE classic Kaggle problem, as a serving-path scenario. -Reproduces the canonical Kaggle-Titanic public-kernel feature pipeline as it -looks at SERVE time: scalar expressions over the raw passenger row plus LEFT -JOINs to prepare-time fitted tables (per-pclass fare medians, per-title age -medians and group mean fares, sex-x-pclass target-mean encoding, embarked -target encoding). +Reproduces the canonical public-kernel feature pipeline at serve time: +scalar expressions over the passenger row + LEFT JOINs to fitted encoding +tables (a target-mean encoding IS a join table at serve time). """ +from __future__ import annotations + +import math import random from collections.abc import Callable import pyarrow as pa -NAME = "titanic_survival_features" +NAME = "titanic" KAGGLE = ( - "Kaggle Titanic (survival prediction) — canonical public-kernel tricks: " - "family_size/is_alone, fare_per_person with per-pclass median imputation, " - "cabin deck letter + has_cabin, title-based age imputation with missing " - "flag, age*pclass interaction, age bins, embarked one-hots, sex-x-pclass " - "target-mean encoding and title-GROUP mean fare as fitted join tables." + "Titanic: Machine Learning from Disaster — the canonical public-kernel recipe: " + "FamilySize/IsAlone, FarePerPerson, deck letter from Cabin + HasCabin, " + "Age-imputation flag + Age*Pclass interaction, age bins, Embarked one-hots, " + "sex x pclass survival target-mean encoding and title-group mean-fare/survival " + "encodings served as fitted join tables (incl. the famous unseen-'Dona' miss)." ) N_INPUT_COLS = 10 -N_OUTPUT_COLS = 20 +N_OUTPUT_COLS = 24 ROW_SCHEMA = { "passenger_id": "int", "pclass": "int", "sex": "str", + "title": "str", "age": "float?", "sibsp": "int", "parch": "int", "fare": "float?", "cabin": "str?", "embarked": "str?", - "title": "str", } -# Fitted fallback constants a real pipeline bakes into its generated SQL. -_GLOBAL_MEDIAN_AGE = 29.7 -_GLOBAL_TITLE_FARE = 32.2 -_MODE_PORT_RATE = 0.339 - -SQL = """ -SELECT - passenger_id, - sibsp + parch + 1 AS family_size, - CASE WHEN sibsp + parch = 0 THEN 1 ELSE 0 END AS is_alone, - coalesce(fare, pclass_dim.median_fare) AS fare_filled, - coalesce(fare, pclass_dim.median_fare) / (sibsp + parch + 1) AS fare_per_person, - coalesce(upper(substr(trim(cabin), 1, 1)), 'U') AS deck, - CASE WHEN cabin IS NULL THEN 0 ELSE 1 END AS has_cabin, - CASE WHEN age IS NULL THEN 1 ELSE 0 END AS age_missing, - coalesce(age, title_stats.median_age, 29.7) AS age_filled, - coalesce(age, title_stats.median_age, 29.7) * pclass AS age_class, - CASE - WHEN coalesce(age, title_stats.median_age, 29.7) < 13.0 THEN 'child' - WHEN coalesce(age, title_stats.median_age, 29.7) < 20.0 THEN 'teen' - WHEN coalesce(age, title_stats.median_age, 29.7) < 60.0 THEN 'adult' - ELSE 'senior' - END AS age_bin, - CASE WHEN coalesce(embarked, 'S') = 'C' THEN 1 ELSE 0 END AS embarked_c, - CASE WHEN coalesce(embarked, 'S') = 'Q' THEN 1 ELSE 0 END AS embarked_q, - CASE WHEN coalesce(embarked, 'S') = 'S' THEN 1 ELSE 0 END AS embarked_s, - CASE WHEN sex = 'female' THEN 1 ELSE 0 END AS sex_female, - sex || '-' || CAST(pclass AS VARCHAR) AS sex_pclass, - sex_pclass_enc.survival_rate AS sex_pclass_rate, - coalesce(title_stats.group_mean_fare, 32.2) AS title_fare_mean, - pclass_dim.pclass_survival_rate AS pclass_rate, - coalesce(embarked_enc.port_survival_rate, 0.339) AS port_rate -FROM __THIS__ -LEFT JOIN pclass_dim ON pclass = pclass_dim.pc -LEFT JOIN title_stats ON title = title_stats.t_title -LEFT JOIN sex_pclass_enc - ON sex = sex_pclass_enc.sp_sex AND pclass = sex_pclass_enc.sp_pclass -LEFT JOIN embarked_enc ON embarked = embarked_enc.port -""" +# Fitted constants a pipeline would bake into the serving query as literals. +AGE_MEDIAN = 28.0 +FARE_MEDIAN = 14.4542 +GLOBAL_RATE = 0.383838 -# Title -> (group median age, denormalized title-GROUP mean fare). The -# grouping {Mr, Mrs, Miss, Master, Rare} happened at fit time; the fitted -# table is keyed on the raw title. Capt/Jonkheer/Dona are deliberately NOT -# here: unseen-at-fit titles cause serve-time LEFT JOIN misses. -_TITLES = { - "Mr": (30.0, 24.4), - "Mrs": (35.0, 45.0), - "Miss": (21.0, 43.8), - "Master": (3.5, 37.0), - "Dr": (46.5, 40.9), - "Rev": (46.5, 40.9), - "Col": (46.5, 40.9), - "Major": (46.5, 40.9), - "Mlle": (21.0, 43.8), - "Ms": (21.0, 43.8), - "Mme": (35.0, 45.0), - "Lady": (46.5, 40.9), - "Sir": (46.5, 40.9), - "Countess": (46.5, 40.9), +# Title -> group, as the canonical kernels collapse rare titles. "Dona" is +# deliberately NOT here: it appears only in the test set and is the classic +# unseen-category trap — served as a LEFT JOIN miss + coalesce fallback. +_TITLE_TO_GROUP = { + "Mr": "Mr", + "Mrs": "Mrs", + "Mme": "Mrs", + "Miss": "Miss", + "Mlle": "Miss", + "Ms": "Miss", + "Master": "Master", + "Dr": "Rare", + "Rev": "Rare", + "Col": "Rare", + "Major": "Rare", + "Capt": "Rare", + "Sir": "Rare", + "Lady": "Rare", + "Don": "Rare", + "Countess": "Rare", + "Jonkheer": "Rare", } def make_statics(seed: int) -> dict[str, pa.Table]: - r = random.Random(seed) + rnd = random.Random(seed) - def jit(x: float) -> float: - return round(x + r.uniform(-0.015, 0.015), 6) + def jit(v: float) -> float: + return round(v + rnd.uniform(-0.015, 0.015), 6) - pclass_dim = pa.table( - { - "pc": [1, 2, 3], - "median_fare": [round(jit(60.2875), 4), 14.25, 8.05], - "pclass_survival_rate": [jit(0.6296), jit(0.4728), jit(0.2424)], - } + # sex x pclass survival target-mean encoding (train-set rates + fit noise). + sp_base = [ + ("female", 1, 0.968, 94), + ("female", 2, 0.921, 76), + ("female", 3, 0.500, 144), + ("male", 1, 0.369, 122), + ("male", 2, 0.157, 108), + ("male", 3, 0.135, 347), + ] + sex_pclass_enc = pa.Table.from_pylist( + [ + {"sex": s, "pclass": p, "survival_rate": jit(r), "n": n} + for s, p, r, n in sp_base + ], + schema=pa.schema( + [ + ("sex", pa.string()), + ("pclass", pa.int64()), + ("survival_rate", pa.float64()), + ("n", pa.int64()), + ] + ), ) - titles = sorted(_TITLES) - title_stats = pa.table( - { - "t_title": titles, - "median_age": [_TITLES[t][0] for t in titles], - "group_mean_fare": [jit(_TITLES[t][1]) for t in titles], - } - ) - sex_pclass_enc = pa.table( - { - "sp_sex": ["female", "female", "female", "male", "male", "male"], - "sp_pclass": [1, 2, 3, 1, 2, 3], - "survival_rate": [ - jit(0.9681), - jit(0.9211), - jit(0.5000), - jit(0.3689), - jit(0.1574), - jit(0.1354), - ], - } + + # Title-group encodings: every title carries its GROUP's fitted stats, + # exactly as a materialized groupby-join would. + group_fare = { + k: jit(v) + for k, v in [ + ("Mr", 24.44), + ("Mrs", 45.14), + ("Miss", 43.80), + ("Master", 37.98), + ("Rare", 33.50), + ] + } + group_rate = { + k: jit(v) + for k, v in [ + ("Mr", 0.157), + ("Mrs", 0.792), + ("Miss", 0.703), + ("Master", 0.575), + ("Rare", 0.444), + ] + } + title_dim = pa.Table.from_pylist( + [ + { + "title": t, + "title_group": g, + "group_mean_fare": group_fare[g], + "group_survival_rate": group_rate[g], + } + for t, g in _TITLE_TO_GROUP.items() + ], + schema=pa.schema( + [ + ("title", pa.string()), + ("title_group", pa.string()), + ("group_mean_fare", pa.float64()), + ("group_survival_rate", pa.float64()), + ] + ), ) - embarked_enc = pa.table( - { - "port": ["S", "C", "Q"], - "port_survival_rate": [jit(0.3370), jit(0.5539), jit(0.3896)], - } + + embarked_dim = pa.Table.from_pylist( + [ + {"embarked": "S", "survival_rate": jit(0.339)}, + {"embarked": "C", "survival_rate": jit(0.554)}, + {"embarked": "Q", "survival_rate": jit(0.390)}, + ], + schema=pa.schema([("embarked", pa.string()), ("survival_rate", pa.float64())]), ) + return { - "pclass_dim": pclass_dim, - "title_stats": title_stats, "sex_pclass_enc": sex_pclass_enc, - "embarked_enc": embarked_enc, + "title_dim": title_dim, + "embarked_dim": embarked_dim, } +_MALE_TITLES = [ + "Mr", + "Master", + "Dr", + "Rev", + "Col", + "Major", + "Capt", + "Sir", + "Don", + "Jonkheer", +] +_MALE_W = [80, 8, 3, 2, 2, 1, 1, 1, 1, 1] +# "Dona" (~3% of women) is unseen by title_dim -> the LEFT JOIN miss path. +_FEMALE_TITLES = ["Miss", "Mrs", "Mlle", "Mme", "Ms", "Lady", "Countess", "Dona"] +_FEMALE_W = [47, 41, 3, 2, 2, 1, 1, 3] +_DECKS_BY_CLASS = {1: "ABCDE", 2: "DEF", 3: "EFG"} +_CABIN_MISS_P = {1: 0.20, 2: 0.75, 3: 0.94} + + def make_rows(seed: int, n: int) -> list[dict]: - r = random.Random(seed) - male_titles = ["Mr"] * 90 + ["Master"] * 6 + ["Dr", "Rev", "Capt", "Jonkheer"] - female_titles = ( - ["Miss"] * 44 - + ["Mrs"] * 44 - + ["Mlle", "Ms", "Mme", "Lady", "Countess", "Dona"] * 2 - ) - decks = {1: "ABCDE", 2: "DEF", 3: "EFG"} - fare_base = {1: 84.15, 2: 20.66, 3: 13.68} - rows = [] + rnd = random.Random(seed) + rows: list[dict] = [] for i in range(n): - pclass = r.choices([1, 2, 3], weights=[24, 21, 55])[0] - sex = "male" if r.random() < 0.65 else "female" - title = r.choice(male_titles if sex == "male" else female_titles) - if r.random() < 0.20: + pclass = rnd.choices([1, 2, 3], [24, 21, 55])[0] + sex = "male" if rnd.random() < 0.65 else "female" + if sex == "male": + title = rnd.choices(_MALE_TITLES, _MALE_W)[0] + else: + title = rnd.choices(_FEMALE_TITLES, _FEMALE_W)[0] + + if rnd.random() < 0.199: # 177/891 missing in the train set age = None elif title == "Master": - age = round(r.uniform(0.5, 12.0) * 2) / 2 + age = round(rnd.uniform(0.42, 12.0) * 2) / 2 else: - age = min(80.0, max(14.0, round(r.gauss(30.0, 13.0) * 2) / 2)) - sibsp = r.choices([0, 1, 2, 3, 4, 8], weights=[68, 23, 4, 3, 1, 1])[0] - parch = r.choices([0, 1, 2, 5], weights=[76, 13, 9, 2])[0] - if r.random() < 0.01: + mu, sd = { + "Miss": (22.0, 10.0), + "Mrs": (36.0, 12.0), + "Mr": (32.0, 12.0), + }.get(title, (45.0, 10.0)) + age = round(min(80.0, max(0.42, rnd.gauss(mu, sd))) * 2) / 2 + + sibsp = rnd.choices([0, 1, 2, 3, 4, 5, 8], [68, 23, 3, 2, 2, 1, 1])[0] + parch = rnd.choices([0, 1, 2, 3, 4, 5, 6], [76, 13, 8, 1, 1, 1, 1])[0] + + if rnd.random() < 0.008: # the lone missing Fare is a test-set classic fare = None - elif r.random() < 0.015: - fare = 0.0 else: - fare = round(fare_base[pclass] * (0.35 + r.random() * 2.2), 4) - if r.random() < 0.77: + mu, sd = {1: (4.2, 0.7), 2: (2.7, 0.4), 3: (2.1, 0.45)}[pclass] + fare = round(math.exp(rnd.gauss(mu, sd)), 4) + + if rnd.random() < _CABIN_MISS_P[pclass]: cabin = None else: - letter = r.choice(decks[pclass]) - if r.random() < 0.15: - letter = letter.lower() - cabin = f"{letter}{r.randint(1, 130)}" - if r.random() < 0.10: - cabin = f" {cabin} " + deck = rnd.choice(_DECKS_BY_CLASS[pclass]) + cabin = f"{deck}{rnd.randint(1, 130)}" + if rnd.random() < 0.08: # multi-cabin families: "C23 C25" + cabin = f"{cabin} {deck}{rnd.randint(1, 130)}" + if rnd.random() < 0.15: # messy serving payloads + cabin = cabin.lower() + if rnd.random() < 0.10: + cabin = " " + cabin + if rnd.random() < 0.10: + cabin = cabin + " " + embarked = ( None - if r.random() < 0.02 - else r.choices(["S", "C", "Q"], weights=[72, 19, 9])[0] + if rnd.random() < 0.012 + else rnd.choices(["S", "C", "Q"], [72, 19, 9])[0] ) + rows.append( { - "passenger_id": 1000 + i, + "passenger_id": 892 + i, "pclass": pclass, "sex": sex, + "title": title, "age": age, "sibsp": sibsp, "parch": parch, "fare": fare, "cabin": cabin, "embarked": embarked, - "title": title, } ) return rows +SQL = """ +SELECT + __THIS__.passenger_id AS passenger_id, + __THIS__.pclass AS pclass, + CASE WHEN __THIS__.sex = 'male' THEN 1 ELSE 0 END AS sex_male, + __THIS__.sex || '_' || CAST(__THIS__.pclass AS VARCHAR) AS sex_pclass_key, + __THIS__.sibsp + __THIS__.parch + 1 AS family_size, + CASE WHEN __THIS__.sibsp + __THIS__.parch = 0 THEN 1 ELSE 0 END AS is_alone, + CASE WHEN __THIS__.fare IS NULL THEN 1 ELSE 0 END AS fare_missing, + coalesce(__THIS__.fare, 14.4542) AS fare_filled, + coalesce(__THIS__.fare, 14.4542) / (__THIS__.sibsp + __THIS__.parch + 1) + AS fare_per_person, + CASE WHEN __THIS__.age IS NULL THEN 1 ELSE 0 END AS age_missing, + coalesce(__THIS__.age, 28.0) AS age_filled, + coalesce(__THIS__.age, 28.0) * __THIS__.pclass AS age_x_pclass, + CASE + WHEN __THIS__.age IS NULL THEN 'unknown' + WHEN __THIS__.age < 13.0 THEN 'child' + WHEN __THIS__.age < 20.0 THEN 'teen' + WHEN __THIS__.age < 41.0 THEN 'adult' + WHEN __THIS__.age < 61.0 THEN 'mid' + ELSE 'senior' + END AS age_bin, + CASE WHEN __THIS__.cabin IS NULL THEN 0 ELSE 1 END AS has_cabin, + CASE + WHEN __THIS__.cabin IS NULL THEN 'U' + ELSE upper(substr(trim(__THIS__.cabin), 1, 1)) + END AS deck, + CASE WHEN __THIS__.embarked = 'S' THEN 1 ELSE 0 END AS embarked_s, + CASE WHEN __THIS__.embarked = 'C' THEN 1 ELSE 0 END AS embarked_c, + CASE WHEN __THIS__.embarked = 'Q' THEN 1 ELSE 0 END AS embarked_q, + sp.survival_rate AS sex_pclass_rate, + coalesce(td.title_group, 'Rare') AS title_group, + coalesce(td.group_survival_rate, 0.383838) AS title_rate, + td.group_mean_fare AS title_fare_mean, + coalesce(__THIS__.fare, 14.4542) - td.group_mean_fare AS fare_minus_title_mean, + coalesce(em.survival_rate, 0.383838) AS embarked_rate +FROM __THIS__ +LEFT JOIN sex_pclass_enc AS sp + ON __THIS__.sex = sp.sex AND __THIS__.pclass = sp.pclass +LEFT JOIN title_dim AS td ON __THIS__.title = td.title +LEFT JOIN embarked_dim AS em ON __THIS__.embarked = em.embarked +""" + + def handcrafted(statics: dict[str, pa.Table]) -> Callable[[dict], dict]: - def cols(t: pa.Table) -> dict[str, list]: - return {name: t.column(name).to_pylist() for name in t.column_names} - - p = cols(statics["pclass_dim"]) - pclass_lut = { - k: (mf, sr) - for k, mf, sr in zip( - p["pc"], p["median_fare"], p["pclass_survival_rate"], strict=True - ) + """What a competent engineer hand-writes for a Python microservice: + plain-dict lookups prepared once, a per-row closure after that.""" + sp = { + (r["sex"], r["pclass"]): r["survival_rate"] + for r in statics["sex_pclass_enc"].to_pylist() } - t = cols(statics["title_stats"]) - title_lut = { - k: (ma, gf) - for k, ma, gf in zip( - t["t_title"], t["median_age"], t["group_mean_fare"], strict=True - ) + td = {r["title"]: r for r in statics["title_dim"].to_pylist()} + em = { + r["embarked"]: r["survival_rate"] for r in statics["embarked_dim"].to_pylist() } - s = cols(statics["sex_pclass_enc"]) - sp_lut = { - (sx, pc): sr - for sx, pc, sr in zip( - s["sp_sex"], s["sp_pclass"], s["survival_rate"], strict=True - ) - } - e = cols(statics["embarked_enc"]) - port_lut = dict(zip(e["port"], e["port_survival_rate"], strict=True)) - def fn(row: dict) -> dict: - sibsp, parch, pclass = row["sibsp"], row["parch"], row["pclass"] - family_size = sibsp + parch + 1 - median_fare, pclass_rate = pclass_lut.get(pclass, (None, None)) + def infer(row: dict) -> dict: + sex = row["sex"] + pclass = row["pclass"] + age = row["age"] fare = row["fare"] - fare_filled = fare if fare is not None else median_fare - fare_per_person = None if fare_filled is None else fare_filled / family_size cabin = row["cabin"] - deck = "U" if cabin is None else cabin.strip()[:1].upper() - median_age, group_fare = title_lut.get(row["title"], (None, None)) - age = row["age"] - if age is not None: - age_filled = age - elif median_age is not None: - age_filled = median_age - else: - age_filled = _GLOBAL_MEDIAN_AGE - if age_filled < 13.0: + emb = row["embarked"] + sibsp = row["sibsp"] + parch = row["parch"] + + family_size = sibsp + parch + 1 + fare_filled = fare if fare is not None else 14.4542 + age_filled = age if age is not None else 28.0 + t = td.get(row["title"]) + + if age is None: + age_bin = "unknown" + elif age < 13.0: age_bin = "child" - elif age_filled < 20.0: + elif age < 20.0: age_bin = "teen" - elif age_filled < 60.0: + elif age < 41.0: age_bin = "adult" + elif age < 61.0: + age_bin = "mid" else: age_bin = "senior" - embarked = row["embarked"] - emb = embarked if embarked is not None else "S" - port_rate = port_lut.get(embarked) if embarked is not None else None - sex = row["sex"] + return { "passenger_id": row["passenger_id"], + "pclass": pclass, + "sex_male": 1 if sex == "male" else 0, + "sex_pclass_key": sex + "_" + str(pclass), "family_size": family_size, "is_alone": 1 if sibsp + parch == 0 else 0, + "fare_missing": 1 if fare is None else 0, "fare_filled": fare_filled, - "fare_per_person": fare_per_person, - "deck": deck, - "has_cabin": 0 if cabin is None else 1, + "fare_per_person": fare_filled / family_size, "age_missing": 1 if age is None else 0, "age_filled": age_filled, - "age_class": age_filled * pclass, + "age_x_pclass": age_filled * pclass, "age_bin": age_bin, + "has_cabin": 0 if cabin is None else 1, + "deck": "U" if cabin is None else cabin.strip()[:1].upper(), + "embarked_s": 1 if emb == "S" else 0, "embarked_c": 1 if emb == "C" else 0, "embarked_q": 1 if emb == "Q" else 0, - "embarked_s": 1 if emb == "S" else 0, - "sex_female": 1 if sex == "female" else 0, - "sex_pclass": f"{sex}-{pclass}", - "sex_pclass_rate": sp_lut.get((sex, pclass)), - "title_fare_mean": group_fare - if group_fare is not None - else _GLOBAL_TITLE_FARE, - "pclass_rate": pclass_rate, - "port_rate": port_rate if port_rate is not None else _MODE_PORT_RATE, + "sex_pclass_rate": sp[(sex, pclass)], + "title_group": t["title_group"] if t is not None else "Rare", + "title_rate": t["group_survival_rate"] if t is not None else 0.383838, + "title_fare_mean": t["group_mean_fare"] if t is not None else None, + "fare_minus_title_mean": ( + fare_filled - t["group_mean_fare"] if t is not None else None + ), + "embarked_rate": em.get(emb, 0.383838), } - return fn + return infer diff --git a/pyproject.toml b/pyproject.toml index 55dd67c..b8142df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,9 +54,10 @@ known-third-party = ["datafusion", "duckdb"] "_state.py" = ["S608"] # Doc generator: emits long markdown lines verbatim; wrapping would corrupt output. "scripts/gen_datafusion_catalogue.py" = ["E501"] -# Bench fixtures: seeded pseudo-random data generators (S311 is about crypto) -# and dense SQL/data-gen literals where wrapping hurts readability. -"benchmarks/serving_scenarios/*.py" = ["S311", "E501"] +# Bench fixtures: seeded pseudo-random data generators (S311 is about +# crypto), dense SQL/data-gen literals where wrapping hurts readability, +# and SQL assembled from module-internal constants (S608, no user input). +"benchmarks/serving_scenarios/*.py" = ["S311", "E501", "S608"] [tool.pytest.ini_options] python_files = ["*_test.py", "test_*.py"] From 0c6df0977cfa86d12997351f320c885929e7200d Mon Sep 17 00:00:00 2001 From: AmirHossein Roozbahani Date: Sun, 26 Jul 2026 11:19:53 +0200 Subject: [PATCH 11/11] fix(bench): titanic handcrafted twin trims spaces only, matching SQL trim semantics (reviewer note) Co-Authored-By: Claude Fable 5 --- benchmarks/serving_scenarios/titanic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/serving_scenarios/titanic.py b/benchmarks/serving_scenarios/titanic.py index 7fcb2f3..588d62d 100644 --- a/benchmarks/serving_scenarios/titanic.py +++ b/benchmarks/serving_scenarios/titanic.py @@ -344,7 +344,7 @@ def infer(row: dict) -> dict: "age_x_pclass": age_filled * pclass, "age_bin": age_bin, "has_cabin": 0 if cabin is None else 1, - "deck": "U" if cabin is None else cabin.strip()[:1].upper(), + "deck": "U" if cabin is None else cabin.strip(" ")[:1].upper(), "embarked_s": 1 if emb == "S" else 0, "embarked_c": 1 if emb == "C" else 0, "embarked_q": 1 if emb == "Q" else 0,