use proptest::prelude::*;
use super::*;
use crate::repr::{Batch, BatchBuilder};
use crate::schema::{SchemaColumn, SchemaDescriptor, TypeCode};
use crate::test_support::{
join_reference, make_batch, make_batch_opk, make_schema_i64pk_i64, make_schema_u128_i64, make_schema_u64_i64,
opk_pk, pk_only_schema, pk_payload_schema, rekey_plan, trace_cursor, zset_of, TestTrace,
};
use gnitz_wire::read_i64_le;
fn plan(
kind: JoinKind,
delta_is_right: bool,
delta_schema: &SchemaDescriptor,
trace_schema: &SchemaDescriptor,
) -> JoinPlan {
JoinPlan::from_wire(kind, delta_is_right, delta_schema, trace_schema).expect("fixture join plan is well-formed")
}
fn join(
kind: JoinKind,
delta_is_right: bool,
delta_schema: &SchemaDescriptor,
trace_schema: &SchemaDescriptor,
delta: &Batch,
cursor: &mut ReadCursor,
) -> Batch {
let p = plan(kind, delta_is_right, delta_schema, trace_schema);
op_join_delta_trace(&delta.to_consolidated(), cursor, &p.out_schema, &p.probe)
}
#[test]
fn the_output_lays_out_the_key_then_both_payloads() {
use TypeCode::*;
let u128_string = SchemaDescriptor::new(
&[SchemaColumn::new(U128, false), SchemaColumn::new(String, false)],
&[0],
);
let compound = SchemaDescriptor::new(&[SchemaColumn::new(U64, false); 4], &[1, 2]);
type Case = (
JoinKind,
SchemaDescriptor,
SchemaDescriptor,
&'static [u32],
&'static [TypeCode],
);
let cases: [Case; 3] = [
(
JoinKind::Equi,
make_schema_u128_i64(),
u128_string,
&[0],
&[U128, I64, String],
),
(
JoinKind::Cross,
make_schema_u128_i64(),
u128_string,
&[0, 1],
&[U128, U128, I64, String],
),
(
JoinKind::Equi,
compound,
pk_payload_schema(&[U64; 2]),
&[0, 1],
&[U64, U64, U64, U64, I64],
),
];
for (kind, left, right, pk, types) in cases {
let joined = plan(kind, false, &left, &right).out_schema;
assert_eq!(
plan(kind, true, &right, &left).out_schema,
joined,
"{kind:?}: the delta's port"
);
assert_eq!(joined.pk_cols(), pk, "{kind:?}");
let got: Vec<TypeCode> = (0..joined.num_columns()).map(|c| joined.columns[c].type_code).collect();
assert_eq!(got, types, "{kind:?}");
}
}
#[test]
fn a_keyed_join_refuses_mismatched_pk_types() {
let signed = pk_payload_schema(&[TypeCode::I64]);
let unsigned = pk_payload_schema(&[TypeCode::U64]);
let narrow = pk_payload_schema(&[TypeCode::U32]);
for kind in [JoinKind::Equi, JoinKind::Range { rel: RangeRel::Lt }] {
for (a, b) in [(&signed, &unsigned), (&unsigned, &narrow)] {
let err = JoinPlan::from_wire(kind, false, a, b)
.err()
.expect("a keyed join must refuse mismatched PK types")
.to_string();
assert!(err.contains("PK column types differ"), "{kind:?}: {err}");
}
}
assert!(JoinPlan::from_wire(JoinKind::Cross, false, &signed, &narrow).is_ok());
}
#[test]
fn equi_join_products_the_trace_group_at_every_pk_shape() {
let shapes: [(&str, SchemaDescriptor, &[u128], &[u128]); 3] = [
("i32", pk_payload_schema(&[TypeCode::I32]), &[-7i64 as u128], &[1]),
("u64", pk_payload_schema(&[TypeCode::U64]), &[1], &[2]),
(
"3xu64",
pk_payload_schema(&[TypeCode::U64; 3]),
&[1, 1, 2],
&[1, 1, 1 << 56],
),
];
for (name, schema, held, absent) in shapes {
let (held, absent) = (opk_pk(&schema, held), opk_pk(&schema, absent));
let trace = make_batch_opk(&schema, &[(&held, 1, 100), (&held, 2, 200)]);
let delta = make_batch_opk(
&schema,
&[(&held, 1, 10), (&held, 1, 20), (&held, 1, 30), (&absent, 1, 40)],
);
let rows = assert_matches_reference(JoinKind::Equi, false, schema, schema, &delta, &trace, name);
assert_eq!(rows, 6, "{name}");
}
}
#[test]
fn an_equi_join_claims_the_output_it_writes_in_order() {
let schema = make_schema_u64_i64();
let trace = make_batch(&schema, &[(1, 1, 100), (1, 1, 200), (2, 1, 300)]);
type Case = (&'static str, &'static [(u64, i64, i64)], bool, bool);
let cases: [Case; 4] = [
("one delta row per key", &[(1, 1, 10), (2, 1, 20)], true, true),
("a delta run on a one-row key", &[(2, 1, 20), (2, 1, 21)], true, true),
("a delta run on a two-row key", &[(1, 1, 10), (1, 1, 11)], false, true),
("no match", &[(3, 1, 30)], false, false),
];
for (what, rows, left, right) in cases {
let delta = make_batch(&schema, rows);
for (delta_is_right, claimed) in [(false, left), (true, right)] {
let mut cursor = trace_cursor(Batch::clone(&trace));
let out = join(JoinKind::Equi, delta_is_right, &schema, &schema, &delta, &mut cursor);
assert_eq!(
out.is_consolidated(),
claimed || out.is_empty(),
"{what}, right={delta_is_right}"
);
let (want, _) = join_reference(JoinKind::Equi, delta_is_right, &schema, &schema, &delta, &trace);
assert_eq!(zset_of(&out, out.schema()), want, "{what}, right={delta_is_right}");
}
}
}
fn range_join(schema: &SchemaDescriptor, rel: RangeRel, delta: &Batch, cursor: &mut ReadCursor) -> Batch {
let kind = JoinKind::Range { rel };
join(kind, true, schema, schema, delta, cursor)
}
#[test]
fn range_join_cuts_the_span_each_rel_names() {
let schema = make_schema_u64_i64();
for (rel, want) in [
(RangeRel::Lt, vec![110]), (RangeRel::Le, vec![110, 120]), (RangeRel::Gt, vec![130]), (RangeRel::Ge, vec![120, 130]), ] {
let mut ch = trace_cursor(make_batch(&schema, &[(10, 1, 110), (20, 1, 120), (30, 1, 130)]));
let delta = make_batch(&schema, &[(20, 1, 200)]);
let out = range_join(&schema, rel, &delta, &mut ch);
let got: Vec<i64> = out_triples(&out).into_iter().map(|(t, _, _)| t).collect();
assert_eq!(got, want, "rel {rel:?}");
for r in 0..out.count {
assert_eq!(out.get_pk(r) as u64, 20);
assert_eq!(read_i64_le(out.col_data(1), r * 8), 200);
}
}
}
#[test]
fn range_join_orders_a_signed_key_by_its_opk_image() {
let schema = make_schema_i64pk_i64();
let trace_rows = [(-100i64 as u64, 1, 1), (0, 1, 2), (50, 1, 3)];
let delta = make_batch(&schema, &[(0, 1, 9)]);
for (rel, want) in [(RangeRel::Gt, vec![3]), (RangeRel::Lt, vec![1])] {
let mut ch = trace_cursor(make_batch(&schema, &trace_rows));
let out = range_join(&schema, rel, &delta, &mut ch);
let got: Vec<i64> = out_triples(&out).into_iter().map(|(t, _, _)| t).collect();
assert_eq!(got, want, "rel {rel:?}");
}
}
#[test]
fn a_used_trace_cursor_yields_the_fresh_cursor_output() {
let schema = make_range_schema(1, false);
let delta = make_range_batch(&schema, &[(vec![1], 5, 1, 1), (vec![3], 5, 1, 3)]);
let trace_rows = [
(vec![1u64], 0u64, 1i64, 10i64),
(vec![1], 9, 1, 19),
(vec![3], 0, 1, 30),
(vec![3], 9, 1, 39),
];
let ranges = RangeRel::ALL.iter().map(|&rel| JoinKind::Range { rel });
for kind in ranges.chain([JoinKind::Cross]) {
let out_schema = plan(kind, true, &schema, &schema).out_schema;
let cursor = || trace_cursor(make_range_batch(&schema, &trace_rows).into_consolidated());
let want = join(kind, true, &schema, &schema, &delta, &mut cursor());
for park_past_end in [false, true] {
let mut ch = cursor();
ch.advance_to(&opk_pk(&schema, &[3, 9]));
if park_past_end {
ch.advance();
assert!(!ch.valid);
}
let got = join(kind, true, &schema, &schema, &delta, &mut ch);
assert_eq!(got.count, want.count, "{kind:?} past_end={park_past_end}");
assert_eq!(
zset_of(&got, &out_schema),
zset_of(&want, &out_schema),
"{kind:?} past_end={park_past_end}",
);
}
}
}
type RangeCase = (&'static str, usize, RangeRows, RangeRows, usize);
type RangeRows = &'static [(&'static [u64], u64, i64, i64)];
#[test]
fn range_join_fixtures_match_the_reference() {
let cases: &[RangeCase] = &[
(
"maximal slot, no eq prefix",
0,
&[(&[], u64::MAX, 1, 9)],
&[(&[], 0, 1, 100), (&[], 50, 1, 150), (&[], u64::MAX, 1, 199)],
3,
),
(
"maximal slot inside an eq group",
1,
&[(&[1], u64::MAX, 1, 9)],
&[(&[1], 0, 1, 100), (&[1], 50, 1, 150), (&[2], 0, 1, 200)],
3,
),
(
"maximal slot present in the trace group",
1,
&[(&[1], u64::MAX, 1, 9)],
&[(&[1], 0, 1, 100), (&[1], u64::MAX, 1, 199), (&[2], 0, 1, 200)],
3,
),
(
"narrow span over a large trace",
0,
&[(&[], 25, 1, 1), (&[], 26, 1, 2)],
LARGE_TRACE,
50,
),
];
for &(name, n_eq, delta_rows, trace_rows, min_rows) in cases {
let schema = make_range_schema(n_eq, false);
let delta = make_range_batch(&schema, &owned(delta_rows));
let trace = make_range_batch(&schema, &owned(trace_rows));
let mut total = 0;
for &rel in RangeRel::ALL {
let kind = JoinKind::Range { rel };
total += assert_matches_reference(kind, true, schema, schema, &delta, &trace, name);
}
assert!(
total >= min_rows,
"{name}: emitted {total} rows, expected at least {min_rows}"
);
}
}
const LARGE_TRACE: RangeRows = &[
(&[], 0, 1, 100),
(&[], 1, 1, 101),
(&[], 2, 1, 102),
(&[], 3, 1, 103),
(&[], 4, 1, 104),
(&[], 5, 1, 105),
(&[], 10, 1, 110),
(&[], 15, 1, 115),
(&[], 20, 1, 120),
(&[], 24, 1, 124),
(&[], 25, 1, 125),
(&[], 26, 1, 126),
(&[], 27, 1, 127),
(&[], 28, 1, 128),
(&[], 29, 1, 129),
];
fn owned(rows: RangeRows) -> Vec<(Vec<u64>, u64, i64, i64)> {
rows.iter().map(|&(eq, d, w, v)| (eq.to_vec(), d, w, v)).collect()
}
fn arb_range_rows(n_eq: usize) -> impl Strategy<Value = Vec<(Vec<u64>, u64, i64, i64)>> {
let slot = prop_oneof![9 => 0u64..4, 1 => Just(u64::MAX)];
let row = (prop::collection::vec(0u64..4, n_eq), slot, -2i64..=2i64, -2i64..2i64);
prop::collection::vec(row, 0..8).prop_map(|mut rows| {
rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)).then(a.3.cmp(&b.3)));
rows
})
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(512))]
#[test]
fn join_dt_matches_reference(
(kind, delta_is_right, d_wide, t_wide, (d_eq, t_eq, delta_rows, trace_rows)) in (
prop_oneof![
Just(JoinKind::Equi),
prop::sample::select(RangeRel::ALL).prop_map(|rel| JoinKind::Range { rel }),
Just(JoinKind::Cross),
],
any::<bool>(), any::<bool>(), any::<bool>(),
).prop_flat_map(|(kind, delta_is_right, d_wide, t_wide)| {
let arities = match kind {
JoinKind::Cross => (0usize..2, 0usize..2).boxed(),
_ => (0usize..3).prop_map(|n| (n, n)).boxed(),
};
let rows = arities.prop_flat_map(|(d_eq, t_eq)| (Just(d_eq), Just(t_eq), arb_range_rows(d_eq), arb_range_rows(t_eq)));
(Just(kind), Just(delta_is_right), Just(d_wide), Just(t_wide), rows)
}),
) {
let (d_schema, t_schema) = (make_range_schema(d_eq, d_wide), make_range_schema(t_eq, t_wide));
let delta = make_range_batch(&d_schema, &delta_rows);
let trace = make_range_batch(&t_schema, &trace_rows);
assert_matches_reference(kind, delta_is_right, d_schema, t_schema, &delta, &trace, "proptest");
}
}
const FIXTURE_STRINGS: [&[u8]; 4] = [
b"",
b"abc",
b"a-long-string-past-the-inline-limit",
b"another-long-payload-string",
];
fn make_range_schema(n_eq: usize, wide: bool) -> SchemaDescriptor {
let mut cols: Vec<SchemaColumn> = (0..n_eq + 1).map(|_| SchemaColumn::new(TypeCode::U64, false)).collect();
cols.push(SchemaColumn::new(TypeCode::I64, false)); if wide {
cols.push(SchemaColumn::new(TypeCode::I64, true));
cols.push(SchemaColumn::new(TypeCode::String, false));
}
let pk: Vec<u32> = (0..n_eq as u32 + 1).collect();
SchemaDescriptor::new(&cols, &pk)
}
fn make_range_batch(schema: &SchemaDescriptor, rows: &[(Vec<u64>, u64, i64, i64)]) -> Batch {
let wide = schema.num_payload_cols() > 1;
let mut b = BatchBuilder::new(schema);
for (eq, range, w, val) in rows {
let mut vals: Vec<u128> = eq.iter().map(|&x| x as u128).collect();
vals.push(*range as u128);
b.begin_row_natives(&vals, *w);
b.put_int(*val as u128);
if wide {
match *val < 0 {
true => b.put_null(),
false => b.put_int(val.wrapping_mul(3) as u128),
}
b.put_blob(FIXTURE_STRINGS[val.rem_euclid(FIXTURE_STRINGS.len() as i64) as usize]);
}
b.end_row();
}
b.finish()
}
fn out_triples(out: &Batch) -> Vec<(i64, i64, i64)> {
(0..out.count)
.map(|r| {
(
read_i64_le(out.col_data(0), r * 8),
read_i64_le(out.col_data(1), r * 8),
out.get_weight(r),
)
})
.collect()
}
fn assert_matches_reference(
kind: JoinKind,
delta_is_right: bool,
delta_schema: SchemaDescriptor,
trace_schema: SchemaDescriptor,
delta: &Batch,
trace: &Batch,
what: &str,
) -> usize {
let p = plan(kind, delta_is_right, &delta_schema, &trace_schema);
let delta = &delta.to_consolidated();
let folded = Batch::clone(trace).into_consolidated();
let (want, want_rows) = join_reference(kind, delta_is_right, &delta_schema, &trace_schema, delta, &folded);
let cursors = [
("one run", trace_cursor(folded)),
("three runs", TestTrace::dealt(trace, 3).cursor()),
];
for (runs, mut ch) in cursors {
let out = op_join_delta_trace(delta, &mut ch, &p.out_schema, &p.probe);
let at = format!("{what}: kind={kind:?} delta_is_right={delta_is_right}, {runs}");
assert_eq!(out.count, want_rows, "{at}: row count");
assert_eq!(zset_of(&out, &p.out_schema), want, "{at}: z-set");
}
want_rows
}
fn cuts(eq: &[u8], d: &[u8], rel: RangeRel) -> Option<(Vec<u8>, Option<Vec<u8>>)> {
let mut pk = eq.to_vec();
pk.extend_from_slice(d);
let schema = pk_only_schema(&vec![TypeCode::U8; pk.len()]);
RangeProbe::new(&schema, rel, true)
.cut_points(&pk)
.map(|(s, e)| (s.pk_bytes().to_vec(), e.map(|e| e.pk_bytes().to_vec())))
}
#[test]
fn cut_points_bound_every_rel_within_its_eq_group() {
type Cut = Option<(&'static [u8], Option<&'static [u8]>)>;
let cases: &[(&[u8], &[u8], [Cut; 4])] = &[
(
&[],
&[0x05],
[
Some((&[0x00], Some(&[0x05]))),
Some((&[0x00], Some(&[0x06]))),
Some((&[0x06], None)),
Some((&[0x05], None)),
],
),
(
&[],
&[0x00],
[
None,
Some((&[0x00], Some(&[0x01]))),
Some((&[0x01], None)),
Some((&[0x00], None)),
],
),
(
&[],
&[0xFF],
[
Some((&[0x00], Some(&[0xFF]))),
Some((&[0x00], None)),
None,
Some((&[0xFF], None)),
],
),
(
&[0x07],
&[0x05],
[
Some((&[0x07, 0x00], Some(&[0x07, 0x05]))),
Some((&[0x07, 0x00], Some(&[0x07, 0x06]))),
Some((&[0x07, 0x06], Some(&[0x08, 0x00]))),
Some((&[0x07, 0x05], Some(&[0x08, 0x00]))),
],
),
(
&[0x07],
&[0xFF],
[
Some((&[0x07, 0x00], Some(&[0x07, 0xFF]))),
Some((&[0x07, 0x00], Some(&[0x08, 0x00]))),
None,
Some((&[0x07, 0xFF], Some(&[0x08, 0x00]))),
],
),
(
&[0xFF],
&[0x05],
[
Some((&[0xFF, 0x00], Some(&[0xFF, 0x05]))),
Some((&[0xFF, 0x00], Some(&[0xFF, 0x06]))),
Some((&[0xFF, 0x06], None)),
Some((&[0xFF, 0x05], None)),
],
),
];
for (eq, d, want) in cases {
for (rel, w) in RangeRel::ALL.iter().zip(want) {
let want = w.map(|(s, e)| (s.to_vec(), e.map(<[u8]>::to_vec)));
assert_eq!(cuts(eq, d, *rel), want, "eq={eq:02x?} d={d:02x?} rel={rel:?}");
}
}
}
fn keyed_source() -> SchemaDescriptor {
use TypeCode::*;
SchemaDescriptor::new(
&[
SchemaColumn::new(I64, false),
SchemaColumn::new(U64, false),
SchemaColumn::new(String, false),
SchemaColumn::new(I64, true),
],
&[0, 1],
)
}
fn keyed_source_batch(rows: &[(i64, u64, i64)]) -> Batch {
let schema = keyed_source();
let mut b = BatchBuilder::new(&schema);
for &(a, k, w) in rows {
b.begin_row_natives(&[a as u128, k as u128], w);
b.put_string(&format!("a string long enough for the heap, number {}", k % 3));
b.put_opt_int((k % 3 != 0).then_some((a * 100 + k as i64) as u128));
b.end_row();
}
b.finish().into_consolidated()
}
#[test]
fn a_join_over_its_traces_source_is_the_join_over_the_trace() {
let source_schema = keyed_source();
let source = keyed_source_batch(&[
(-3, 1, 1),
(-3, 2, 1),
(1, 0, 2),
(1, 4, 1),
(1, 5, -1),
(2, 9, 1),
(7, 3, 1),
(7, 6, 1),
]);
let cases: &[(&[u32], &[u32])] = &[
(&[0], &[1, 2, 3]),
(&[0], &[3, 2]),
(&[0], &[]),
(&[0, 1], &[2, 3]),
(&[0, 1], &[3, 1, 0, 2]),
];
for &(key, keep) in cases {
let mut map = rekey_plan(&source_schema, key, keep);
let trace = map.evaluate_map_batch(&source).into_consolidated();
let trace_schema = *map.out_schema();
let delta_schema = SchemaDescriptor::new(
&key.iter()
.map(|&c| source_schema.columns[c as usize])
.chain([SchemaColumn::new(TypeCode::I64, false)])
.collect::<Vec<_>>(),
&(0..key.len() as u32).collect::<Vec<_>>(),
);
let mut d = BatchBuilder::new(&delta_schema);
for (a, k, w) in [
(-9i64, 0u64, 1i64),
(-3, 2, 2),
(1, 4, -1),
(1, 5, 1),
(5, 5, 1),
(7, 6, 3),
(8, 0, 1),
] {
let pk: Vec<u128> = [a as u128, k as u128][..key.len()].to_vec();
d.begin_row_natives(&pk, w);
d.put_int((a * 7) as u128);
d.end_row();
}
let delta = d.finish().into_consolidated();
for delta_is_right in [false, true] {
let stored = plan(JoinKind::Equi, delta_is_right, &delta_schema, &trace_schema);
let want = op_join_delta_trace(
&delta,
&mut trace_cursor(trace.clone()),
&stored.out_schema,
&stored.probe,
);
let over = JoinPlan::over_source(delta_is_right, &delta_schema, &source_schema, &map).unwrap();
assert!(over.out_schema.same_layout(&stored.out_schema));
assert!(!want.is_empty(), "premise: key {key:?} matches something");
for runs in [1, 3] {
let got = op_join_delta_trace(
&delta,
&mut TestTrace::dealt(&source, runs).cursor(),
&over.out_schema,
&over.probe,
);
assert_eq!(
zset_of(&got, &over.out_schema),
zset_of(&want, &stored.out_schema),
"key {key:?} keep {keep:?} delta_is_right={delta_is_right} runs={runs}"
);
}
}
}
}
#[test]
fn over_source_refuses_a_trace_its_source_does_not_prefix() {
let source = keyed_source();
let col = |c: u32| source.columns[c as usize];
for key in [&[1u32][..], &[1, 0], &[3]] {
let delta = SchemaDescriptor::new(
&key.iter()
.map(|&c| SchemaColumn::new(col(c).type_code, false))
.collect::<Vec<_>>(),
&(0..key.len() as u32).collect::<Vec<_>>(),
);
let map = rekey_plan(&source, key, &[2]);
assert!(
JoinPlan::over_source(false, &delta, &source, &map).is_err(),
"key {key:?}"
);
}
let delta = SchemaDescriptor::new(&[col(0)], &[0]);
assert!(JoinPlan::over_source(false, &delta, &source, &rekey_plan(&source, &[0], &[2])).is_ok());
}