use super::*;
use crate::repr::{Batch, BatchBuilder};
use crate::schema::{ColumnTable, SchemaColumn, TypeCode};
use crate::test_support::{arb_fold_case, cell, fold_batch, fold_schemas, le_cell, opk_pk, pk_only_schema};
use proptest::prelude::*;
#[test]
fn single_natural_col_keys_by_its_opk_from_either_side() {
let key = |tc: TypeCode, le: &[u8], col1_is_pk: bool| -> (u128, Vec<u8>) {
let cols = [SchemaColumn::new(TypeCode::U64, false), SchemaColumn::new(tc, false)];
let schema = SchemaDescriptor::new(&cols, if col1_is_pk { &[0, 1] } else { &[0] });
let mut b = BatchBuilder::new(&schema);
match col1_is_pk {
true => b.begin_row_natives(&[0, le_cell(le)], 1),
false => {
b.begin_row(0, 1);
b.put_int(le_cell(le));
}
}
b.end_row();
let b = b.finish();
let (key, _) = GroupOutKey::new(&schema, &[1], []).unwrap();
let mb = b.as_mem_batch();
(key.identity(&mb, 0), key.out_pk(&mb, 0).bytes().to_vec())
};
for (tc, vals) in [
(TypeCode::I32, vec![1i128, -1, 100, i32::MIN as i128, i32::MAX as i128]),
(TypeCode::I64, vec![0, -1, i64::MIN as i128, i64::MAX as i128]),
(TypeCode::U16, vec![0, 1, 0xBEEF, u16::MAX as i128]),
(TypeCode::U64, vec![0, 1, u64::MAX as i128]),
] {
let width = SchemaColumn::new(tc, false).size() as usize;
for v in vals {
let le = &(v as u128).to_le_bytes()[..width];
let want_pk = opk_pk(&pk_only_schema(&[tc]), &[v as u128]);
let want = (gnitz_wire::widen_pk_be(&want_pk), want_pk);
assert_eq!(key(tc, le, true), want, "PK-column key for {tc} v={v}");
assert_eq!(key(tc, le, false), want, "payload-column key for {tc} v={v}");
}
}
}
#[test]
fn sorted_ordinals_number_groups_by_key_and_keep_the_first_row() {
let groups = GroupOrdinals::sorted(12, |i| 2 - i % 3);
assert_eq!(groups.ord, [2, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1, 0]);
assert_eq!(groups.first, [2, 1, 0]);
assert_eq!(groups.by_pk, [0, 1, 2]);
let groups = GroupOrdinals::sorted(5, |_| 0u64);
assert_eq!((groups.ord, groups.first, groups.by_pk), (vec![0; 5], vec![0], vec![0]));
}
fn assert_groups_follow_out_pk(raw: &Batch) -> Result<(), TestCaseError> {
let schema = *raw.schema();
let consolidated = Batch::clone(raw).into_consolidated();
let pk: Vec<u32> = schema.pk_cols().to_vec();
let all: Vec<u32> = (0..schema.num_columns() as u32).collect();
let n = schema.num_columns() as u32;
let mut forms = vec![vec![], all];
forms.extend((1..=pk.len()).map(|lead| pk[..lead].to_vec()));
forms.extend((0..n).map(|c| vec![c]));
forms.extend((0..n).flat_map(|a| (0..n).filter(move |&b| b != a).map(move |b| vec![a, b])));
for cols in &forms {
let (key, _) = GroupOutKey::new(&schema, cols, []).unwrap();
for batch in [raw, &consolidated] {
let mb = batch.as_mem_batch();
let out_pk = |row: usize| key.out_pk(&mb, row).bytes().to_vec();
let groups = key.ordinals(batch);
prop_assert_eq!(groups.ord.len(), batch.count, "{:?}", cols);
for (row, &g) in groups.ord.iter().enumerate() {
let first = groups.first[g as usize] as usize;
prop_assert!(first <= row, "{:?}: row {} precedes its group's first", cols, row);
prop_assert_eq!(out_pk(first), out_pk(row), "{:?}: a group mixes keys", cols);
}
let mut seen = groups.by_pk.clone();
seen.sort_unstable();
prop_assert_eq!(seen, (0..groups.len() as u32).collect::<Vec<_>>(), "{:?}", cols);
let pks: Vec<Vec<u8>> = groups
.by_pk
.iter()
.map(|&g| out_pk(groups.first[g as usize] as usize))
.collect();
prop_assert!(pks.windows(2).all(|w| w[0] < w[1]), "{:?}: groups out of order", cols);
let runs = key.runs(batch);
if let Some(runs) = &runs {
prop_assert_eq!(runs.len(), groups.len(), "{:?}", cols);
for (g, run) in runs.iter().enumerate() {
prop_assert!(
run.clone().all(|row| groups.ord[row] == g as u32),
"{:?}: run {} is not group {}",
cols,
g,
g
);
}
}
let bytes: usize = cols.iter().map(|&c| schema.columns[c as usize].size() as usize).sum();
let in_place = *cols == pk || bytes <= GROUP_PK_BYTES;
if std::ptr::eq(batch, &consolidated) && !cols.is_empty() && pk.starts_with(cols) && in_place {
prop_assert!(runs.is_some(), "{:?}: a PK prefix of a consolidated batch", cols);
}
let group =
|r: usize| -> Vec<_> { cols.iter().map(|&c| cell(&mb, schema.locate(c as usize), r)).collect() };
for a in 0..batch.count {
for b in a + 1..batch.count {
let same_out_pk = out_pk(a) == out_pk(b);
prop_assert_eq!(group(a) == group(b), same_out_pk, "{:?}: rows {} and {}", cols, a, b);
prop_assert_eq!(
groups.ord[a] == groups.ord[b],
same_out_pk,
"{:?}: rows {} and {}",
cols,
a,
b
);
prop_assert_eq!(
key.identity(&mb, a) == key.identity(&mb, b),
same_out_pk,
"{:?}: rows {} and {}",
cols,
a,
b
);
}
}
}
}
Ok(())
}
proptest! {
#[test]
fn groups_follow_out_pk_at_every_pk_width((si, rows) in arb_fold_case()) {
assert_groups_follow_out_pk(&fold_batch(&fold_schemas()[si], &rows))?;
}
}
#[test]
fn groups_follow_out_pk_over_signed_wide_and_nullable_columns() {
let schema = SchemaDescriptor::new(
&[
SchemaColumn::new(TypeCode::U32, false),
SchemaColumn::new(TypeCode::I32, false),
SchemaColumn::new(TypeCode::I64, false),
SchemaColumn::new(TypeCode::U128, false),
SchemaColumn::new(TypeCode::I64, true),
],
&[0, 1],
);
let mut bb = BatchBuilder::new(&schema);
for i in 0..64u64 {
let m = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 40;
bb.begin_row_natives(&[(m % 5) as u128, (m % 7) as i32 as i64 as u128], 1);
bb.put_int(((m % 9) as i64 - 4) as u128);
bb.put_int(u128::from(m % 3) << 100);
match m % 4 {
0 => bb.put_null(),
k => bb.put_int((k as i64 - 2) as u128),
}
bb.end_row();
}
assert_groups_follow_out_pk(&bb.finish()).unwrap();
}