use mongreldb_core::columnar::NativeColumn;
use mongreldb_core::schema::{ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
use mongreldb_core::{Condition, Cursor, Table, Value};
use tempfile::tempdir;
fn schema() -> Schema {
Schema {
schema_id: 1,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "cost".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty(),
},
],
indexes: vec![IndexDef {
name: "cost_bm".into(),
column_id: 2,
kind: IndexKind::Bitmap,
predicate: None,
}],
colocation: vec![],
constraints: Default::default(),
clustered: false,
}
}
fn i64s(col: &NativeColumn) -> Vec<i64> {
match col {
NativeColumn::Int64 { data, .. } => data.clone(),
_ => vec![],
}
}
fn drain_id_cost(
db: &Table,
snap: mongreldb_core::Snapshot,
conditions: &[Condition],
) -> Vec<(i64, i64)> {
let mut cur = db
.native_multi_run_cursor(
snap,
vec![(1, TypeId::Int64), (2, TypeId::Int64)],
conditions,
)
.unwrap()
.expect("multi-run cursor should fire");
let mut out = Vec::new();
while let Some(batch) = cur.next_batch().unwrap() {
let ids = i64s(&batch[0]);
let costs = i64s(&batch[1]);
for (a, b) in ids.iter().zip(costs.iter()) {
out.push((*a, *b));
}
}
out
}
#[test]
fn multi_run_cursor_equals_visible_rows() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
db.set_mutable_run_spill_bytes(1);
let mut run1_ids = Vec::new();
for i in 0..200i64 {
let rid = db
.put(vec![(1, Value::Int64(i)), (2, Value::Int64(i * 10))])
.unwrap();
run1_ids.push(rid);
}
db.flush().unwrap();
for i in 200..400i64 {
db.put(vec![(1, Value::Int64(i)), (2, Value::Int64(i * 10))])
.unwrap();
}
db.flush().unwrap();
db.delete(run1_ids[50]).unwrap();
db.flush().unwrap();
let run2_r250 = {
db.put(vec![(1, Value::Int64(250)), (2, Value::Int64(2500))])
.unwrap();
let pk = db.lookup_pk(&Value::Int64(250).encode_key()).unwrap();
db.delete(pk).unwrap();
pk
};
for i in 400..420i64 {
db.put(vec![(1, Value::Int64(i)), (2, Value::Int64(i * 10))])
.unwrap();
}
db.commit().unwrap();
let _ = run2_r250;
let snap = db.snapshot();
assert!(db.run_count() >= 3, "expected a multi-run layout");
let mut expect: Vec<(i64, i64)> = db
.visible_rows(snap)
.unwrap()
.into_iter()
.map(|r| {
let id = match r.columns.get(&1) {
Some(Value::Int64(v)) => *v,
_ => i64::MIN,
};
let cost = match r.columns.get(&2) {
Some(Value::Int64(v)) => *v,
_ => i64::MIN,
};
(id, cost)
})
.collect();
expect.sort_unstable();
let got = drain_id_cost(&db, snap, &[]);
let mut ordered = got.clone();
ordered.sort_unstable();
assert_eq!(
got, ordered,
"cursor output must be in ascending RowId order"
);
assert_eq!(got, expect, "multi-run cursor must equal visible_rows");
assert!(
got.iter().all(|(id, _)| *id != 50),
"tombstoned id 50 must be dropped"
);
}
#[test]
fn multi_run_cursor_applies_predicate() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
db.set_mutable_run_spill_bytes(1);
for run in 0..3i64 {
for i in 0..100i64 {
let id = run * 1000 + i;
db.put(vec![(1, Value::Int64(id)), (2, Value::Int64(id))])
.unwrap();
}
db.flush().unwrap();
}
assert!(db.run_count() >= 3);
let snap = db.snapshot();
let cond = Condition::BitmapEq {
column_id: 2,
value: Value::Int64(1050).encode_key(), };
let got = drain_id_cost(&db, snap, std::slice::from_ref(&cond));
assert_eq!(
got,
vec![(1050, 1050)],
"predicate survivor in a multi-run layout"
);
let cur = db
.native_multi_run_cursor(snap, vec![(1, TypeId::Int64)], &[cond])
.unwrap()
.unwrap();
assert_eq!(cur.remaining_rows(), 1);
}
#[test]
fn multi_run_cursor_limits_short_circuits() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
db.set_mutable_run_spill_bytes(1);
for run in 0..4i64 {
for i in 0..1000i64 {
let id = run * 10_000 + i;
db.put(vec![(1, Value::Int64(id)), (2, Value::Int64(id))])
.unwrap();
}
db.flush().unwrap();
}
let snap = db.snapshot();
let mut cur = db
.native_multi_run_cursor(snap, vec![(1, TypeId::Int64)], &[])
.unwrap()
.unwrap();
let first = cur.next_batch().unwrap().expect("at least one batch");
let n = i64s(&first.into_iter().next().unwrap()).len();
assert!(
n > 0 && n <= 65_536,
"first batch within the merge page bound: {n}"
);
}