diff --git a/CHANGELOG.md b/CHANGELOG.md index 5613a377..924a7ea2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ All notable changes to this project are documented in this file. ### Fixed - **`number_of_cases` is now emitted as a JSON integer instead of source text.** The slot is typed `int` on biolink-model's `EntityToDiseaseAssociation` / `EntityToPhenotypicFeatureAssociation`, but it was absent from `lib.numeric_columns`' exact set — since 15.1's `STUDY_SIZE_EXEMPT_PATTERN` (#119) stopped it being *renamed* onto `study_size`, nothing cast its *values*, so the raw TSV cell (`"1"`) shipped on the edge NDJSON. Pydantic's lax validation coerced the string back to int inside `validate_kgx`, so the record still validated, and the Rust `uuid_on_collision: merge` recompute already wrote a real int — leaving the shipped graph type-inconsistent edge to edge. `number_of_cases` now rides the same `clean_numeric` / `format_numeric` machinery as `study_size`: `numeric_slot_kind` reads the `int` typing off the installed model, fractional / negative / non-numeric counts become null, and the release-mode `drop_low_number_of_cases` filter is untouched (it already cast inline). +### Performance +- **The `uuid_on_collision: merge` edge pass is now near-linear in divergent rows per edge instead of quadratic, and its canonical serializer no longer clones per record.** The merge fold re-canonicalized and re-sorted every list item already stored on an edge on each incoming row — O(n²) in rows × `supporting_case_ids` — and `canonical_json_bytes` deep-cloned into a sorted-key `Value` tree before writing, a fixed cost paid once per record on every pass. Array fields now carry their items' canonical bytes, computed once when an item arrives, so each fold is O(1) per incoming item (hash-set membership, one deferred stable sort at write-out, and only for lists a real union touched), while the sorted-key form is written directly into the output buffer with scalars delegated to `serde_json`'s own writers. Measured on release builds (`uv run maturin develop --release`) over the seeded corpora (seed 42) of the committed harness — `TABLASSERT_BENCH=1 uv run pytest tests/bench_merge_bench.py -s -n 0 -q` — on an AMD Ryzen AI 9 HX 370 with Python 3.13: **37.45 s → 1.311 s (28.6x)** for 100,000 lines of 1k triples × 100 divergent rows at 50 case ids per row, **20.69 s → 2.155 s (9.6x)** for 100,000 lines of 5k triples × 20 rows at 100 case ids per row, **6.71 s → 5.184 s (1.29x)** for 600,000 lines of 200k triples × 3 rows at 10–20 case ids per row, and **0.59 s → 0.495 s (1.19x)** for 100,000 byte-identical rows, so the no-fold path does not pay for the new bookkeeping either. Merged-record and conflict counters are unchanged on every scenario (400,000 / 1,199,899; 95,000 / 284,979; 99,000 / 296,969), and the emitted bytes are identical to the previous fold — pinned by a seeded fuzz against a frozen verbatim copy of the old implementation (`cargo test merge_fold_matches_reference_on_fuzz`), an end-to-end pure-Python reference (`tests/test_rs.py::test_dedup_edges_merge_matches_python_reference`), and an in-process guard that fails below a 5x bound against that frozen quadratic reference (`cargo test fold_speedup`). The buffered merge state does carry per-list-item bookkeeping memory — one canonical-bytes copy per distinct list item, a single allocation shared between the membership set and the ordered write-out list — which is the price of the O(1) membership check and one more reason the mode stays opt-in: sharing that allocation rather than storing two copies drops peak RSS on the harness's 1k triples × 100 rows corpus (~5M stored list items) from 1.93 GB to 1.26 GB. ([#143](https://github.com/SkyeAv/Tablassert/pull/143)) + ## 16.6.1 - 2026-09-04 ### Performance diff --git a/rust/src/json.rs b/rust/src/json.rs index d9332231..8bd2ed53 100644 --- a/rust/src/json.rs +++ b/rust/src/json.rs @@ -79,30 +79,61 @@ pub fn emitted_json_bytes(value: &Value) -> serde_json::Result> { /// different bytes. The edge deduper compares records for equality, so it needs a form /// where "same content" means "same bytes"; array order is preserved because it is /// semantic. +/// +/// US-003: the sorted-key form is written DIRECTLY into the output buffer instead of +/// first deep-cloning into a canonical `Value` tree -- this runs once per record on +/// every pass, so the clone was the fixed per-record cost. Scalars are delegated to +/// serde_json's own writers, keeping escaping and number formatting byte-identical to +/// the old clone-then-`to_vec` form (pinned by `canonical_json_bytes_matches_reference_on_fuzz`). pub fn canonical_json_bytes(value: &Value) -> serde_json::Result> { - serde_json::to_vec(&canonical_value(value)) + let mut buf: Vec = Vec::with_capacity(128); + write_canonical(value, &mut buf)?; + Ok(buf) } -fn canonical_value(value: &Value) -> Value { +fn write_canonical(value: &Value, buf: &mut Vec) -> serde_json::Result<()> { match value { Value::Object(entries) => { let mut keys: Vec<&String> = entries.keys().collect(); keys.sort_unstable(); - let mut sorted: Map = Map::with_capacity(entries.len()); + buf.push(b'{'); + let mut first: bool = true; for key in keys { - sorted.insert(key.clone(), canonical_value(&entries[key])); + if !first { + buf.push(b','); + } + first = false; + serde_json::to_writer(&mut *buf, key)?; + buf.push(b':'); + write_canonical(entries.get(key).expect("key came from the same map"), buf)?; } - Value::Object(sorted) + buf.push(b'}'); + } + Value::Array(items) => { + buf.push(b'['); + let mut first: bool = true; + for item in items { + if !first { + buf.push(b','); + } + first = false; + write_canonical(item, buf)?; + } + buf.push(b']'); + } + // null / bool / number / string: serde_json's own writer emits exactly the bytes + // `to_vec` would for the same scalar (ryu/itoa numbers, full string escaping). + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => { + serde_json::to_writer(&mut *buf, value)? } - Value::Array(items) => Value::Array(items.iter().map(canonical_value).collect()), - _ => value.clone(), } + Ok(()) } #[cfg(test)] mod tests { - use super::strip_nulls; - use serde_json::json; + use super::{canonical_json_bytes, strip_nulls}; + use serde_json::{json, Map, Value}; #[test] fn strip_nulls_removes_absent_and_null_like_values() { @@ -153,4 +184,271 @@ mod tests { }) ); } + + /// Byte-verbatim copy of the pre-US-003 canonical serializer: deep-clone into a + /// key-sorted `Value` at every depth, then `to_vec` the clone. Kept ONLY as the + /// equivalence oracle for `canonical_json_bytes_matches_reference_on_fuzz` -- + /// production code must never call it (the deep clone was the fixed per-record cost + /// US-003 removed). + fn canonical_json_bytes_reference(value: &Value) -> serde_json::Result> { + serde_json::to_vec(&canonical_value_reference(value)) + } + + fn canonical_value_reference(value: &Value) -> Value { + match value { + Value::Object(entries) => { + let mut keys: Vec<&String> = entries.keys().collect(); + keys.sort_unstable(); + let mut sorted: Map = Map::with_capacity(entries.len()); + for key in keys { + sorted.insert(key.clone(), canonical_value_reference(&entries[key])); + } + Value::Object(sorted) + } + Value::Array(items) => { + Value::Array(items.iter().map(canonical_value_reference).collect()) + } + _ => value.clone(), + } + } + + /// Tiny deterministic PRNG (splitmix64): the fuzz stream must be reproducible across + /// machines without pulling in a `rand` dependency. + struct Rng(u64); + + impl Rng { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z: u64 = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + fn below(&mut self, bound: usize) -> usize { + (self.next_u64() % u64::try_from(bound).expect("bound fits in u64")) as usize + } + + fn one_in(&mut self, odds: usize) -> bool { + self.below(odds) == 0 + } + } + + fn shuffle(rng: &mut Rng, items: &mut [T]) { + for high in (1..items.len()).rev() { + let low: usize = rng.below(high + 1); + items.swap(high, low); + } + } + + /// Strings loaded with escape-sensitive content: quotes, backslashes, control chars, + /// DEL, NEL/line separators, BMP and non-BMP unicode. + const FUZZ_STRINGS: [&str; 14] = [ + "", + "plain", + "two words", + "quote \" inside", + "back\\slash", + "\"both\" and \\ together", + "tab\tnewline\nreturn\r", + "\u{0}\u{1}\u{1f}", + "\u{7f}", + "\u{85}\u{2028}\u{2029}", + "unicode: \u{e9} \u{fc} \u{df} \u{6f22}\u{5b57}", + "non-bmp: \u{1d11e} \u{1f600} \u{1f980}", + "\u{10ffff}", + "trailing backslash\\", + ]; + + /// Keys chosen to stress the byte-wise key sort: digits-vs-letters, case, punctuation + /// above and below letters, spaces, and multi-byte UTF-8. + const FUZZ_KEYS: [&str; 12] = [ + "a", "b", "A", "B", "_", "~", "10", "2", "z", "\u{e9}", "zz", "a b", + ]; + + /// Shapes the fuzz stream actually covered, so a passing run cannot be vacuous. + #[derive(Default)] + struct FuzzStats { + multi_key_objects: usize, + arrays: usize, + nested_containers: usize, + escaped_strings: usize, + } + + fn fuzz_number(rng: &mut Rng) -> Value { + match rng.below(4) { + 0 => Value::Number(rng.next_u64().into()), // u64, incl. above i64::MAX + 1 => Value::Number((rng.next_u64() as i64).into()), // i64, incl. negatives + 2 => { + let pool: [f64; 10] = [ + 0.0, + -0.0, + 0.1, + -0.1, + 1e30, + 1.5e-7, + 0.3333333333333333, + 5e-324, + 1.7976931348623157e308, + 987654321.5, + ]; + Value::Number( + serde_json::Number::from_f64(pool[rng.below(pool.len())]) + .expect("pool values are finite"), + ) + } + _ => { + // Random finite fraction with exponent in [-8, 8]: exercises ryu's + // decimal-vs-exponent formatting boundary and negative signs. + let mantissa: f64 = (rng.next_u64() % 1_000_000) as f64; + let exponent: i32 = rng.below(17) as i32 - 8; + let mut magnitude: f64 = mantissa * 10f64.powi(exponent); + if rng.one_in(2) { + magnitude = -magnitude; + } + Value::Number(serde_json::Number::from_f64(magnitude).expect("finite")) + } + } + } + + fn fuzz_value(rng: &mut Rng, depth: usize, stats: &mut FuzzStats) -> Value { + // Beyond depth 4 only scalars: recursion stays bounded while still nesting deep + // enough to exercise the recursive sorted-key write. + match rng.below(if depth >= 4 { 4 } else { 8 }) { + 0 => Value::Null, + 1 => Value::Bool(rng.one_in(2)), + 2 => fuzz_number(rng), + 3 => { + let text: &str = FUZZ_STRINGS[rng.below(FUZZ_STRINGS.len())]; + if text.chars().any(|c| c == '"' || c == '\\' || c < ' ') { + stats.escaped_strings += 1; + } + Value::String(text.to_string()) + } + 4 => { + stats.arrays += 1; + if depth > 0 { + stats.nested_containers += 1; + } + let mut items: Vec = Vec::new(); + for _ in 0..rng.below(5) { + items.push(fuzz_value(rng, depth + 1, stats)); + } + Value::Array(items) + } + _ => { + let mut entries: Map = Map::new(); + for _ in 0..rng.below(5) { + let key: String = FUZZ_KEYS[rng.below(FUZZ_KEYS.len())].to_string(); + entries.insert(key, fuzz_value(rng, depth + 1, stats)); + } + if entries.len() >= 2 { + stats.multi_key_objects += 1; + } + if depth > 0 { + stats.nested_containers += 1; + } + Value::Object(entries) + } + } + } + + #[test] + fn canonical_json_bytes_matches_reference_on_fuzz() { + // WHY: US-003 rewrote `canonical_json_bytes` from "deep-clone into a canonical + // Value, then to_vec" to "write the sorted-key JSON directly into the output + // buffer". Node dedup (`record_if_new`) keys on these exact bytes and the + // default-mode edge content hash hashes them, so even one drifted byte (string + // escaping, number formatting, key sort, separators) silently shifts node ids, + // dedup decisions, and edge hashes. This seeded fuzz drives the new writer and + // a byte-verbatim copy of the old serializer over nested objects at depth, + // arrays, every number shape (u64/i64/negative/fractional/exponent), bools, + // nulls, escape-heavy unicode strings, and empty containers, and demands + // byte-identical output. + let mut rng = Rng(0x5EED_2024_0000_0003); + let mut stats = FuzzStats::default(); + + for _ in 0..4096 { + let value: Value = fuzz_value(&mut rng, 0, &mut stats); + let direct: Vec = canonical_json_bytes(&value).expect("canonical"); + let reference: Vec = canonical_json_bytes_reference(&value).expect("reference"); + assert_eq!(direct, reference, "divergence on: {value}"); + } + + // Key-order permutations of the same object: canonical bytes must not depend on + // insertion order (serde_json's preserve_order leaks it to plain to_vec). + let mut order: Vec = (0..FUZZ_KEYS.len()).collect(); + for _ in 0..128 { + shuffle(&mut rng, &mut order); + let count: usize = 1 + rng.below(FUZZ_KEYS.len()); + let entries: Vec<(String, Value)> = order[..count] + .iter() + .map(|index| { + ( + FUZZ_KEYS[*index].to_string(), + fuzz_value(&mut rng, 1, &mut stats), + ) + }) + .collect(); + let mut forward: Map = Map::new(); + for (key, entry) in &entries { + forward.insert(key.clone(), entry.clone()); + } + let mut reversed: Map = Map::new(); + for (key, entry) in entries.iter().rev() { + reversed.insert(key.clone(), entry.clone()); + } + let forward_value = Value::Object(forward); + let reversed_value = Value::Object(reversed); + let direct_forward: Vec = canonical_json_bytes(&forward_value).expect("canonical"); + let direct_reversed: Vec = + canonical_json_bytes(&reversed_value).expect("canonical"); + let reference_forward: Vec = + canonical_json_bytes_reference(&forward_value).expect("reference"); + assert_eq!(direct_forward, direct_reversed, "key order leaked"); + assert_eq!( + direct_forward, reference_forward, + "divergence on permutation" + ); + } + + // Duplicate keys in raw JSON: serde_json's parser collapses them last-wins + // before either serializer sees the Value; pin that both implementations agree + // on the collapsed form (plus the degenerate scalars/empties). + let raw_cases: [&str; 8] = [ + r#"{"a":1,"a":2}"#, + r#"{"b":{"x":[1,2],"x":null},"b":{}}"#, + "{}", + "[]", + "null", + "true", + "\"\"", + "-0.0", + ]; + for raw in raw_cases { + let value: Value = serde_json::from_str(raw).expect("parse raw case"); + let direct: Vec = canonical_json_bytes(&value).expect("canonical"); + let reference: Vec = canonical_json_bytes_reference(&value).expect("reference"); + assert_eq!(direct, reference, "raw case: {raw}"); + } + + // Non-vacuity: the seeded stream must actually exercise the risky shapes or the + // equivalence check could pass on a trivially degenerate input. + assert!( + stats.multi_key_objects >= 100, + "expected multi-key objects, got {}", + stats.multi_key_objects + ); + assert!(stats.arrays >= 100, "expected arrays, got {}", stats.arrays); + assert!( + stats.nested_containers >= 50, + "expected nested containers, got {}", + stats.nested_containers + ); + assert!( + stats.escaped_strings >= 50, + "expected escape-heavy strings, got {}", + stats.escaped_strings + ); + } } diff --git a/rust/src/ndjson.rs b/rust/src/ndjson.rs index 9f329386..6e7e9721 100644 --- a/rust/src/ndjson.rs +++ b/rust/src/ndjson.rs @@ -2,12 +2,13 @@ use crate::json::{canonical_json_bytes, emitted_json_bytes, strip_nulls}; use crate::uuid::uuid_for_json_object; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use serde_json::Value; use std::collections::hash_map::Entry; use std::fs::File; use std::io::{BufRead, BufReader, BufWriter, Write}; use std::path::{Path, PathBuf}; +use std::rc::Rc; use uuid::Uuid; use xxhash_rust::xxh64::xxh64; @@ -192,8 +193,8 @@ fn not_a_key_error(output: &Path, id: &str, incoming: &Value, fields: Option<&[S )) } -/// Merge-mode dedup state: derived id -> (hashes of every absorbed record, merged record), -/// plus the first-seen id order so the buffered output is deterministic. +/// Merge-mode dedup state: derived id -> the record that first claimed it plus its fold +/// bookkeeping, plus the first-seen id order so the buffered output is deterministic. /// /// Unlike the default `EdgeIndex` (24 bytes per edge, streaming writes), this retains one /// COMPLETE record per unique id and writes nothing until end-of-stream, because a @@ -201,12 +202,309 @@ fn not_a_key_error(output: &Path, id: &str, incoming: &Value, fields: Option<&[S /// memory cost is exactly why merge mode is opt-in (`uuid_on_collision: merge`). #[derive(Default)] struct MergeIndex { - records: FxHashMap<[u8; 16], (Vec, Value)>, + records: FxHashMap<[u8; 16], MergedRecord>, order: Vec<[u8; 16]>, merged: u64, scalar_conflicts: u64, } +/// The canonical bytes of one list item, shared by every structure that needs them. +/// +/// ONE heap allocation per distinct item instead of one per consumer: `Rc`'s `Hash`, +/// `Eq`, and `Ord` all delegate to `T`, so set membership still compares FULL canonical +/// bytes (never a bare hash -- a hash-only membership set would merge two distinct items +/// on a collision) and the deferred write-out sort still orders by those same bytes. +/// `Rc<[u8]>` rather than `Rc>` because the slice form stores the bytes inline in +/// the reference-count block: one exactly-sized allocation per item instead of a count +/// block plus a `Vec` buffer that keeps `canonical_json_bytes`' 128-byte starting +/// capacity. At scenario-C shapes (~5M stored list items, since every row contributes its +/// own `supporting_case_ids`) that is roughly 64 resident bytes per short item instead of +/// ~320 for two full copies -- measured on the committed harness, scenario C's peak RSS +/// falls from 1.93 GB to 1.26 GB. +type SharedBytes = Rc<[u8]>; + +/// Union bookkeeping for one array field of one buffered record. +/// +/// Canonical bytes are computed ONCE per item -- when the item first enters the record, +/// either with the record itself or appended by a fold -- and are reused for membership +/// and the final sort: each fold costs O(1) per incoming item instead of re-canonicalizing +/// every stored item on every fold (the pre-US-002 quadratic). The two structures below +/// SHARE that one allocation per item (`SharedBytes`), so the bookkeeping costs one +/// canonical-bytes copy per distinct item, not two. +struct ListState { + /// Canonical bytes of every stored item. Membership oracle for INCOMING items only: + /// duplicates inside the first-seen record stay in the list but still reject an equal + /// incoming item, exactly like the former linear `seen` scan. + seen: FxHashSet, + /// Canonical bytes parallel to the live items, so the deferred write-out sort never + /// re-canonicalizes anything. The SAME allocation `seen` holds -- a refcount, not a + /// second copy. + bytes: Vec, + /// Set when a real union ran (BOTH sides carried an array). Only then does the + /// write-out sort the list; a list merely copied from a later record keeps its source + /// order, byte-for-byte with the pre-US-002 fold. + unioned: bool, +} + +impl ListState { + /// Seed from an array that just entered the record (first-seen, or copied from a later + /// record): keep EVERY item -- duplicates included -- and record each item's canonical + /// bytes once. + fn from_items(items: &[Value]) -> PyResult { + let mut state = Self { + seen: FxHashSet::with_capacity_and_hasher(items.len(), Default::default()), + bytes: Vec::with_capacity(items.len()), + unioned: false, + }; + for item in items { + let bytes: SharedBytes = canonical_json_bytes(item).map_err(runtime_error)?.into(); + state.seen.insert(Rc::clone(&bytes)); + state.bytes.push(bytes); + } + Ok(state) + } +} + +/// One buffered edge record plus the O(1) bookkeeping its folds need. +/// +/// Replaces the pre-US-002 `(Vec, Value)` tuple: exact-repeat suppression and list +/// membership are hash sets, and the per-fold re-canonicalization + re-sort of every +/// stored list item is gone -- lists sort ONCE, at write-out, and only if a union ran. +struct MergedRecord { + /// The live record: scalars first-wins, lists in original/append order until the + /// deferred write-out sort. + value: Value, + /// Content hashes of every record absorbed into this id: O(1) exact-repeat + /// suppression. Bare hashes keep the pre-US-002 `Vec` semantics exactly (a false + /// hit could only skip re-folding a record, matching the old `contains`); node dedup + /// stays full-bytes because a collision there would drop a DISTINCT record. + hashes: FxHashSet, + /// Union state for every field whose current value is an array. + lists: FxHashMap, +} + +impl MergedRecord { + /// Buffer the record that first claimed its id, seeding union state for every array + /// field (items kept verbatim, canonical bytes recorded once). + fn new(value: Value, content: u64) -> PyResult { + let mut lists: FxHashMap = FxHashMap::default(); + if let Some(map) = value.as_object() { + for (key, item) in map { + if let Value::Array(items) = item { + lists.insert(key.clone(), ListState::from_items(items)?); + } + } + } + let mut hashes: FxHashSet = FxHashSet::default(); + hashes.insert(content); + Ok(Self { + value, + hashes, + lists, + }) + } + + /// Apply the one deferred sort just before the write: every list that went through a + /// real union is sorted by canonical bytes (stable, so equal-byte items keep arrival + /// order -- byte-identical to the pre-US-002 per-fold sort). Lists that were only + /// copied keep their source order. Returns the finished record. + fn finish(mut self) -> PyResult { + let Some(map) = self.value.as_object_mut() else { + return Err(runtime_error("expected JSON object")); + }; + for (key, state) in &mut self.lists { + if !state.unioned { + continue; + } + // The `number_of_cases` recompute may have replaced a unioned array with a + // number after the union; only arrays are sortable. + let Some(Value::Array(items)) = map.get_mut(key) else { + continue; + }; + let bytes: Vec = std::mem::take(&mut state.bytes); + // Fail loudly on a desync instead of silently corrupting the record: `zip` + // truncates to the shorter side while `drain(..)` empties the WHOLE array, so + // ANY length mismatch drops entries without a trace -- live items when `bytes` + // is short, canonical-byte entries when `items` is short. The invariant + // (`bytes` is parallel to the live items) holds by construction -- every push + // into one is paired with a push into the other -- but this repo's standard is + // fail-loudly: hash-only keying once silently dropped records the same way. + // Deliberately a plain runtime check, not a `debug_assert_eq!`: a debug assert + // would panic first in test/debug builds, so the structured error below could + // never be observed or tested there. This fires in EVERY build profile. + if bytes.len() != items.len() { + return Err(runtime_error(format!( + "merge-state-desync: list field {key:?} carries {} canonical-byte entries \ + for {} live items; refusing to write the record, because pairing them would \ + silently truncate {} entry/entries", + bytes.len(), + items.len(), + items.len().abs_diff(bytes.len()) + ))); + } + let mut keyed: Vec<(SharedBytes, Value)> = + bytes.into_iter().zip(items.drain(..)).collect(); + keyed.sort_by(|left, right| left.0.cmp(&right.0)); + items.extend(keyed.into_iter().map(|(_, item)| item)); + } + Ok(self.value) + } +} + +/// Fold `incoming` into `stored`, field-wise. Returns the number of conflicting scalar +/// fields (kept first-wins) so the caller can report them. +/// +/// Near-linear implementation of the semantics frozen by `merge_records_reference` and +/// policed by `merge_fold_matches_reference_on_fuzz`: +/// +/// - list fields: union, deduped by canonical JSON bytes (so two `sources` objects that +/// differ only in key order collapse). Each incoming item is canonicalized ONCE and +/// checked against the field's `ListState` set in O(1); stored items are never +/// re-canonicalized, and the sort by canonical bytes is deferred to write-out +/// (`MergedRecord::finish`), where it runs only for fields that saw a real union; +/// - scalar fields: first-wins on conflict, counted; +/// - fields only on `incoming`: copied over (only a conflict when both sides disagree); +/// - `id` is never touched: both sides carry the same one by construction. +/// +/// `number_of_cases` has one hardcoded exception to first-wins: when the MERGED record +/// carries `supporting_case_ids` (a build-internal `list[str]` of the case IDs behind +/// the count -- allowed onto edge frames so it reaches this pass, then stripped before +/// write) and either side carried a count, the count is recomputed as the length of the +/// unioned ID list. A case ID shared by both records is one case, so first-wins and +/// summing both over- and under-report; the union length is the exact count. The +/// superseded divergence is NOT reported as a scalar conflict. When neither side +/// carries the list, `number_of_cases` stays an ordinary first-wins scalar. +fn merge_records(stored: &mut MergedRecord, incoming: &Value) -> PyResult { + let Some(incoming_map) = incoming.as_object() else { + return Err(runtime_error("expected JSON object")); + }; + let mut conflicts: u64 = 0; + // Read both counts BEFORE the fold: the fold may copy incoming's over a stored side + // that lacks it, and the recompute rule below needs to know each side contributed one. + let stored_cases: Option = stored.value.get("number_of_cases").cloned(); + let incoming_cases: Option = incoming_map.get("number_of_cases").cloned(); + let Some(stored_map) = stored.value.as_object_mut() else { + return Err(runtime_error("expected JSON object")); + }; + for (key, incoming_value) in incoming_map { + if key == "id" { + continue; + } + match stored_map.get_mut(key) { + None => { + // First time this field appears on the stored record: copy it over. An + // array gets union state seeded so LATER folds can union into it in O(1) + // (this copy itself is NOT a union: it keeps its source order). + if let Value::Array(items) = incoming_value { + stored + .lists + .insert(key.clone(), ListState::from_items(items)?); + } + stored_map.insert(key.clone(), incoming_value.clone()); + } + Some(stored_value) => { + if let (Value::Array(stored_items), Value::Array(incoming_items)) = + (&mut *stored_value, incoming_value) + { + let Some(state) = stored.lists.get_mut(key) else { + // Invariant: every stored array field carries union state. + return Err(runtime_error("list field without union state")); + }; + // BOTH sides arrays -> a union happened: this field sorts at write-out + // even when no new item survives membership (the pre-US-002 fold + // sorted on every such fold). + state.unioned = true; + for item in incoming_items { + let bytes: SharedBytes = + canonical_json_bytes(item).map_err(runtime_error)?.into(); + if state.seen.insert(Rc::clone(&bytes)) { + state.bytes.push(bytes); + stored_items.push(item.clone()); + } + } + } else if stored_value != incoming_value { + conflicts += 1; + } + } + } + } + // WHY: exact-unique `number_of_cases` semantics (see the docstring). Guarded on the + // merged record actually carrying the ID list: a one-sided carrier is fine (the union + // is just that side's list), while a carrier-less merge never recomputes. + if let Some(union_len) = stored + .value + .get("supporting_case_ids") + .and_then(Value::as_array) + .map(Vec::len) + { + if stored_cases.is_some() || incoming_cases.is_some() { + // Undo the loop's conflict count when it fired on this very field: the + // recompute supersedes first-wins, so the divergence is not a conflict. + // Mirrors the loop's condition exactly -- both sides present, unequal, and + // not both arrays (two arrays took the union path and were never counted). + if let (Some(left), Some(right)) = (&stored_cases, &incoming_cases) { + if left != right && !(left.is_array() && right.is_array()) { + conflicts -= 1; + } + } + let Some(map) = stored.value.as_object_mut() else { + return Err(runtime_error("expected JSON object")); + }; + map.insert( + "number_of_cases".to_string(), + Value::Number(serde_json::Number::from(union_len)), + ); + } + } + Ok(conflicts) +} + +/// Remove build-internal carrier fields from an edge record before it is written. +/// +/// `supporting_case_ids` exists only so merge mode can recompute `number_of_cases` +/// (see `merge_records`); it must never ship in the final NDJSON, so EVERY edge write +/// path drops it -- the buffering merge pass and the default streaming path alike. +/// Nodes never carry it and are untouched. +fn strip_internal_edge_fields(value: &mut Value) { + if let Some(map) = value.as_object_mut() { + map.remove("supporting_case_ids"); + } +} + +impl MergeIndex { + fn absorb(&mut self, id: [u8; 16], value: Value, content: u64) -> PyResult<()> { + match self.records.entry(id) { + Entry::Vacant(slot) => { + slot.insert(MergedRecord::new(value, content)?); + self.order.push(id); + Ok(()) + } + Entry::Occupied(mut slot) => { + let record: &mut MergedRecord = slot.get_mut(); + // Exact repeat of ANY record already folded into this id -- including a + // divergent one -- is suppressed, so the conflict summary never + // double-counts a re-seen row. + if record.hashes.contains(&content) { + return Ok(()); + } + self.merged += 1; + self.scalar_conflicts += merge_records(record, &value)?; + record.hashes.insert(content); + Ok(()) + } + } + } +} + +/// FROZEN EQUIVALENCE ORACLE -- BYTE-VERBATIM copy of the pre-US-002 `merge_records` +/// fold, compiled ONLY under `cfg(test)`. US-002 may rewrite the production fold for +/// speed, but this copy stays the exact algorithm of record: `merge_fold_matches_ +/// reference_on_fuzz` drives both and demands identical output and counters. Never edit +/// this copy to match an optimization -- edit the production code and let the fuzz test +/// arbitrate. See the module docs on `merge_fold_reference` below. +/// +/// Original semantics doc: +/// /// Fold `incoming` into `stored`, field-wise. Returns the number of conflicting scalar /// fields (kept first-wins) so the caller can report them. /// @@ -225,7 +523,8 @@ struct MergeIndex { /// summing both over- and under-report; the union length is the exact count. The /// superseded divergence is NOT reported as a scalar conflict. When neither side /// carries the list, `number_of_cases` stays an ordinary first-wins scalar. -fn merge_records(stored: &mut Value, incoming: &Value) -> PyResult { +#[cfg(test)] +fn merge_records_reference(stored: &mut Value, incoming: &Value) -> PyResult { let Some(incoming_map) = incoming.as_object() else { return Err(runtime_error("expected JSON object")); }; @@ -302,19 +601,20 @@ fn merge_records(stored: &mut Value, incoming: &Value) -> PyResult { Ok(conflicts) } -/// Remove build-internal carrier fields from an edge record before it is written. -/// -/// `supporting_case_ids` exists only so merge mode can recompute `number_of_cases` -/// (see `merge_records`); it must never ship in the final NDJSON, so EVERY edge write -/// path drops it -- the buffering merge pass and the default streaming path alike. -/// Nodes never carry it and are untouched. -fn strip_internal_edge_fields(value: &mut Value) { - if let Some(map) = value.as_object_mut() { - map.remove("supporting_case_ids"); - } +/// FROZEN EQUIVALENCE ORACLE -- BYTE-VERBATIM copy of the pre-US-002 `MergeIndex` absorb +/// path (state shape included), compiled ONLY under `cfg(test)`; drives +/// `merge_records_reference`. Treat as read-only -- see that fn's docs. +#[cfg(test)] +#[derive(Default)] +struct MergeIndexReference { + records: FxHashMap<[u8; 16], (Vec, Value)>, + order: Vec<[u8; 16]>, + merged: u64, + scalar_conflicts: u64, } -impl MergeIndex { +#[cfg(test)] +impl MergeIndexReference { fn absorb(&mut self, id: [u8; 16], value: Value, content: u64) -> PyResult<()> { match self.records.entry(id) { Entry::Vacant(slot) => { @@ -331,7 +631,7 @@ impl MergeIndex { return Ok(()); } self.merged += 1; - self.scalar_conflicts += merge_records(stored_value, &value)?; + self.scalar_conflicts += merge_records_reference(stored_value, &value)?; hashes.push(content); Ok(()) } @@ -364,9 +664,11 @@ fn dedup_edges_merge( index.absorb(edge_id_bytes(&value)?, value, content)?; } for id in std::mem::take(&mut index.order) { - let Some((_, mut value)) = index.records.remove(&id) else { + let Some(record) = index.records.remove(&id) else { continue; }; + // The one deferred sort: unioned lists sort here, everything else ships as-is. + let mut value: Value = record.finish()?; strip_internal_edge_fields(&mut value); writer .write_all(&emitted_json_bytes(&value).map_err(runtime_error)?) @@ -1224,4 +1526,817 @@ mod tests { .expect_err("unknown mode"); assert!(error.to_string().contains("bogus"), "{error}"); } + + #[test] + fn merge_mode_sorts_only_lists_that_went_through_a_union() { + // WHY: the US-002 fold defers sorting to write-out and sorts ONLY fields that saw + // a real union (both sides arrays). A list copied from a later record is evidence + // in its source order and must NOT be sorted -- byte parity with the pre-US-002 + // fold, pinned here directly in addition to the fuzz oracle. + let dir = tempdir().expect("tempdir"); + let (input, output) = write_edges( + dir.path(), + concat!( + // Group 1 (subject A): first record has no `tags`; the second copies its + // deliberately unsorted list in -- never unioned, so it stays unsorted. + "{\"subject\":\"A\",\"predicate\":\"r\",\"object\":\"B\"}\n", + "{\"subject\":\"A\",\"predicate\":\"r\",\"object\":\"B\",\"tags\":[\"z\",\"m\"]}\n", + // Group 2 (subject X): both records carry `tags`, so a real union runs + // and the write-out sorts. + "{\"subject\":\"X\",\"predicate\":\"r\",\"object\":\"Y\",\"tags\":[\"z\",\"m\"]}\n", + "{\"subject\":\"X\",\"predicate\":\"r\",\"object\":\"Y\",\"tags\":[\"a\"]}\n" + ), + ); + let (merged, conflicts) = dedup_ndjson( + input, + output.clone(), + true, + None, + spo_fields(), + Some("merge".to_string()), + ) + .expect("merge dedup"); + + assert_eq!((merged, conflicts), (2, 0)); + let lines: Vec = fs::read_to_string(output) + .expect("read output") + .lines() + .map(|line| serde_json::from_str(line).expect("json")) + .collect(); + assert_eq!(lines.len(), 2); + // Copied, never unioned: source order preserved. + assert_eq!(lines[0]["tags"], serde_json::json!(["z", "m"])); + // Unioned: sorted by canonical bytes. + assert_eq!(lines[1]["tags"], serde_json::json!(["a", "m", "z"])); + } +} + +/// Seeded fuzz equivalence gate for the merge-mode fold. +/// +/// Drives the CURRENT production path (through the real `dedup_ndjson` entry point) and +/// the frozen `merge_records_reference` / `MergeIndexReference` oracle with the SAME +/// reproducible randomized stream and demands byte-identical output, identical output +/// order, and identical `(merged, scalar_conflicts)` counters. Trivially green while the +/// reference IS the current algorithm; it becomes the tripwire for the US-002 rewrite. +#[cfg(test)] +mod merge_fold_reference { + use super::{ + dedup_ndjson, edge_id_bytes, finalize_record, runtime_error, strip_internal_edge_fields, + Finalized, MergeIndexReference, + }; + use crate::json::{canonical_json_bytes, emitted_json_bytes}; + use pyo3::prelude::*; + use serde_json::Value; + use std::fs; + use tempfile::tempdir; + + /// The same pass as `dedup_edges_merge`, but every fold runs through the frozen + /// reference oracle instead of the production `MergeIndex`. + fn reference_pipeline( + lines: &[String], + domain: &str, + fields: &[String], + ) -> PyResult<(Vec, u64, u64)> { + let mut index: MergeIndexReference = MergeIndexReference::default(); + for line in lines { + if line.trim().is_empty() { + continue; + } + let value: Value = serde_json::from_str(line).map_err(runtime_error)?; + let Some(Finalized { value, content }) = + finalize_record(value, true, domain, Some(fields))? + else { + continue; + }; + index.absorb(edge_id_bytes(&value)?, value, content)?; + } + let mut output: Vec = Vec::new(); + for id in std::mem::take(&mut index.order) { + let Some((_, mut value)) = index.records.remove(&id) else { + continue; + }; + strip_internal_edge_fields(&mut value); + output.extend_from_slice(&emitted_json_bytes(&value).map_err(runtime_error)?); + output.push(b'\n'); + } + Ok((output, index.merged, index.scalar_conflicts)) + } + + /// Tiny deterministic PRNG (splitmix64): the fuzz stream must be reproducible across + /// machines without pulling in a `rand` dependency. + struct Rng(u64); + + impl Rng { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z: u64 = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + fn below(&mut self, bound: usize) -> usize { + (self.next_u64() % u64::try_from(bound).expect("bound fits in u64")) as usize + } + + fn one_in(&mut self, odds: usize) -> bool { + self.below(odds) == 0 + } + } + + fn shuffle(rng: &mut Rng, items: &mut [T]) { + for high in (1..items.len()).rev() { + let low: usize = rng.below(high + 1); + items.swap(high, low); + } + } + + const SUBJECTS: [&str; 4] = ["MONDO:1", "MONDO:2", "MONDO:3", "MONDO:4"]; + const PREDICATES: [&str; 2] = ["biolink:related_to", "biolink:associated_with"]; + const OBJECTS: [&str; 3] = ["NCBIGene:1", "NCBIGene:2", "NCBIGene:3"]; + const TAGS: [&str; 5] = ["tag:a", "tag:b", "tag:c", "tag:d", "tag:e"]; + const CASES: [&str; 6] = ["case:1", "case:2", "case:3", "case:4", "case:5", "case:6"]; + const P_VALUES: [&str; 3] = ["0.01", "0.05", "0.99"]; + + /// One of two logical `sources` entries, emitted in a random key order: the fold must + /// collapse key-order variants of the same object via canonical bytes. + fn source_object(rng: &mut Rng) -> Value { + let (id, role): (&str, &str) = if rng.one_in(2) { + ("infores:one", "primary_knowledge_source") + } else { + ("infores:two", "aggregator_knowledge_source") + }; + let mut map = serde_json::Map::new(); + if rng.one_in(2) { + map.insert("resource_id".to_string(), Value::String(id.to_string())); + map.insert("resource_role".to_string(), Value::String(role.to_string())); + } else { + map.insert("resource_role".to_string(), Value::String(role.to_string())); + map.insert("resource_id".to_string(), Value::String(id.to_string())); + } + Value::Object(map) + } + + fn fuzz_record(rng: &mut Rng, group: usize, record_index: usize) -> Value { + let mut map = serde_json::Map::new(); + // The identity triple, inserted in a random key order: all records of a group + // derive the same id while their bytes diverge, so the fold decides the outcome. + let triple: [(&str, &str); 3] = [ + ("subject", SUBJECTS[rng.below(SUBJECTS.len())]), + ("predicate", PREDICATES[rng.below(PREDICATES.len())]), + ("object", OBJECTS[rng.below(OBJECTS.len())]), + ]; + let mut order: [usize; 3] = [0, 1, 2]; + shuffle(rng, &mut order); + for index in order { + let (key, value) = triple[index]; + map.insert(key.to_string(), Value::String(value.to_string())); + } + // Scalar fields drawn from small pools -> first-wins conflicts. + map.insert( + "p_value".to_string(), + Value::String(P_VALUES[rng.below(P_VALUES.len())].to_string()), + ); + map.insert( + "effect_size".to_string(), + Value::Number(serde_json::Number::from(rng.below(4) as u64)), + ); + if rng.one_in(4) { + map.insert("negated".to_string(), Value::Bool(rng.one_in(2))); + } + // String list field drawn WITH replacement: a record may repeat an item inside its + // own array (stored-side duplicates survive; incoming-side ones collapse). + if rng.one_in(2) { + let mut tags: Vec = Vec::new(); + for _ in 0..rng.below(4) { + tags.push(Value::String(TAGS[rng.below(TAGS.len())].to_string())); + } + map.insert("tags".to_string(), Value::Array(tags)); + } + // Object list field with key-order variants -> canonical-bytes union + sort. + if rng.one_in(2) { + let mut sources: Vec = Vec::new(); + for _ in 0..1 + rng.below(2) { + sources.push(source_object(rng)); + } + map.insert("sources".to_string(), Value::Array(sources)); + } + // Scalar-vs-array conflict on one field. + if rng.one_in(3) { + let mode: Value = if rng.one_in(2) { + Value::String("solo".to_string()) + } else { + serde_json::json!(["solo", "extra"]) + }; + map.insert("mode".to_string(), mode); + } + // A field only LATER records of the group carry (fold must copy it, not conflict). + if record_index > 0 && (group.is_multiple_of(3) || rng.one_in(2)) { + map.insert("late".to_string(), Value::String(format!("late:{group}"))); + } + // The same "only a later record carries it" shape, but LIST-valued and on exactly + // ONE record per group: the field is COPIED into the stored record and never + // unioned, so `unioned` stays false and the deferred write-out sort must leave it + // in its SOURCE order. The items are drawn strictly DESCENDING by canonical bytes + // (`TAGS` reversed), which is what makes a stray sort observable -- + // `merge_fold_matches_reference_on_fuzz` counts the descending survivors. Before + // US-006 no fuzz dataset produced this shape at all, so the `unioned == false` gate + // (the subtlest semantic in the fold) rested on one hand-written test. + if group.is_multiple_of(2) && record_index == 1 { + let descending: Vec = TAGS + .iter() + .rev() + .take(2 + rng.below(3)) + .map(|tag| Value::String(tag.to_string())) + .collect(); + map.insert("late_list".to_string(), Value::Array(descending)); + } + // The `number_of_cases` carrier pair in its three shapes: count + list, list only, + // or absent. The count is deliberately wrong sometimes -- the recompute supersedes + // it and must also undo the scalar conflict it would otherwise have counted. + if rng.one_in(2) { + let mut case_ids: Vec = Vec::new(); + for _ in 0..1 + rng.below(4) { + case_ids.push(Value::String(CASES[rng.below(CASES.len())].to_string())); + } + map.insert( + "supporting_case_ids".to_string(), + Value::Array(case_ids.clone()), + ); + if !rng.one_in(4) { + let count: u64 = if rng.one_in(3) { + case_ids.len() as u64 + } else { + (case_ids.len() + 1 + rng.below(3)) as u64 + }; + map.insert( + "number_of_cases".to_string(), + Value::Number(serde_json::Number::from(count)), + ); + } + } + Value::Object(map) + } + + fn fuzz_stream(rng: &mut Rng) -> Vec { + // 24 groups of 1-8 divergent same-id records, arrival order shuffled so first-seen + // id order and fold order disagree. + let mut records: Vec = Vec::new(); + for group in 0..24 { + for record_index in 0..1 + rng.below(8) { + records.push(fuzz_record(rng, group, record_index)); + } + } + shuffle(rng, &mut records); + let mut lines: Vec = records + .iter() + .map(|record| serde_json::to_string(record).expect("serialize fuzz record")) + .collect(); + // Exact byte repeats -- including of records that already diverged -- which the + // content-hash membership must suppress without merging or counting. + let unique: usize = lines.len(); + for _ in 0..12 { + lines.push(lines[rng.below(unique)].clone()); + } + shuffle(rng, &mut lines); + // Blank lines and empty objects are legal stream noise the pass must skip. + let mut stream: Vec = Vec::new(); + for line in lines { + if rng.one_in(14) { + stream.push(String::new()); + } + if rng.one_in(20) { + stream.push(" ".to_string()); + } + if rng.one_in(16) { + stream.push("{}".to_string()); + } + stream.push(line); + } + stream + } + + /// Count the emitted records that carry `late_list`, and how many of those kept it in + /// strictly DESCENDING canonical order -- proof that a list merely COPIED from one + /// record was never sorted at write-out (the `unioned == false` gate). + fn late_list_order(output: &[u8]) -> PyResult<(usize, usize)> { + let mut carriers: usize = 0; + let mut preserved: usize = 0; + for line in String::from_utf8_lossy(output).lines() { + if line.trim().is_empty() { + continue; + } + let value: Value = serde_json::from_str(line).map_err(runtime_error)?; + let Some(items) = value.get("late_list").and_then(Value::as_array) else { + continue; + }; + carriers += 1; + let mut bytes: Vec> = Vec::with_capacity(items.len()); + for item in items { + bytes.push(canonical_json_bytes(item).map_err(runtime_error)?); + } + if bytes.len() >= 2 && bytes.windows(2).all(|pair| pair[0] > pair[1]) { + preserved += 1; + } + } + Ok((carriers, preserved)) + } + + #[test] + fn merge_fold_matches_reference_on_fuzz() { + // WHY: US-002 will rewrite the merge fold for speed. This seeded fuzz drives the + // CURRENT fold through the real `dedup_ndjson` entry point and the frozen + // `merge_records_reference` oracle with the SAME randomized stream -- divergent + // same-id groups of varying sizes, list unions over object and scalar items, + // key-order variants, fields only later records carry, scalar-vs-array conflicts, + // exact byte repeats, empty objects, and blank lines -- and demands byte-identical + // records in identical order plus identical `(merged, scalar_conflicts)` counters. + let mut rng = Rng(0x5EED_2024_0000_0001); + let lines: Vec = fuzz_stream(&mut rng); + let domain: String = "infores:multiomicskg".to_string(); + let fields: Vec = ["subject", "predicate", "object"] + .iter() + .map(ToString::to_string) + .collect(); + + let dir = tempdir().expect("tempdir"); + let input = dir.path().join("fuzz.ndjson.tmp"); + let output = dir.path().join("fuzz.ndjson"); + fs::write(&input, lines.join("\n") + "\n").expect("write fuzz input"); + let current: (u64, u64) = dedup_ndjson( + input, + output.clone(), + true, + Some(domain.clone()), + Some(fields.clone()), + Some("merge".to_string()), + ) + .expect("merge dedup"); + let current_bytes: Vec = fs::read(&output).expect("read merged output"); + + let (reference_bytes, reference_merged, reference_conflicts): (Vec, u64, u64) = + reference_pipeline(&lines, &domain, &fields).expect("reference pipeline"); + + // Non-vacuity: the seeded stream must actually exercise the fold -- records merge, + // scalars conflict, several ids survive, and the carrier never leaks -- or the + // equivalence check could pass on a trivially degenerate input. + let output_lines: usize = current_bytes.iter().filter(|byte| **byte == b'\n').count(); + assert!( + current.0 >= 10, + "expected real folding, merged={}", + current.0 + ); + assert!( + current.1 >= 1, + "expected scalar conflicts, got {}", + current.1 + ); + assert!( + output_lines >= 8, + "expected distinct ids, got {output_lines}" + ); + assert!( + !current_bytes + .windows("supporting_case_ids".len()) + .any(|window| window == b"supporting_case_ids"), + "build-internal carrier leaked into the merged output" + ); + + assert_eq!( + current, + (reference_merged, reference_conflicts), + "counters diverged from the frozen reference" + ); + assert_eq!( + current_bytes, reference_bytes, + "merged output diverged from the frozen reference" + ); + + // The copied-never-unioned gate, asserted on the shape `fuzz_record` was extended + // to emit (US-006): `late_list` arrives strictly descending, so any write-out that + // sorted a merely COPIED list would flip it ascending and drop this count. It is a + // lower bound, not an exact count, because the identity triple is drawn per RECORD, + // so two carriers can land on one derived id -- those legitimately union and sort + // (and both implementations still agree, checked above). Zero would mean the + // unsorted-copy shape never reached the output at all. + let (carriers, preserved): (usize, usize) = + late_list_order(¤t_bytes).expect("parse merged output"); + println!("copied-never-unioned `late_list`: {carriers} emitted carriers, {preserved} kept their source (descending) order"); + assert!( + preserved >= 3, + "expected the copied, never-unioned `late_list` to keep its source (descending) \ + order on at least 3 emitted records, got {preserved} of {carriers}" + ); + } +} + +/// Performance bound gate for the merge-mode fold. +/// +/// US-002 rewrote the merge fold from quadratic to near-linear. Equivalence with the +/// frozen pre-US-002 oracle is policed by `merge_fold_reference` above; this module +/// polices the SPEED: it drives the production `MergeIndex` and the frozen +/// `MergeIndexReference` through the SAME seeded quadratic-shaped workload in-process +/// and demands the production fold win by a fixed ratio. A ratio -- not an absolute +/// wall-clock limit -- is machine-independent, because both legs share the same core, +/// allocator, and input. +/// +/// The module is deliberately self-contained (own splitmix64 `Rng`, own generators): +/// the frozen equivalence harness above must stay byte-verbatim. +#[cfg(test)] +mod merge_fold_speedup { + use super::{runtime_error, strip_internal_edge_fields, MergeIndex, MergeIndexReference}; + use crate::json::{canonical_json_bytes, emitted_json_bytes}; + use pyo3::prelude::*; + use serde_json::Value; + use std::time::{Duration, Instant}; + use uuid::Uuid; + use xxhash_rust::xxh64::xxh64; + + /// Workload shape: `GROUPS` ids, each accumulating `RECORDS_PER_GROUP` divergent + /// records whose unioned lists grow into the hundreds of items. Sized so the frozen + /// quadratic leg takes ~1-5s in debug builds and the near-linear leg milliseconds; + /// the whole test stays in single-digit seconds. + const GROUPS: usize = 16; + const RECORDS_PER_GROUP: usize = 280; + /// The near-linear fold must beat the frozen quadratic oracle by at least this + /// factor. The observed margin is far larger (the reference re-canonicalizes and + /// re-sorts EVERY stored list item on EVERY fold), so the bound leaves ample headroom + /// for machine noise while still tripping on any regression back toward quadratic. + const BOUND: f64 = 5.0; + /// Timing attempts for the NEW-fold leg (the reference leg stays single-shot). The new + /// fold finishes in milliseconds, where one shot measures mostly scheduler and + /// allocator noise; the MINIMUM over three attempts is the stable estimate of the leg's + /// own cost, because noise only ever ADDS time. + const NEW_FOLD_ATTEMPTS: usize = 3; + const WORKLOAD_SEED: u64 = 0x5EED_2024_0000_0003; + const WARMUP_SEED: u64 = 0x5EED_2024_0000_0004; + + const SUBJECTS: [&str; 4] = ["MONDO:1", "MONDO:2", "MONDO:3", "MONDO:4"]; + const PREDICATES: [&str; 2] = ["biolink:related_to", "biolink:associated_with"]; + const OBJECTS: [&str; 3] = ["NCBIGene:1", "NCBIGene:2", "NCBIGene:3"]; + const P_VALUES: [&str; 3] = ["0.01", "0.05", "0.99"]; + const CASES: [&str; 6] = ["case:1", "case:2", "case:3", "case:4", "case:5", "case:6"]; + + /// Tiny deterministic PRNG (splitmix64): a COPY of `merge_fold_reference::Rng` -- + /// sharing it would mean editing the frozen harness, so the copy is deliberate. + struct Rng(u64); + + impl Rng { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z: u64 = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + fn below(&mut self, bound: usize) -> usize { + (self.next_u64() % u64::try_from(bound).expect("bound fits in u64")) as usize + } + + fn one_in(&mut self, odds: usize) -> bool { + self.below(odds) == 0 + } + } + + fn shuffle(rng: &mut Rng, items: &mut [T]) { + for high in (1..items.len()).rev() { + let low: usize = rng.below(high + 1); + items.swap(high, low); + } + } + + /// One of many logical `sources` entries, emitted in a random key order: the union + /// must collapse key-order variants of the same object via canonical bytes. + fn source_variant(rng: &mut Rng, identity: usize) -> Value { + let resource_id: String = format!("infores:bulk{identity}"); + let role: &str = if identity.is_multiple_of(2) { + "primary_knowledge_source" + } else { + "aggregator_knowledge_source" + }; + let mut map = serde_json::Map::new(); + if rng.one_in(2) { + map.insert("resource_id".to_string(), Value::String(resource_id)); + map.insert("resource_role".to_string(), Value::String(role.to_string())); + } else { + map.insert("resource_role".to_string(), Value::String(role.to_string())); + map.insert("resource_id".to_string(), Value::String(resource_id)); + } + Value::Object(map) + } + + /// One record ready for `absorb`: the derived id bytes, the labeled record, and the + /// xxh64 of its canonical id-free content -- exactly what the production + /// `dedup_edges_merge` path computes before each fold, minus the file IO and parsing + /// that are identical for both legs and would only dilute the ratio. + #[derive(Clone)] + struct FoldCase { + id: [u8; 16], + record: Value, + content: u64, + } + + /// The quadratic-shaped workload: many divergent records per id, unioned list fields + /// growing into the hundreds of items (string items, object items with key-order + /// variants), scalar conflicts, scalar-vs-array conflicts, the `number_of_cases` + /// carrier, exact byte repeats, and shuffled arrival order. Deterministic in `seed`, + /// so two calls with the same arguments yield identical streams. + fn fold_workload( + seed: u64, + groups: usize, + records_per_group: usize, + ) -> PyResult> { + let mut rng = Rng(seed); + // A tag pool large enough that intra-record dedup is rare: each id's unioned + // `tags` list grows near-linearly toward the hundreds of items. + let tag_pool: Vec = (0..records_per_group * 4) + .map(|index| format!("tag:{index}")) + .collect(); + let mut cases: Vec = Vec::with_capacity(groups * records_per_group); + for group in 0..groups { + let mut id_bytes = [0u8; 16]; + for byte in &mut id_bytes { + *byte = rng.below(256) as u8; + } + let subject: &str = SUBJECTS[rng.below(SUBJECTS.len())]; + let predicate: &str = PREDICATES[rng.below(PREDICATES.len())]; + let object: &str = OBJECTS[rng.below(OBJECTS.len())]; + for record_index in 0..records_per_group { + let mut map = serde_json::Map::new(); + // The identity triple is constant within a group, so every record of the + // group derives the same id and the fold decides the outcome. + map.insert("subject".to_string(), Value::String(subject.to_string())); + map.insert( + "predicate".to_string(), + Value::String(predicate.to_string()), + ); + map.insert("object".to_string(), Value::String(object.to_string())); + // Scalars from small pools -> first-wins conflicts on most folds. + map.insert( + "p_value".to_string(), + Value::String(P_VALUES[rng.below(P_VALUES.len())].to_string()), + ); + map.insert( + "effect_size".to_string(), + Value::Number(serde_json::Number::from(rng.below(4) as u64)), + ); + if rng.one_in(4) { + map.insert("negated".to_string(), Value::Bool(rng.one_in(2))); + } + // The growing string list: 2-5 items from the large pool per record. + let mut tags: Vec = Vec::new(); + for _ in 0..2 + rng.below(4) { + tags.push(Value::String(tag_pool[rng.below(tag_pool.len())].clone())); + } + map.insert("tags".to_string(), Value::Array(tags)); + // The growing object list with key-order variants: a mix of new items and + // canonical-byte duplicates of earlier ones. + if rng.one_in(2) { + let mut sources: Vec = Vec::new(); + for _ in 0..1 + rng.below(3) { + let identity: usize = rng.below(records_per_group); + sources.push(source_variant(&mut rng, identity)); + } + map.insert("sources".to_string(), Value::Array(sources)); + } + // Scalar-vs-array conflict on one field. + if rng.one_in(3) { + let mode: Value = if rng.one_in(2) { + Value::String("solo".to_string()) + } else { + serde_json::json!(["solo", "extra"]) + }; + map.insert("mode".to_string(), mode); + } + // A field only LATER records carry (fold copies it, no conflict). + if record_index > 0 && rng.one_in(2) { + map.insert("late".to_string(), Value::String(format!("late:{group}"))); + } + // The `number_of_cases` carrier pair, count deliberately wrong sometimes. + if rng.one_in(2) { + let mut case_ids: Vec = Vec::new(); + for _ in 0..1 + rng.below(5) { + case_ids.push(Value::String(CASES[rng.below(CASES.len())].to_string())); + } + map.insert( + "supporting_case_ids".to_string(), + Value::Array(case_ids.clone()), + ); + if !rng.one_in(4) { + let count: u64 = case_ids.len() as u64 + + if rng.one_in(3) { + 0 + } else { + 1 + rng.below(3) as u64 + }; + map.insert( + "number_of_cases".to_string(), + Value::Number(serde_json::Number::from(count)), + ); + } + } + let mut record: Value = Value::Object(map); + // Content hashes the canonical id-free record exactly like + // `finalize_record`, so an exact byte repeat carries the same content. + let content: u64 = xxh64(&canonical_json_bytes(&record).map_err(runtime_error)?, 0); + record.as_object_mut().expect("record is an object").insert( + "id".to_string(), + Value::String(Uuid::from_bytes(id_bytes).to_string()), + ); + cases.push(FoldCase { + id: id_bytes, + record, + content, + }); + } + } + // Exact byte repeats (~10%): content-hash suppression must keep them out of BOTH + // folds without counting a merge. + let unique: usize = cases.len(); + for _ in 0..unique / 10 { + cases.push(cases[rng.below(unique)].clone()); + } + // Arrival order shuffled so first-seen id order and fold order disagree. + shuffle(&mut rng, &mut cases); + Ok(cases) + } + + /// Drive the production fold over a prepared workload: absorb every record, then + /// `finish` in first-seen order (the one deferred union sort) -- exactly + /// `dedup_edges_merge` minus the file IO, parsing, and finalization that are + /// identical for both legs. + fn run_merge_index(cases: Vec) -> PyResult<(Vec, u64, u64)> { + let mut index: MergeIndex = MergeIndex::default(); + for case in cases { + index.absorb(case.id, case.record, case.content)?; + } + let mut output: Vec = Vec::new(); + for id in std::mem::take(&mut index.order) { + let Some(record) = index.records.remove(&id) else { + continue; + }; + let mut value: Value = record.finish()?; + strip_internal_edge_fields(&mut value); + output.extend_from_slice(&emitted_json_bytes(&value).map_err(runtime_error)?); + output.push(b'\n'); + } + Ok((output, index.merged, index.scalar_conflicts)) + } + + /// Drive the frozen quadratic oracle over a prepared workload: the absorb path of + /// `reference_pipeline` minus the file IO, parsing, and finalization. + fn run_merge_index_reference(cases: Vec) -> PyResult<(Vec, u64, u64)> { + let mut index: MergeIndexReference = MergeIndexReference::default(); + for case in cases { + index.absorb(case.id, case.record, case.content)?; + } + let mut output: Vec = Vec::new(); + for id in std::mem::take(&mut index.order) { + let Some((_, mut value)) = index.records.remove(&id) else { + continue; + }; + strip_internal_edge_fields(&mut value); + output.extend_from_slice(&emitted_json_bytes(&value).map_err(runtime_error)?); + output.push(b'\n'); + } + Ok((output, index.merged, index.scalar_conflicts)) + } + + /// WHY: US-002 rewrote the merge fold from quadratic to near-linear, and equivalence + /// with the frozen oracle is already policed by + /// `merge_fold_matches_reference_on_fuzz` -- this test guards the SPEED half of that + /// work. It drives the production `MergeIndex` and the frozen quadratic + /// `MergeIndexReference` through the SAME seeded quadratic-shaped workload + /// in-process and demands the new fold beat the reference by at least `BOUND`. + /// If a future change silently regresses the fold back toward quadratic, this + /// trips. The bound is a same-process RATIO, not an absolute wall-clock limit: + /// both legs share the same core, allocator, and input, so the assertion is + /// machine-independent and needs no per-CI-box tuning. + #[test] + fn merge_fold_speedup_bound_vs_reference() { + // Warmup: lazy allocator/paging work must not bill to whichever leg runs first. + run_merge_index(fold_workload(WARMUP_SEED, 2, 4).expect("warmup workload")) + .expect("warmup new fold"); + run_merge_index_reference(fold_workload(WARMUP_SEED, 2, 4).expect("warmup workload")) + .expect("warmup reference fold"); + + // Build the same seeded workload per leg so each one consumes its own owned + // records and the timed region clones nothing: the ratio measures the fold + // algorithm alone, not input preparation. + // + // The NEW fold is timed FIRST (cold caches bill against it, so a bound that passes + // anyway is conservative) and BEST-OF-`NEW_FOLD_ATTEMPTS` (US-006): a millisecond + // leg timed once is dominated by noise, and noise in a ratio's denominator is how + // a >=5x bound turns flaky. The minimum of the attempts is compared, and each + // attempt rebuilds its workload so no timed region clones. + let mut new_elapsed: Duration = Duration::MAX; + let mut new_outcome: Option<(Vec, u64, u64)> = None; + for _ in 0..NEW_FOLD_ATTEMPTS { + let cases: Vec = + fold_workload(WORKLOAD_SEED, GROUPS, RECORDS_PER_GROUP).expect("workload"); + let started: Instant = Instant::now(); + let outcome: (Vec, u64, u64) = run_merge_index(cases).expect("new fold"); + let elapsed: Duration = started.elapsed(); + if elapsed < new_elapsed { + new_elapsed = elapsed; + new_outcome = Some(outcome); + } + } + let (new_bytes, new_merged, new_conflicts): (Vec, u64, u64) = + new_outcome.expect("at least one new-fold attempt ran"); + + // The frozen quadratic leg stays SINGLE-shot: it already runs for seconds, so its + // timing is stable, and repeating it would multiply this test's runtime without + // reducing noise. + let reference_cases: Vec = + fold_workload(WORKLOAD_SEED, GROUPS, RECORDS_PER_GROUP).expect("workload"); + let started: Instant = Instant::now(); + let (reference_bytes, reference_merged, reference_conflicts): (Vec, u64, u64) = + run_merge_index_reference(reference_cases).expect("reference fold"); + let reference_elapsed = started.elapsed(); + + // Identical input must yield identical outcomes: the ratio measures the + // algorithm, not an input mismatch. + assert_eq!( + (new_merged, new_conflicts), + (reference_merged, reference_conflicts), + "fold counters diverged on the identical workload" + ); + assert_eq!( + new_bytes, reference_bytes, + "fold outputs diverged on the identical workload" + ); + + // Non-vacuity: the workload must actually fold heavily, or a ratio on a + // degenerate input would be meaningless. + assert!( + new_merged >= (GROUPS * RECORDS_PER_GROUP * 9 / 10) as u64, + "expected heavy folding, merged={new_merged}" + ); + assert!( + new_conflicts >= 1, + "expected scalar conflicts, got {new_conflicts}" + ); + let output_lines: usize = new_bytes.iter().filter(|byte| **byte == b'\n').count(); + assert_eq!(output_lines, GROUPS, "expected one merged record per id"); + + let speedup: f64 = reference_elapsed.as_secs_f64() / new_elapsed.as_secs_f64(); + println!( + "merge fold speedup: new {new_elapsed:.3?} vs frozen quadratic reference \ + {reference_elapsed:.3?} -> {speedup:.1}x (bound {BOUND}x)" + ); + assert!( + speedup >= BOUND, + "the near-linear merge fold lost its speed margin: only {speedup:.2}x faster \ + than the frozen quadratic reference (new fold {new_elapsed:.3?}, reference \ + {reference_elapsed:.3?}); expected >= {BOUND}x on the identical workload" + ); + } +} + +#[cfg(test)] +mod merge_state_desync { + use super::MergedRecord; + use serde_json::json; + + /// WHY this test exists: `MergedRecord::finish`'s desync guard is load-bearing + /// fail-loudly precedent, not decoration. Hash-only edge keying once silently dropped + /// DISTINCT records at scale (the collision class `record_if_new_suppresses_only_exact_ + /// byte_duplicates` polices); the same silent-data-loss class lurks here if the + /// `bytes`-parallel-to-`items` invariant ever breaks, because `zip` truncates to the + /// shorter side while `drain(..)` empties the whole array. This forces that desync and + /// asserts `finish` REFUSES to write -- returning a structured `merge-state-desync` + /// error -- proving the guard is a real runtime check observable in EVERY build + /// profile, not a `debug_assert_eq!` that panics first in tests and can never surface + /// the structured error path. + #[test] + fn merge_state_desync_is_a_structured_error_not_silent_truncation() { + // Reading a `PyErr`'s message needs an initialized interpreter (pyo3 is built + // without `auto-initialize`); idempotent, so this is safe alongside the full suite. + pyo3::Python::initialize(); + + // Seed a record whose `ids` array holds three live items; `new` records a matching + // three-entry canonical-`bytes` list (unioned = false, so `finish` skips it as-is). + let mut record = + MergedRecord::new(json!({ "ids": [1, 2, 3] }), 0).expect("seed merged record"); + + // Desync it the way a broken parallel-invariant would: drop ONE canonical-byte + // entry (2 byte-entries left for 3 live items) and mark the field unioned so + // `finish` routes it through the zip/drain path the guard protects. + let state = record.lists.get_mut("ids").expect("ids list state"); + state.bytes.pop().expect("a canonical-byte entry to drop"); + state.unioned = true; + + // The guard must fire BEFORE any zip/drain mutates the record: `finish` returns a + // structured error naming the desync, so nothing is silently truncated/written. + let error = record + .finish() + .expect_err("a desynced list must fail loudly, not silently truncate"); + let message = error.to_string(); + assert!( + message.contains("merge-state-desync"), + "expected a structured merge-state-desync error, got: {message}" + ); + } } diff --git a/tests/bench_merge_bench.py b/tests/bench_merge_bench.py new file mode 100644 index 00000000..1a787385 --- /dev/null +++ b/tests/bench_merge_bench.py @@ -0,0 +1,296 @@ +"""In-repo benchmark harness for the Rust `rs.dedup_ndjson` dedup passes. + +WHY: US-002 rewrites the merge fold for speed and needs reproducible before/after +numbers without relying on throwaway /tmp scripts. This harness regenerates the four +merge-mode benchmark corpora -- byte-identical to the original `/tmp/merge_bench` +generators (`gen.py` / `gen_dup.py`, seed 42) -- and times `rs.dedup_ndjson` in merge +mode. US-006 adds scenarios E and N so the two OTHER passes that share the rewritten +canonical serializer are measured too (the audit found the default `error` mode and the +node path had no before/after number at all). + +Usage (opt-in; collected nowhere by default, so it never slows the suite): + + TABLASSERT_BENCH=1 uv run pytest tests/bench_merge_bench.py -s -n 0 -q + +`-n 0` keeps a single process (the datasets are session-scoped; xdist would rebuild +them in every worker), `-s` shows the timings. Generation writes ~1.7 GB of NDJSON +into pytest's basetemp and takes a couple of minutes before any timing starts. + +`TABLASSERT_BENCH_ONLY=` restricts which corpora get +generated AND timed -- the baseline half of a before/after comparison only needs the +scenarios under test, and skipping the rest saves ~1.1 GB of generation: + + TABLASSERT_BENCH=1 TABLASSERT_BENCH_ONLY=e,n uv run pytest tests/bench_merge_bench.py -s -n 0 -q + + scenario pass / mode shape baseline (pre US-002) + A merge 200k triples x 3 rows, 20 cases 6.71s + B merge 5k triples x 20 rows, 100 cases 20.69s + C merge 1k triples x 100 rows, 50 cases 37.45s + DUP merge 100k byte-identical rows 0.59s + E default `error` mode, edges 598k lines (2% exact repeats) see US-006 brief + N node dedup (`is_edges=False`) 600k lines (300k unique) see US-006 brief +""" + +from __future__ import annotations + +import json +import os +import random +import time +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif(os.environ.get("TABLASSERT_BENCH") != "1", reason="benchmark harness; opt in with TABLASSERT_BENCH=1") + +_SEED = 42 +_DOMAIN = "infores:multiomicskg" +_UUID_FIELDS = ["subject", "predicate", "object"] +# (n_triples, rows_per_triple, n_cases) -- exactly the shapes behind /tmp/merge_bench's +# edges_a/b/c.ndjson (verified byte-identical against the originals, seed 42). +_SCENARIOS: dict[str, tuple[int, int, int]] = {"a": (200_000, 3, 20), "b": (5_000, 20, 100), "c": (1_000, 100, 50)} +# Scenario E (default `error` mode). `source_record_urls` is unique per source row, so the +# three divergent rows of a triple derive three DISTINCT ids and the pass never aborts with +# `uuid-fields-not-a-key` -- the discriminating field that error's own message recommends. +_ERROR_UUID_FIELDS = ["subject", "predicate", "object", "source_record_urls"] +_ERROR_TRIPLES = 195_000 +_ERROR_CASES = 20 +# Re-emit every Nth triple's first row byte-identically, so duplicate suppression is real +# work (~2% of the stream) instead of the pass only ever seeing fresh ids. +_ERROR_DUPLICATE_EVERY = 15 +# Scenario N (node dedup): `1 + node_i % 3` copies of each unique node -> 600k lines over +# 300k distinct records, so `record_if_new` suppresses half the stream by byte equality. +_NODE_COUNT = 300_000 +_ALL_SCENARIOS: tuple[str, ...] = ("a", "b", "c", "dup", "e", "n") + + +def _edge(rng: random.Random, triple_i: int, row_j: int, n_cases: int) -> dict[str, object]: + """One synthetic edge; verbatim port of /tmp/merge_bench/gen.py's `edge`.""" + return { + "subject": f"MONDO:{triple_i:07d}", + "predicate": "biolink:associated_with", + "object": f"NCBIGene:{triple_i % 20000:07d}", + "p_value": f"{rng.random() * 0.05:.3e}", + "effect_size": round(rng.random(), 4), + "supporting_text": f"Study row text variant {row_j} for association {triple_i}.", + "publications": [f"PMCID:PMC{(triple_i * 7 + k) % 12_000_000:08d}" for k in range(row_j % 4 + 1)], + "sources": [{"resource_id": "infores:multiomicskg", "resource_role": "primary_knowledge_source"}], + "source_record_urls": [f"https://example.org/table/{triple_i}#row{row_j}"], + "has_supporting_studies": [f"STUDY:{(triple_i + row_j) % 5000:05d}"], + "number_of_cases": n_cases, + "supporting_case_ids": [f"CASE:{triple_i:06d}:{row_j:03d}:{c:04d}" for c in range(n_cases)], + } + + +def _generate_scenario(path: Path, n_triples: int, rows_per_triple: int, n_cases: int) -> int: + """Reproduce /tmp/merge_bench/gen.py byte-for-byte (seed 42); return the line count.""" + rng = random.Random(_SEED) + lines = 0 + with path.open("w") as handle: + for triple_i in range(n_triples): + for row_j in range(rows_per_triple): + handle.write(json.dumps(_edge(rng, triple_i, row_j, n_cases if row_j else max(1, n_cases // 2))) + "\n") + lines += 1 + return lines + + +def _generate_error_edges(path: Path) -> int: + """Scenario E corpus: default-`error`-mode edges plus exact-duplicate rows. + + Reuses `_edge` verbatim (same seed, same 20-case shape as scenario A) so each line is + the realistic ~700-byte record that `finalize_record` canonicalizes once per line on + this pass -- the cost US-003 removed the clone from, and the reason this scenario + exists. Returns the line count (598,000 = 195k x 3 + 13k repeats). + """ + rng = random.Random(_SEED) + lines = 0 + with path.open("w") as handle: + for triple_i in range(_ERROR_TRIPLES): + first_row = "" + for row_j in range(3): + line = json.dumps(_edge(rng, triple_i, row_j, _ERROR_CASES if row_j else max(1, _ERROR_CASES // 2))) + "\n" + if row_j == 0: + first_row = line + handle.write(line) + lines += 1 + if triple_i % _ERROR_DUPLICATE_EVERY == 0: + # Byte-identical repeat of a row already written: same derived id AND same + # content hash -> `EdgeVerdict::Duplicate`, suppressed without aborting. + handle.write(first_row) + lines += 1 + return lines + + +def _node(rng: random.Random, node_i: int) -> dict[str, object]: + """One synthetic KGX node, in the post-coercion shape the node pass actually sees. + + No empty/null-like values: `strip_nulls` would drop them and the corpus would no longer + measure the bytes it claims to. + """ + return { + "id": f"MONDO:{node_i:07d}", + "name": f"Disease term number {node_i}", + "category": ["biolink:Disease" if node_i % 2 else "biolink:PhenotypicFeature"], + "description": f"Synthetic node description variant {rng.randrange(97)} for term {node_i}.", + "synonym": [f"alias {node_i} {k}" for k in range(1 + node_i % 3)], + "xrefs": [f"UMLS:C{(node_i * 3) % 2_000_000:07d}", f"DOID:{node_i % 20_000}"], + "provided_by": [_DOMAIN], + } + + +def _generate_nodes(path: Path) -> int: + """Scenario N corpus: node lines with heavy exact duplication (600k lines, 300k unique). + + Duplicates are written adjacent to their original, which is the common real shape + (records arrive grouped by source table) and keeps generation streaming -- no corpus is + ever held in memory. Returns the line count. + """ + rng = random.Random(_SEED) + lines = 0 + with path.open("w") as handle: + for node_i in range(_NODE_COUNT): + line = json.dumps(_node(rng, node_i)) + "\n" + for _ in range(1 + node_i % 3): + handle.write(line) + lines += 1 + return lines + + +def _generate_duplicates(path: Path) -> int: + """Reproduce /tmp/merge_bench/gen_dup.py byte-for-byte; return the line count.""" + row: dict[str, object] = { + "subject": "MONDO:0000001", + "predicate": "biolink:associated_with", + "object": "NCBIGene:0000001", + "p_value": "1e-5", + "effect_size": 0.5, + "supporting_text": "same", + "publications": ["PMCID:PMC1"], + "sources": [{"resource_id": "infores:x", "resource_role": "primary_knowledge_source"}], + "number_of_cases": 50, + "supporting_case_ids": [f"CASE:{c:04d}" for c in range(50)], + } + line = json.dumps(row) + "\n" + with path.open("w") as handle: + for _ in range(100_000): + handle.write(line) + return 100_000 + + +def _selected_scenarios() -> tuple[str, ...]: + """Scenario names to generate and time: `TABLASSERT_BENCH_ONLY`, or all six. + + Fails loudly on an unknown name -- a typo must not silently time nothing and report a + green benchmark run. + """ + raw = os.environ.get("TABLASSERT_BENCH_ONLY", "") + if not raw.strip(): + return _ALL_SCENARIOS + names = tuple(part.strip().lower() for part in raw.split(",") if part.strip()) + unknown = [name for name in names if name not in _ALL_SCENARIOS] + if unknown: + raise ValueError(f"TABLASSERT_BENCH_ONLY names unknown scenario(s) {unknown}; expected a subset of {list(_ALL_SCENARIOS)}") + return names + + +@pytest.fixture(scope="session") +def bench_datasets(tmp_path_factory: pytest.TempPathFactory) -> dict[str, tuple[Path, int]]: + """Generate the selected corpora once per session (skipped entirely unless opted in). + + Each generator seeds its own `random.Random(_SEED)`, so a corpus is byte-identical no + matter which subset was requested or which worktree generated it -- that is what lets a + baseline checkout and the branch checkout be compared on the same input (verified by + checksum in the US-006 run). + """ + root = tmp_path_factory.mktemp("merge_bench") + selected = _selected_scenarios() + datasets: dict[str, tuple[Path, int]] = {} + for name in _ALL_SCENARIOS: + if name not in selected: + continue + if name in _SCENARIOS: + path = root / f"edges_{name}.ndjson" + datasets[name] = (path, _generate_scenario(path, *_SCENARIOS[name])) + elif name == "dup": + path = root / "edges_dup.ndjson" + datasets[name] = (path, _generate_duplicates(path)) + elif name == "e": + path = root / "edges_error_mode.ndjson" + datasets[name] = (path, _generate_error_edges(path)) + else: # "n" + path = root / "nodes_dedup.ndjson" + datasets[name] = (path, _generate_nodes(path)) + return datasets + + +def _time_dedup( + name: str, datasets: dict[str, tuple[Path, int]], tmp_path: Path, *, is_edges: bool, uuid_fields: list[str] | None, on_collision: str | None +) -> None: + """Time one `rs.dedup_ndjson` run on a generated corpus and report it. + + The returned pair is `(merged divergent records, conflicting scalar fields)` in merge + mode and `(0, 0)` on the other two passes, so the mode is printed beside it to keep all + six scenarios' lines comparable in one table. Output lines are counted OUTSIDE the timed + region: they are the non-vacuity evidence that duplicates were really suppressed. + """ + from tablassert import rs + + if name not in datasets: + pytest.skip(f"scenario {name} was not generated (TABLASSERT_BENCH_ONLY={os.environ.get('TABLASSERT_BENCH_ONLY', '')!r})") + p_in, line_count = datasets[name] + p_out = tmp_path / f"deduped_{name}.ndjson" + start = time.perf_counter() + merged, conflicts = rs.dedup_ndjson(p_in, p_out, is_edges, _DOMAIN, uuid_fields, on_collision) + elapsed = time.perf_counter() - start + assert p_out.stat().st_size > 0, f"scenario {name} produced an empty output" + with p_out.open("rb") as handle: + written = sum(1 for _ in handle) + mode = "nodes" if not is_edges else f"edges/{on_collision or 'error'}" + print( + f"\nbench {name}: mode={mode} lines={line_count:,} written={written:,} merged={merged:,} conflicts={conflicts:,} " + f"elapsed={elapsed:.3f}s ({line_count / elapsed:,.0f} lines/s)" + ) + + +def test_bench_scenario_a(bench_datasets: dict[str, tuple[Path, int]], tmp_path: Path) -> None: + """Scenario A: 200k triples x 3 rows, 20 cases. Baseline pre-US-002: 6.71s.""" + _time_dedup("a", bench_datasets, tmp_path, is_edges=True, uuid_fields=_UUID_FIELDS, on_collision="merge") + + +def test_bench_scenario_b(bench_datasets: dict[str, tuple[Path, int]], tmp_path: Path) -> None: + """Scenario B: 5k triples x 20 rows, 100 cases. Baseline pre-US-002: 20.69s.""" + _time_dedup("b", bench_datasets, tmp_path, is_edges=True, uuid_fields=_UUID_FIELDS, on_collision="merge") + + +def test_bench_scenario_c(bench_datasets: dict[str, tuple[Path, int]], tmp_path: Path) -> None: + """Scenario C: 1k triples x 100 rows, 50 cases. Baseline pre-US-002: 37.45s.""" + _time_dedup("c", bench_datasets, tmp_path, is_edges=True, uuid_fields=_UUID_FIELDS, on_collision="merge") + + +def test_bench_scenario_dup(bench_datasets: dict[str, tuple[Path, int]], tmp_path: Path) -> None: + """Scenario DUP: 100k byte-identical rows (pure suppression path). Baseline: 0.59s.""" + _time_dedup("dup", bench_datasets, tmp_path, is_edges=True, uuid_fields=_UUID_FIELDS, on_collision="merge") + + +def test_bench_scenario_e(bench_datasets: dict[str, tuple[Path, int]], tmp_path: Path) -> None: + """Scenario E: default `error` mode over 598k edge lines (13k of them exact repeats). + + WHY: US-003's allocation-free `canonical_json_bytes` is called once per line by + `finalize_record`'s content hash on the DEFAULT mode too -- not just inside the merge + fold -- so the rewrite lands on every production edge build. The merge-only A-DUP + scenarios never measured that path, which is exactly the gap this scenario closes. + """ + _time_dedup("e", bench_datasets, tmp_path, is_edges=True, uuid_fields=_ERROR_UUID_FIELDS, on_collision=None) + + +def test_bench_scenario_n(bench_datasets: dict[str, tuple[Path, int]], tmp_path: Path) -> None: + """Scenario N: node dedup over 600k node lines (300k unique, so half are suppressed). + + WHY: the node pass is the other leg REQ-PERF-4/DoD-8 asks for. Measured, not assumed, + it turns out NOT to reach `canonical_json_bytes` at all: `finalize_record` returns + before the content hash when `is_edges` is false (`rust/src/ndjson.rs`) and + `record_if_new` keys on `emitted_json_bytes`. So N is the regression guard for the pass + US-003 leaves untouched and its before/after delta is expected to be run-to-run noise, + while E is the leg that actually exercises the rewritten serializer. + """ + _time_dedup("n", bench_datasets, tmp_path, is_edges=False, uuid_fields=None, on_collision=None) diff --git a/tests/test_rs.py b/tests/test_rs.py index 26108bd6..da62096e 100644 --- a/tests/test_rs.py +++ b/tests/test_rs.py @@ -1,7 +1,12 @@ from __future__ import annotations +import itertools import json +import random +import uuid +from collections.abc import Iterable from pathlib import Path +from typing import Any from uuid import UUID @@ -52,3 +57,330 @@ def test_dedup_ndjson_labels_edges(tmp_path: Path) -> None: row: dict = json.loads(p_out.read_text()) UUID(row["id"]) # edge id is a valid UUID + + +# --------------------------------------------------------------------------- +# Pure-Python reference of the CURRENT rust merge-mode semantics (frozen oracle +# for the US-002 rewrite). Mirrors `strip_nulls` / `uuid_for_json_object` / +# `merge_records` / `dedup_edges_merge` from `rust/src/{json,uuid,ndjson}.rs`. +# Intentionally simple (lists, not hash sets) because correctness against the +# rust output matters more than speed here. +# --------------------------------------------------------------------------- + +_BAD_TOKENS = frozenset({"", "na", "nan", "null", "none"}) +_MISSING: Any = object() +_NIL_NAMESPACE = uuid.UUID(int=0) + + +def _canonical_json_bytes(value: Any) -> bytes: + """Mirror of rust `canonical_json_bytes`: keys sorted recursively, compact bytes. + + ASCII-only datasets only (the sibling of `_fuzz_record`'s int-only constraint): + `ensure_ascii=True` orders non-ASCII list items by their `\\uXXXX` escapes while rust + orders by raw UTF-8 bytes, so a non-ASCII item could sort differently here. + """ + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True).encode("utf-8") + + +def _strip_nulls(record: dict[str, Any]) -> dict[str, Any]: + """Mirror of rust `strip_nulls` (json.rs): drop absent/null-like values only. + + Deliberately NOT Python truthiness -- `0`, `0.0` and `false` are meaningful values + and stay; nested dicts recurse (an emptied nested dict survives as `{}`), and list + items recurse only when they are dicts. + """ + + def is_present(value: Any) -> bool: + if value is None: + return False + if isinstance(value, bool | int | float): + return True + if isinstance(value, str): + return value != "" + return len(value) > 0 # lists and dicts + + def passes_bad_check(value: Any) -> bool: + return not (isinstance(value, str) and value.strip().lower() in _BAD_TOKENS) + + def transform(value: Any) -> Any: + if isinstance(value, list): + return [_strip_nulls(item) if isinstance(item, dict) else item for item in value] + if isinstance(value, dict): + return _strip_nulls(value) + return value + + return {key: transform(value) for key, value in record.items() if is_present(value) and passes_bad_check(value)} + + +def _uuid_part(value: Any) -> str | None: + """Mirror of rust `uuid_part`: normalize one field value to its hashable form.""" + if value is None: + return None + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int | float): + return str(value) + if isinstance(value, str): + return value if value else None + return _canonical_json_bytes(value).decode("utf-8") + + +def _edge_uuid(domain: str, record: dict[str, Any], fields: list[str]) -> str: + """Mirror of rust `uuid_for_json_object` over declared `uuid_fields`. + + v3 UUIDs: namespaced by `uuid3(NIL, domain)`, over the length-prefixed + `:` join of sorted key/value parts; entries whose value normalizes + to nothing (null, empty) drop out key included. + """ + namespace = uuid.uuid3(_NIL_NAMESPACE, domain) + keys: list[str] = sorted({key for key in fields if key in record}) + parts: list[str] = [] + for key in keys: + part = _uuid_part(record[key]) + if part is not None: + parts.append(key) + parts.append(part) + joined = "".join(f"{len(part.encode('utf-8'))}:{part}" for part in parts) + return str(uuid.uuid3(namespace, joined)) + + +def _merge_records(stored: dict[str, Any], incoming: dict[str, Any]) -> int: + """Mirror of rust `merge_records`: fold `incoming` into `stored` field-wise. + + List fields union by canonical bytes and sort ONLY when both sides carry an array + (a list copied from a later record keeps its order); stored-side duplicates survive, + incoming-side ones collapse. Scalars are first-wins, each difference counts one + conflict; a field only on `incoming` is copied, not a conflict. `id` is untouched. + `number_of_cases` is recomputed to the union length of `supporting_case_ids` when the + merged record carries the carrier list and either side carried a count, and that + superseded divergence is decremented out of the conflict total (both sides present, + unequal, not both arrays -- the exact mirror of the rust loop condition). + """ + conflicts = 0 + stored_cases: Any = stored.get("number_of_cases", _MISSING) + incoming_cases: Any = incoming.get("number_of_cases", _MISSING) + for key, incoming_value in incoming.items(): + if key == "id": + continue + if key not in stored: + stored[key] = incoming_value + continue + stored_value = stored[key] + if isinstance(stored_value, list) and isinstance(incoming_value, list): + seen: list[bytes] = [_canonical_json_bytes(item) for item in stored_value] + for item in incoming_value: + item_bytes = _canonical_json_bytes(item) + if item_bytes not in seen: + seen.append(item_bytes) + stored_value.append(item) + stored_value.sort(key=_canonical_json_bytes) + elif stored_value != incoming_value: + conflicts += 1 + union = stored.get("supporting_case_ids") + if isinstance(union, list) and (stored_cases is not _MISSING or incoming_cases is not _MISSING): + if ( + stored_cases is not _MISSING + and incoming_cases is not _MISSING + and stored_cases != incoming_cases + and not (isinstance(stored_cases, list) and isinstance(incoming_cases, list)) + ): + conflicts -= 1 + stored["number_of_cases"] = len(union) + return conflicts + + +def _swap_remove(record: dict[str, Any], key: str) -> None: + """Mirror of serde_json `Map::remove` under `preserve_order`: indexmap `swap_remove`. + + The LAST key fills the vacated slot -- removal does NOT preserve the order of the + remaining keys. This is part of the pinned emit semantics of `strip_internal_edge_ + fields`, so the reference must reproduce it byte-for-byte. + """ + if key not in record: + return + keys = list(record.keys()) + if keys[-1] == key: + del record[key] + return + rebuilt: dict[str, Any] = {} + for existing in keys: + if existing == key: + rebuilt[keys[-1]] = record[keys[-1]] + elif existing != keys[-1]: + rebuilt[existing] = record[existing] + record.clear() + record.update(rebuilt) + + +def _merge_reference(lines: Iterable[str], domain: str, fields: list[str]) -> tuple[bytes, int, int]: + """Drive the CURRENT rust merge pass over NDJSON lines, entirely in Python. + + Mirror of `dedup_edges_merge`: skip blank lines and empty objects, strip nulls, + label each edge, suppress exact content repeats WITHOUT counting them, fold divergent + same-id records first-wins, then emit one compact line per id in FIRST-SEEN order with + `supporting_case_ids` stripped. Returns `(output bytes, merged, scalar_conflicts)`. + """ + records: dict[str, tuple[list[bytes], dict[str, Any]]] = {} + order: list[str] = [] + merged = 0 + conflicts = 0 + for line in lines: + if not line.strip(): + continue + cleaned = _strip_nulls(json.loads(line)) + if not cleaned: + continue + content = _canonical_json_bytes(cleaned) + cleaned["id"] = _edge_uuid(domain, cleaned, fields) + slot = records.get(cleaned["id"]) + if slot is None: + records[cleaned["id"]] = ([content], cleaned) + order.append(cleaned["id"]) + continue + hashes, stored = slot + if content in hashes: + continue + merged += 1 + conflicts += _merge_records(stored, cleaned) + hashes.append(content) + chunks: list[bytes] = [] + for edge_id in order: + _, value = records[edge_id] + _swap_remove(value, "supporting_case_ids") + chunks.append(json.dumps(value, separators=(",", ":"), ensure_ascii=True).encode("utf-8") + b"\n") + return b"".join(chunks), merged, conflicts + + +_SUBJECTS = [f"MONDO:{index:03d}" for index in range(6)] +_PREDICATES = ["biolink:related_to", "biolink:associated_with"] +_OBJECTS = [f"NCBIGene:{index:03d}" for index in range(4)] +_TAGS = ["tag:a", "tag:b", "tag:c", "tag:d", "tag:e"] +_CASE_IDS = [f"case:{index}" for index in range(8)] +_P_VALUES = ["1e-5", "0.01", "0.99"] + + +def _fuzz_source(rng: random.Random) -> dict[str, str]: + identity, role = ("infores:one", "primary_knowledge_source") if rng.randrange(2) == 0 else ("infores:two", "aggregator_knowledge_source") + if rng.randrange(2) == 0: + return {"resource_id": identity, "resource_role": role} + return {"resource_role": role, "resource_id": identity} # same source, different key order + + +def _fuzz_record(rng: random.Random, group: int, record_index: int) -> dict[str, Any]: + record: dict[str, Any] = {} + # Identity triple in random insertion order: same id, divergent bytes -- the fold, not + # the id derivation, decides the outcome. (Int scalars only: float formatting differs + # between rust and python serializers and is irrelevant to merge semantics.) + triple: list[tuple[str, str]] = [("subject", rng.choice(_SUBJECTS)), ("predicate", rng.choice(_PREDICATES)), ("object", rng.choice(_OBJECTS))] + rng.shuffle(triple) + record.update(triple) + record["p_value"] = rng.choice(_P_VALUES) + record["effect_size"] = rng.randrange(4) + if rng.randrange(4) == 0: + record["negated"] = rng.randrange(2) == 0 + if rng.randrange(2) == 0: + # Drawn WITH replacement: a record may repeat an item inside its own array. + record["tags"] = [rng.choice(_TAGS) for _ in range(rng.randrange(4))] + if rng.randrange(2) == 0: + record["sources"] = [_fuzz_source(rng) for _ in range(1 + rng.randrange(2))] + if rng.randrange(3) == 0: + record["mode"] = "solo" if rng.randrange(2) == 0 else ["solo", "extra"] # scalar-vs-array conflict + if record_index > 0 and (group % 3 == 0 or rng.randrange(2) == 0): + record["late"] = f"late:{group}" # field only later records carry + if group % 2 == 0 and record_index == 1: + # Same "only a later record carries it" shape but LIST-valued, on exactly ONE record + # per group: copied in, never unioned, so the write-out must NOT sort it. Items are + # strictly DESCENDING by canonical bytes (reversed `_TAGS`) so a stray sort is + # observable; the test counts the descending survivors. Mirrors the rust fuzz + # generator's `late_list` (US-006 Fix 4). ASCII-only, per the `_canonical_json_bytes` + # ordering caveat above. + record["late_list"] = list(reversed(_TAGS[: 2 + rng.randrange(3)])) + if rng.randrange(2) == 0: + case_ids = [rng.choice(_CASE_IDS) for _ in range(1 + rng.randrange(4))] + record["supporting_case_ids"] = case_ids + if rng.randrange(4): + # Deliberately wrong sometimes: the recompute supersedes the count and must + # also undo the scalar conflict the divergence would otherwise have counted. + record["number_of_cases"] = len(case_ids) if rng.randrange(3) == 0 else len(case_ids) + 1 + rng.randrange(3) + return record + + +def _fuzz_dataset(rng: random.Random) -> list[str]: + records: list[dict[str, Any]] = [] + for group in range(36): # divergent same-id groups of 1-7 records + for record_index in range(1 + rng.randrange(7)): + records.append(_fuzz_record(rng, group, record_index)) + rng.shuffle(records) + lines = [json.dumps(record) for record in records] + unique = len(lines) + for _ in range(10): # exact byte repeats -> content-hash suppression, no counters + lines.append(lines[rng.randrange(unique)]) + rng.shuffle(lines) + stream: list[str] = [] + for line in lines: + if rng.randrange(14) == 0: + stream.append("") # blank lines are legal stream noise + if rng.randrange(20) == 0: + stream.append(" ") + if rng.randrange(16) == 0: + stream.append("{}") # empty objects are skipped + stream.append(line) + return stream + + +def _is_strictly_descending(items: Any) -> bool: + """True when a list's canonical bytes strictly decrease -- i.e. it is NOT sorted. + + WHY: the fuzz datasets emit `late_list` in strictly descending order, so this predicate + is how the test recognizes a list that was merely COPIED into a record (never unioned) + and therefore must have survived the write-out in its source order. + """ + if not isinstance(items, list) or len(items) < 2: + return False + keys = [_canonical_json_bytes(item) for item in items] + return all(left > right for left, right in itertools.pairwise(keys)) + + +def test_dedup_edges_merge_matches_python_reference(tmp_path: Path) -> None: + """Merge mode must byte-match an independent Python port of its own semantics. + + WHY: US-002 will rewrite the rust merge fold for speed. This test pins the CURRENT + semantics from a SECOND implementation: a seeded randomized dataset (divergent same-id + groups, list unions over object and scalar items, key-order variants, late fields, + including one list a single record carries so it is copied and never unioned, + scalar-vs-array conflicts, exact repeats, empty objects, blank lines) is deduped by + `rs.dedup_ndjson(..., on_collision="merge")` and by the pure-Python reference above. + The emitted bytes -- first-seen id order, folded records, `supporting_case_ids` + stripped -- plus the `(merged, scalar_conflicts)` counters must agree exactly, so any + semantic drift in the rewrite fails here. + """ + from tablassert import rs + + rng = random.Random(1337) + lines = _fuzz_dataset(rng) + domain = "infores:multiomicskg" + fields = ["subject", "predicate", "object"] + p_in: Path = tmp_path / "edges.ndjson.tmp" + p_out: Path = tmp_path / "edges.ndjson" + p_in.write_text("\n".join(lines) + "\n") + + merged, conflicts = rs.dedup_ndjson(p_in, p_out, True, domain, fields, "merge") + expected, expected_merged, expected_conflicts = _merge_reference(lines, domain, fields) + + # Non-vacuity: the seeded dataset must actually exercise the fold. + assert merged > 0, "seeded dataset must actually fold records" + assert conflicts > 0, "seeded dataset must produce scalar conflicts" + actual = p_out.read_bytes() + assert b"supporting_case_ids" not in actual + assert actual == expected + assert (merged, conflicts) == (expected_merged, expected_conflicts) + + # The copied-never-unioned gate (US-006 Fix 4): `late_list` is generated strictly + # descending, so a write-out that sorted a merely COPIED list would flip it ascending + # and drop this count to zero. A lower bound rather than an exact count because the + # identity triple is drawn per RECORD, so two carriers can share a derived id -- those + # legitimately union and sort, and the byte comparison above already proves both + # implementations agree on them. + carriers = [row["late_list"] for row in map(json.loads, actual.decode().splitlines()) if "late_list" in row] + preserved = sum(1 for items in carriers if _is_strictly_descending(items)) + assert preserved >= 3, f"expected >= 3 copied-but-never-unioned lists to keep their source order, got {preserved} of {len(carriers)}"