use arrow::array::Array;
use mongreldb_core::schema::{ColumnDef, ColumnFlags, Schema, TypeId};
use mongreldb_core::{Table, Value};
use mongreldb_query::MongrelSession;
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: "destination".into(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty(),
},
ColumnDef {
id: 3,
name: "departure".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty(),
},
ColumnDef {
id: 4,
name: "cost".into(),
ty: TypeId::Float64,
flags: ColumnFlags::empty(),
},
ColumnDef {
id: 5,
name: "rating".into(),
ty: TypeId::Float64,
flags: ColumnFlags::empty(),
},
],
indexes: vec![
mongreldb_core::schema::IndexDef {
name: "dest_bitmap".into(),
column_id: 2,
kind: mongreldb_core::schema::IndexKind::Bitmap,
},
mongreldb_core::schema::IndexDef {
name: "dest_fm".into(),
column_id: 2,
kind: mongreldb_core::schema::IndexKind::FmIndex,
},
],
colocation: vec![],
constraints: Default::default(),
}
}
fn total_rows(batches: &[arrow::record_batch::RecordBatch]) -> usize {
batches.iter().map(|b| b.num_rows()).sum()
}
async fn setup() -> (tempfile::TempDir, MongrelSession) {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
for i in 0..100i64 {
db.put(vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{i}").into_bytes())),
(3, Value::Int64(1_700_000_000 + i * 86_400)),
(4, Value::Float64(199.99 + i as f64)),
(5, Value::Float64(3.5 + (i % 3) as f64)),
])
.unwrap();
}
db.flush().unwrap();
let session = MongrelSession::new(db);
session.register("travel_trips").await.unwrap();
(dir, session)
}
#[tokio::test]
async fn select_star_returns_all_rows() {
let (_dir, session) = setup().await;
let batches = session.run("select * from travel_trips").await.unwrap();
assert_eq!(total_rows(&batches), 100);
let expected_cols = 5;
assert_eq!(
batches[0].schema().fields().len(),
expected_cols,
"select * should expose the user columns"
);
}
#[tokio::test]
async fn count_star_is_correct() {
let (_dir, session) = setup().await;
let batches = session
.run("select count(*) as n from travel_trips")
.await
.unwrap();
assert_eq!(total_rows(&batches), 1);
let col = batches[0].column(0);
let val = col
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
assert_eq!(val.value(0), 100);
}
#[tokio::test]
async fn where_filter_prunes_rows() {
let (_dir, session) = setup().await;
let batches = session
.run("select * from travel_trips where cost < 250")
.await
.unwrap();
assert_eq!(total_rows(&batches), 51);
}
#[tokio::test]
async fn projection_and_limit() {
let (_dir, session) = setup().await;
let batches = session
.run("select destination, cost from travel_trips order by cost desc limit 5")
.await
.unwrap();
assert_eq!(total_rows(&batches), 5);
assert_eq!(batches[0].schema().fields().len(), 2);
}
#[tokio::test]
async fn predicate_pushdown_filtered_query() {
let (_dir, session) = setup().await;
let batches = session
.run("select * from travel_trips where destination = 'City50'")
.await
.unwrap();
assert_eq!(total_rows(&batches), 1, "filtered to exactly one row");
let batches2 = session
.run("select id from travel_trips where id = 42")
.await
.unwrap();
assert_eq!(total_rows(&batches2), 1);
}
#[tokio::test]
async fn result_cache_returns_same_batches_on_repeat() {
let (_dir, session) = setup().await;
let b1 = session
.run("select count(*) as n from travel_trips")
.await
.unwrap();
let b2 = session
.run("select count(*) as n from travel_trips")
.await
.unwrap();
assert_eq!(total_rows(&b1), total_rows(&b2));
session.clear_cache();
let b3 = session
.run("select count(*) as n from travel_trips")
.await
.unwrap();
assert_eq!(total_rows(&b1), total_rows(&b3));
}
#[tokio::test]
async fn aggregation_groups_by() {
let (_dir, session) = setup().await;
let batches = session
.run("select rating, count(*) as n from travel_trips group by rating order by rating")
.await
.unwrap();
assert_eq!(total_rows(&batches), 3);
}
#[tokio::test]
async fn range_int64_pushdown() {
let (_dir, session) = setup().await;
let lo = 1_700_000_000i64;
let hi = 1_700_000_000 + 10 * 86_400;
let sql = format!(
"select id from travel_trips where departure >= {lo} and departure <= {hi} order by id"
);
let batches = session.run(&sql).await.unwrap();
assert_eq!(total_rows(&batches), 11);
let col = batches[0].column(0);
let vals = col
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
assert_eq!(vals.value(0), 0);
assert_eq!(vals.value(vals.len() - 1), 10);
}
#[tokio::test]
async fn range_between_pushdown() {
let (_dir, session) = setup().await;
let lo = 1_700_000_000i64 + 5 * 86_400;
let hi = 1_700_000_000 + 8 * 86_400;
let sql = format!("select id from travel_trips where departure between {lo} and {hi}");
let batches = session.run(&sql).await.unwrap();
assert_eq!(total_rows(&batches), 4);
}
#[tokio::test]
async fn range_float64_pushdown() {
let (_dir, session) = setup().await;
let batches = session
.run("select id from travel_trips where cost < 250 order by id")
.await
.unwrap();
assert_eq!(total_rows(&batches), 51);
let col = batches[0].column(0);
let vals = col
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
assert_eq!(vals.value(vals.len() - 1), 50);
}
#[tokio::test]
async fn like_fm_pushdown() {
let (_dir, session) = setup().await;
let batches = session
.run("select id from travel_trips where destination like 'City5%' order by id")
.await
.unwrap();
assert_eq!(total_rows(&batches), 11);
}
#[tokio::test]
async fn bitmap_intersect_range() {
let (_dir, session) = setup().await;
let sql = "select id from travel_trips \
where destination = 'City5' and departure >= 0 and departure <= 1701000000";
let batches = session.run(sql).await.unwrap();
assert_eq!(total_rows(&batches), 1);
}
#[tokio::test]
async fn or_of_equalities_same_column_unions() {
let (_dir, session) = setup().await;
let two = session
.run("select id from travel_trips where destination = 'City5' or destination = 'City7'")
.await
.unwrap();
assert_eq!(total_rows(&two), 2);
let three = session
.run(
"select id from travel_trips \
where destination = 'City1' or destination = 'City2' or destination = 'City3'",
)
.await
.unwrap();
assert_eq!(total_rows(&three), 3);
}
#[tokio::test]
async fn is_null_and_is_not_null_partition_rows() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
for i in 0..60i64 {
db.put(vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{i}").into_bytes())),
(3, Value::Int64(1_700_000_000 + i)),
(4, Value::Float64(1.0 + i as f64)),
(5, Value::Float64(2.0)),
])
.unwrap();
}
db.commit().unwrap();
db.flush().unwrap();
db.add_column(
"note",
TypeId::Bytes,
ColumnFlags::empty().with(ColumnFlags::NULLABLE),
)
.unwrap();
for i in 60..100i64 {
db.put(vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{i}").into_bytes())),
(3, Value::Int64(1_700_000_000 + i)),
(4, Value::Float64(1.0 + i as f64)),
(5, Value::Float64(2.0)),
(6, Value::Bytes(format!("note{i}").into_bytes())),
])
.unwrap();
}
db.commit().unwrap();
db.flush().unwrap();
let session = MongrelSession::new(db);
session.register("travel_trips").await.unwrap();
let nulls = session
.run("select id from travel_trips where note is null")
.await
.unwrap();
assert_eq!(total_rows(&nulls), 60);
let not_nulls = session
.run("select id from travel_trips where note is not null")
.await
.unwrap();
assert_eq!(total_rows(¬_nulls), 40);
}
fn vec_schema() -> Schema {
Schema {
schema_id: 2,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "vec".into(),
ty: TypeId::Embedding { dim: 8 },
flags: ColumnFlags::empty(),
},
],
indexes: vec![mongreldb_core::schema::IndexDef {
name: "vec_ann".into(),
column_id: 2,
kind: mongreldb_core::schema::IndexKind::Ann,
}],
colocation: vec![],
constraints: Default::default(),
}
}
#[tokio::test]
async fn ann_search_pushdown() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), vec_schema(), 2).unwrap();
let proto = [1.0f32, -1.0, 1.0, 1.0, -1.0, 1.0, 1.0, -1.0];
for i in 0..12i64 {
let mut v = proto;
if i > 0 {
v[((i - 1) as usize) % 8] *= -1.0;
}
db.put(vec![
(1, Value::Int64(i)),
(2, Value::Embedding(v.to_vec())),
])
.unwrap();
}
db.flush().unwrap();
let session = MongrelSession::new(db);
session.register("items").await.unwrap();
let sql = "select id from items where ann_search(vec, '[1,-1,1,1,-1,1,1,-1]', 3) order by id";
let batches = session.run(sql).await.unwrap();
assert_eq!(total_rows(&batches), 3);
let col = batches[0].column(0);
let vals = col
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
let ids: Vec<i64> = vals.values().iter().copied().collect();
assert!(
ids.contains(&0),
"top-k must include the exact match: {ids:?}"
);
}
fn cities_schema() -> Schema {
Schema {
schema_id: 3,
columns: vec![
ColumnDef {
id: 1,
name: "city_name".into(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "country".into(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty(),
},
],
indexes: vec![mongreldb_core::schema::IndexDef {
name: "country_bitmap".into(),
column_id: 2,
kind: mongreldb_core::schema::IndexKind::Bitmap,
}],
colocation: vec![],
constraints: Default::default(),
}
}
#[tokio::test]
async fn multi_table_join() {
let (dir, session) = setup().await;
let cities_dir = tempdir().unwrap();
let mut cities = Table::create(cities_dir.path(), cities_schema(), 3).unwrap();
for i in 0..10i64 {
let country = if i % 2 == 0 { "North" } else { "South" };
cities
.put(vec![
(1, Value::Bytes(format!("City{i}").into_bytes())),
(2, Value::Bytes(country.as_bytes().to_vec())),
])
.unwrap();
}
cities.flush().unwrap();
session.register_db("cities", cities).await.unwrap();
let sql = "select t.destination, c.country \
from travel_trips t join cities c on t.destination = c.city_name \
where c.country = 'North' \
order by t.destination";
let batches = session.run(sql).await.unwrap();
assert_eq!(total_rows(&batches), 5);
let _ = dir;
}
#[tokio::test]
async fn fk_join_composes_pk_and_fk_predicates() {
let (dir, session) = setup().await;
let cities_dir = tempdir().unwrap();
let mut cities = Table::create(cities_dir.path(), cities_schema(), 3).unwrap();
for i in 0..10i64 {
let country = if i % 2 == 0 { "North" } else { "South" };
cities
.put(vec![
(1, Value::Bytes(format!("City{i}").into_bytes())),
(2, Value::Bytes(country.as_bytes().to_vec())),
])
.unwrap();
}
cities.flush().unwrap();
session.register_db("cities", cities).await.unwrap();
let sql = "select t.destination, c.country \
from travel_trips t join cities c on t.destination = c.city_name \
where c.country = 'North' and t.departure >= 1700432000 \
order by t.destination";
let batches = session.run(sql).await.unwrap();
assert_eq!(
total_rows(&batches),
2,
"only City6 and City8 pass both filters"
);
let arr = batches[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::StringArray>()
.unwrap();
assert_eq!(arr.value(0), "City6");
let _ = dir;
}
#[tokio::test]
async fn streaming_scan_emits_multiple_batches() {
let dir = tempdir().unwrap();
let minimal = Schema {
schema_id: 2,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "v".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty(),
},
],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
};
let mut db = Table::create(dir.path(), minimal, 1).unwrap();
let n: i64 = 65_536 + 1000;
for i in 0..n {
db.put(vec![(1, Value::Int64(i)), (2, Value::Int64(i * 2))])
.unwrap();
}
db.flush().unwrap();
let session = MongrelSession::new(db);
session.register("nums").await.unwrap();
let batches = session.run("select * from nums").await.unwrap();
assert_eq!(total_rows(&batches), n as usize);
assert!(
batches.len() >= 2,
"expected multiple streamed batches, got {}",
batches.len()
);
assert!(batches.iter().all(|b| b.num_rows() <= 65_536));
for b in &batches[..batches.len() - 1] {
assert_eq!(b.num_rows(), 65_536);
}
let c = session.run("select count(*) as n from nums").await.unwrap();
let col = c[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
assert_eq!(col.value(0), n);
let lim = session.run("select * from nums limit 10").await.unwrap();
assert_eq!(total_rows(&lim), 10);
let half = session
.run("select * from nums where v < 1000")
.await
.unwrap();
assert_eq!(total_rows(&half), 500);
let _ = dir;
}
#[tokio::test]
async fn multi_run_streams_and_limit_short_circuits() {
use mongreldb_core::schema::{IndexDef, IndexKind};
let dir = tempdir().unwrap();
let multi_schema = Schema {
schema_id: 9,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "v".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty(),
},
],
indexes: vec![IndexDef {
name: "v_bm".into(),
column_id: 2,
kind: IndexKind::Bitmap,
}],
colocation: vec![],
constraints: Default::default(),
};
let mut db = Table::create(dir.path(), multi_schema, 1).unwrap();
db.set_mutable_run_spill_bytes(1); for run in 0..3i64 {
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();
}
assert!(db.run_count() >= 3, "multi-run layout");
let session = MongrelSession::new(db);
session.register("nums").await.unwrap();
let all = session.run("select * from nums").await.unwrap();
assert_eq!(total_rows(&all), 3000);
let lim = session.run("select * from nums limit 10").await.unwrap();
assert_eq!(total_rows(&lim), 10);
let one = session
.run("select * from nums where v = 10500")
.await
.unwrap();
assert_eq!(total_rows(&one), 1);
let _ = dir;
}
#[tokio::test]
async fn cursor_page_pruning_is_exact() {
let dir = tempdir().unwrap();
let minimal = Schema {
schema_id: 3,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "v".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty(),
},
],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
};
let mut db = Table::create(dir.path(), minimal, 1).unwrap();
let n: i64 = 65_536 * 2 + 5000;
for i in 0..n {
db.put(vec![(1, Value::Int64(i)), (2, Value::Int64(i))])
.unwrap();
}
db.flush().unwrap();
let session = MongrelSession::new(db);
session.register("nums").await.unwrap();
let third = session
.run("select * from nums where v >= 131000 and v < 131500")
.await
.unwrap();
assert_eq!(total_rows(&third), 500);
let span = session
.run("select * from nums where v >= 65500 and v < 65600")
.await
.unwrap();
assert_eq!(total_rows(&span), 100);
let none = session
.run("select * from nums where v > 9999999")
.await
.unwrap();
assert_eq!(total_rows(&none), 0);
let c = session.run("select count(*) as n from nums").await.unwrap();
let col = c[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
assert_eq!(col.value(0), n);
let _ = dir;
}
#[tokio::test]
async fn zero_copy_preserves_nulls() {
let dir = tempdir().unwrap();
let nullable = Schema {
schema_id: 4,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
ColumnDef {
id: 2,
name: "v".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
};
let mut db = Table::create(dir.path(), nullable, 1).unwrap();
for i in 0..10i64 {
if i % 2 == 0 {
db.put(vec![(1, Value::Int64(i)), (2, Value::Int64(i * 10))])
.unwrap();
} else {
db.put(vec![(1, Value::Int64(i))]).unwrap();
}
}
db.flush().unwrap();
let session = MongrelSession::new(db);
session.register("nums").await.unwrap();
let batches = session.run("select v from nums").await.unwrap();
assert_eq!(total_rows(&batches), 10);
let mut pos = 0;
for b in &batches {
let arr = b
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap();
for j in 0..arr.len() {
let expected_null = pos % 2 == 1;
assert_eq!(
arr.is_null(j),
expected_null,
"pos {pos} null mismatch (expected {expected_null})"
);
if !expected_null {
assert_eq!(arr.value(j), pos * 10);
}
pos += 1;
}
}
let _ = dir;
}
#[tokio::test]
async fn cursor_handles_schema_evolution() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
for i in 0..100i64 {
db.put(vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{i}").into_bytes())),
(3, Value::Int64(1_700_000_000 + i * 86_400)),
(4, Value::Float64(199.99 + i as f64)),
(5, Value::Float64(3.5 + (i % 3) as f64)),
])
.unwrap();
}
db.flush().unwrap();
db.add_column(
"note",
TypeId::Bytes,
ColumnFlags::empty().with(ColumnFlags::NULLABLE),
)
.unwrap();
let session = MongrelSession::new(db);
session.register("travel_trips").await.unwrap();
let batches = session.run("select note from travel_trips").await.unwrap();
assert_eq!(total_rows(&batches), 100);
for b in &batches {
let arr = b.column(0);
assert_eq!(arr.null_count(), arr.len());
}
let all = session.run("select * from travel_trips").await.unwrap();
assert_eq!(total_rows(&all), 100);
assert_eq!(all[0].schema().fields().len(), 6);
let _ = dir;
}
#[tokio::test]
async fn metadata_aggregates_count_min_max() {
let (_dir, session) = setup().await; let c = session
.run("select count(*) as n from travel_trips")
.await
.unwrap();
let n = c[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
assert_eq!(n, 100);
let mm = session
.run("select min(id) as mn, max(id) as mx from travel_trips")
.await
.unwrap();
let mn = mm[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
let mx = mm[0]
.column(1)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
assert_eq!(mn, 0);
assert_eq!(mx, 99);
let mixed = session
.run("select min(id) as aggregate_min, max(id) as aggregate_max, min(3, 1, 2) as scalar_min, max(3, 1, 2) as scalar_max from travel_trips")
.await
.unwrap();
let aggregate_min = mixed[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
let aggregate_max = mixed[0]
.column(1)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
let scalar_min = mixed[0]
.column(2)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
let scalar_max = mixed[0]
.column(3)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
assert_eq!(aggregate_min, 0);
assert_eq!(aggregate_max, 99);
assert_eq!(scalar_min, 1);
assert_eq!(scalar_max, 3);
}
#[tokio::test]
async fn metadata_aggregates_count_col_and_null_column() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
let rows: Vec<Vec<(u16, Value)>> = (0..100i64)
.map(|i| {
vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{i}").into_bytes())),
(3, Value::Int64(1_700_000_000 + i)),
(4, Value::Float64(10.0 + i as f64)),
(5, Value::Float64(2.0)),
]
})
.collect();
db.bulk_load(rows).unwrap(); db.add_column(
"empty_i",
TypeId::Int64,
ColumnFlags::empty().with(ColumnFlags::NULLABLE),
)
.unwrap(); let session = MongrelSession::new(db);
session.register("travel_trips").await.unwrap();
assert_eq!(
i64_val(&session, "select count(id) as c from travel_trips").await,
100
);
assert_eq!(
i64_val(&session, "select count(empty_i) as c from travel_trips").await,
0
);
assert_eq!(
i64_val(&session, "select min(departure) as m from travel_trips").await,
1_700_000_000
);
assert_eq!(
i64_val(&session, "select max(departure) as m from travel_trips").await,
1_700_000_099
);
assert!(
(f64_val(&session, "select min(cost) as m from travel_trips").await - 10.0).abs() < 1e-9
);
assert!(
(f64_val(&session, "select max(cost) as m from travel_trips").await - 109.0).abs() < 1e-9
);
let nb = session
.run("select min(empty_i) as m from travel_trips")
.await
.unwrap();
assert!(
nb[0].column(0).is_null(0),
"MIN over an all-NULL column must be SQL NULL"
);
}
#[tokio::test]
async fn count_distinct_from_bitmap_partition() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
let rows: Vec<Vec<(u16, Value)>> = (0..100i64)
.map(|i| {
vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{}", i % 5).into_bytes())),
(3, Value::Int64(1_700_000_000 + i)),
(4, Value::Float64(1.0 + i as f64)),
(5, Value::Float64(2.0)),
]
})
.collect();
db.bulk_load(rows).unwrap();
let session = MongrelSession::new(db);
session.register("travel_trips").await.unwrap();
assert_eq!(
i64_val(
&session,
"select count(distinct destination) as c from travel_trips"
)
.await,
5
);
assert_eq!(
i64_val(&session, "select count(distinct id) as c from travel_trips").await,
100
);
}
#[tokio::test]
async fn count_col_excludes_schema_evolution_nulls_on_scan() {
let dir = tempdir().unwrap();
let mut db = Table::create(dir.path(), schema(), 1).unwrap();
for i in 0..40i64 {
db.put(vec![
(1, Value::Int64(i)),
(2, Value::Bytes(format!("City{i}").into_bytes())),
(3, Value::Int64(1_700_000_000 + i)),
(4, Value::Float64(10.0 + i as f64)),
(5, Value::Float64(2.0)),
])
.unwrap();
}
db.commit().unwrap();
db.flush().unwrap(); db.add_column(
"empty_i",
TypeId::Int64,
ColumnFlags::empty().with(ColumnFlags::NULLABLE),
)
.unwrap();
let session = MongrelSession::new(db);
session.register("travel_trips").await.unwrap();
assert_eq!(
i64_val(&session, "select count(empty_i) as c from travel_trips").await,
0
);
assert_eq!(
i64_val(&session, "select count(*) as c from travel_trips").await,
40
);
}
#[tokio::test]
async fn metadata_aggregates_skip_the_scan() {
let (_dir, session) = setup().await;
let explained = session
.run("explain select count(*) as n from travel_trips")
.await
.unwrap();
let plan: String = (0..explained[0].num_rows())
.flat_map(|r| {
explained[0]
.column(1)
.as_any()
.downcast_ref::<arrow::array::StringArray>()
.unwrap()
.value(r)
.lines()
.map(str::to_owned)
.collect::<Vec<_>>()
})
.collect::<Vec<_>>()
.join("\n");
assert!(
!plan.contains("MongrelScanExec"),
"COUNT(*) should be served from stats, not a scan. Plan:\n{plan}"
);
}
#[tokio::test]
async fn explain_query_plan_returns_compatibility_shape() {
let (_dir, session) = setup().await;
let explained = session
.run(
"explain query plan
select count(*) as n
from travel_trips
where destination = 'City 7'
group by destination
order by destination",
)
.await
.unwrap();
assert_eq!(explained[0].schema().field(0).name(), "id");
assert_eq!(explained[0].schema().field(1).name(), "parent");
assert_eq!(explained[0].schema().field(2).name(), "notused");
assert_eq!(explained[0].schema().field(3).name(), "detail");
let detail = explained[0]
.column(3)
.as_any()
.downcast_ref::<arrow::array::StringArray>()
.unwrap();
assert!(
(0..detail.len()).any(|idx| detail.value(idx) == "SEARCH travel_trips USING MONGREL INDEX"),
"expected MongrelDB SEARCH detail"
);
assert!(
(0..detail.len()).any(|idx| detail.value(idx) == "USE TEMP B-TREE FOR GROUP BY"),
"expected GROUP BY detail"
);
assert!(
(0..detail.len()).any(|idx| detail.value(idx) == "USE TEMP B-TREE FOR ORDER BY"),
"expected ORDER BY detail"
);
assert!(
(0..detail.len()).any(|idx| detail.value(idx).contains("logical_plan")
|| detail.value(idx).contains("physical_plan")),
"expected DataFusion plan detail"
);
}
#[tokio::test]
async fn native_aggregates_sum_min_max_avg_count() {
let (_dir, session) = setup().await;
let sum_id = i64_val(&session, "select sum(id) as s from travel_trips").await;
assert_eq!(sum_id, (0..100i64).sum::<i64>());
let min_id = i64_val(&session, "select min(id) as m from travel_trips").await;
assert_eq!(min_id, 0);
let max_id = i64_val(&session, "select max(id) as m from travel_trips").await;
assert_eq!(max_id, 99);
let cnt = i64_val(&session, "select count(*) as c from travel_trips").await;
assert_eq!(cnt, 100);
let avg_id = f64_val(&session, "select avg(id) as a from travel_trips").await;
assert!((avg_id - 49.5).abs() < 1e-9);
let sum_cost = f64_val(&session, "select sum(cost) as s from travel_trips").await;
let expected_cost: f64 = (0..100).map(|i| 199.99 + i as f64).sum();
assert!((sum_cost - expected_cost).abs() < 1e-6);
let sum_filt = i64_val(
&session,
"select sum(id) as s from travel_trips where id < 10",
)
.await;
assert_eq!(sum_filt, 45);
let cnt_filt = i64_val(
&session,
"select count(*) as c from travel_trips where id < 10",
)
.await;
assert_eq!(cnt_filt, 10);
}
async fn i64_val(session: &MongrelSession, sql: &str) -> i64 {
let b = session.run(sql).await.unwrap();
b[0].column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0)
}
async fn f64_val(session: &MongrelSession, sql: &str) -> f64 {
let b = session.run(sql).await.unwrap();
b[0].column(0)
.as_any()
.downcast_ref::<arrow::array::Float64Array>()
.unwrap()
.value(0)
}
#[tokio::test]
async fn native_aggregate_rejects_like_filter() {
let (_dir, session) = setup().await; let n = session
.run("select count(*) as c from travel_trips where destination like '%City_1%'")
.await
.unwrap();
let c = n[0]
.column(0)
.as_any()
.downcast_ref::<arrow::array::Int64Array>()
.unwrap()
.value(0);
assert_eq!(c, 9, "LIKE wildcard semantics must be applied exactly");
}
#[tokio::test]
async fn direct_dispatch_equality_and_range_match_datafusion() {
let (_dir, session) = setup().await;
let (b, trace) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE id = 7")
.await
.unwrap();
assert_eq!(total_rows(&b), 1);
assert_eq!(
trace.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch,
"simple single-table eq SELECT should direct-dispatch"
);
let (b, trace) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE id >= 90")
.await
.unwrap();
assert_eq!(total_rows(&b), 10);
assert_eq!(
trace.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch
);
assert_eq!(b[0].schema().fields().len(), 1);
assert_eq!(b[0].schema().field(0).name(), "id");
let (b, trace) = session
.run_sql_traced("SELECT * FROM travel_trips WHERE id < 3")
.await
.unwrap();
assert_eq!(total_rows(&b), 3);
assert_eq!(
trace.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch
);
assert_eq!(
b[0].schema().fields().len(),
5,
"SELECT * projects all columns"
);
}
#[tokio::test]
async fn direct_dispatch_bitmap_equality_on_indexed_column() {
let (_dir, session) = setup().await;
let (b, trace) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE destination = 'City5'")
.await
.unwrap();
assert_eq!(total_rows(&b), 1);
assert_eq!(
trace.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch
);
}
#[tokio::test]
async fn direct_dispatch_in_list_and_null_fall_through_correctly() {
let (_dir, session) = setup().await;
let (b, _t) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE destination IN ('City1','City2')")
.await
.unwrap();
assert_eq!(total_rows(&b), 2);
let (b, _t) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE destination IS NOT NULL")
.await
.unwrap();
assert_eq!(total_rows(&b), 100);
}
#[tokio::test]
async fn like_and_limit_fall_through_to_datafusion() {
let (_dir, session) = setup().await;
let (b, trace) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE destination LIKE 'City1%'")
.await
.unwrap();
assert_eq!(total_rows(&b), 11);
assert_ne!(
trace.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch,
"LIKE must fall through (inexact)"
);
let (b, trace) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE id = 5 LIMIT 1")
.await
.unwrap();
assert_eq!(total_rows(&b), 1);
assert_ne!(
trace.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch
);
}
#[tokio::test]
async fn direct_dispatch_is_memoized_by_result_cache() {
let (_dir, session) = setup().await;
let (b1, t1) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE id = 4")
.await
.unwrap();
assert_eq!(total_rows(&b1), 1);
assert_eq!(
t1.scan_mode,
mongreldb_core::trace::ScanMode::DirectDispatch
);
let (b2, t2) = session
.run_sql_traced("SELECT id FROM travel_trips WHERE id = 4")
.await
.unwrap();
assert_eq!(total_rows(&b2), 1);
assert_eq!(t2.scan_mode, mongreldb_core::trace::ScanMode::Unknown);
}
#[tokio::test]
async fn recursive_cte_works() {
let (_tmp, session) = setup().await;
let sql = "WITH RECURSIVE chain AS (
SELECT id, 0 AS depth FROM travel_trips WHERE id = 1
UNION ALL
SELECT t.id, c.depth + 1 FROM travel_trips t JOIN chain c ON t.id = c.id + 1
) SELECT id, depth FROM chain ORDER BY id";
let batches = session.run(sql).await.unwrap();
let rows = total_rows(&batches);
assert!(rows > 0, "recursive CTE should return rows, got {rows}");
}
#[tokio::test]
async fn window_function_works() {
let (_tmp, session) = setup().await;
let sql = "SELECT id, ROW_NUMBER() OVER (ORDER BY id) AS rn FROM travel_trips";
let batches = session.run(sql).await.unwrap();
let rows = total_rows(&batches);
assert_eq!(rows, 100, "window function should return 100 rows");
}
#[tokio::test]
async fn regexp_function_works() {
let (_tmp, session) = setup().await;
let sql = "SELECT id FROM travel_trips WHERE regexp('^City[0-5]$', destination) = 1";
let batches = session.run(sql).await.unwrap();
let rows = total_rows(&batches);
assert_eq!(rows, 6, "regexp should match City0-City5, got {rows}");
}
#[tokio::test]
async fn information_schema_lists_tables() {
let (_tmp, session) = setup().await;
let batches = session
.run("SELECT type, name FROM information_schema.tables ORDER BY name")
.await
.unwrap();
let rows = total_rows(&batches);
assert!(rows >= 1, "information_schema.tables should list the travel_trips table, got {rows}");
}
#[tokio::test]
async fn attach_database_enables_cross_db_query() {
let dir2 = tempdir().unwrap();
{
let db2 = mongreldb_core::Database::create(dir2.path()).unwrap();
db2.create_table("items", schema()).unwrap();
let handle = db2.table("items").unwrap();
let mut g = handle.lock();
g.put(vec![
(1, Value::Int64(999)),
(2, Value::Bytes(b"attached".to_vec())),
(3, Value::Int64(1_700_000_000)),
(4, Value::Float64(42.0)),
(5, Value::Float64(1.0)),
])
.unwrap();
g.flush().unwrap();
}
let dir1 = tempdir().unwrap();
let db1 = std::sync::Arc::new(mongreldb_core::Database::create(dir1.path()).unwrap());
let session = MongrelSession::open(db1).unwrap();
let attach_sql = format!("ATTACH '{}' AS other", dir2.path().display());
session.run(&attach_sql).await.unwrap();
let batches = session
.run("SELECT id FROM other_items")
.await
.unwrap();
let rows = total_rows(&batches);
assert_eq!(rows, 1, "attached table should have 1 row, got {rows}");
}
#[tokio::test]
async fn savepoint_syntax_is_accepted() {
let (_tmp, session) = setup().await;
session.run("SAVEPOINT sp1").await.unwrap();
session.run("RELEASE sp1").await.unwrap();
session.run("SAVEPOINT sp2").await.unwrap();
session.run("ROLLBACK TO sp2").await.unwrap();
}