use crate::helpers::{make_ctx, payload_value, send_ok};
use nodedb_physical::physical_plan::{
AggregateSpec, ColumnarInsertIntent, ColumnarOp, GroupKeySpec, PhysicalPlan, QueryOp,
};
use nodedb_types::config::tuning::QueryTuning;
fn insert_grouped_columnar(
ctx: &mut crate::helpers::TestCtx,
collection: &str,
total: usize,
groups: usize,
) {
let rows: Vec<serde_json::Value> = (0..total)
.map(|i| {
serde_json::json!({
"id": format!("r{i}"),
"g": format!("g{}", i % groups),
"v": i,
})
})
.collect();
let payload = nodedb_types::json_to_msgpack(&serde_json::Value::Array(rows)).unwrap();
send_ok(
&mut ctx.core,
&mut ctx.tx,
&mut ctx.rx,
PhysicalPlan::Columnar(ColumnarOp::Insert {
collection: collection.into(),
payload,
format: "msgpack".into(),
intent: ColumnarInsertIntent::Insert,
on_conflict_updates: Vec::new(),
surrogates: Vec::new(),
schema_bytes: Vec::new(),
provenance: None,
wal_lsn: None,
}),
);
}
#[test]
fn limit_honoured_when_groups_span_multiple_aggregate_chunks() {
let mut ctx = make_ctx();
ctx.core.set_query_tuning(QueryTuning {
aggregate_chunk_size: 2,
..QueryTuning::default()
});
const GROUPS: usize = 10;
const ROWS_PER_GROUP: usize = 3;
const TOTAL: usize = GROUPS * ROWS_PER_GROUP;
insert_grouped_columnar(&mut ctx, "chunk_limit_col", TOTAL, GROUPS);
const LIMIT: usize = 3;
let payload = send_ok(
&mut ctx.core,
&mut ctx.tx,
&mut ctx.rx,
PhysicalPlan::Query(QueryOp::Aggregate {
collection: "chunk_limit_col".into(),
input: None,
group_by: vec![GroupKeySpec::column("g")],
aggregates: vec![AggregateSpec {
function: "count".into(),
alias: "count(*)".into(),
user_alias: None,
field: "*".into(),
expr: None,
}],
filters: Vec::new(),
having: Vec::new(),
limit: LIMIT,
sub_group_by: Vec::new(),
sub_aggregates: Vec::new(),
grouping_sets: Vec::new(),
sort_keys: Vec::new(),
}),
);
let result = payload_value(&payload);
let rows = result
.as_array()
.unwrap_or_else(|| panic!("expected aggregate rows, got {result}"));
assert!(
rows.len() <= LIMIT,
"LIMIT {LIMIT} must cap the GROUP BY result; got {} rows. \
A result of {GROUPS} would indicate LIMIT was silently ignored \
when groups span multiple aggregate_chunk_size chunks.",
rows.len()
);
assert_eq!(
rows.len(),
LIMIT,
"expected exactly {LIMIT} rows (LIMIT honoured); got {}",
rows.len()
);
}
#[test]
fn no_limit_returns_all_groups_with_small_chunk_size() {
let mut ctx = make_ctx();
ctx.core.set_query_tuning(QueryTuning {
aggregate_chunk_size: 2,
..QueryTuning::default()
});
const GROUPS: usize = 10;
insert_grouped_columnar(&mut ctx, "chunk_nolimit_col", 30, GROUPS);
let payload = send_ok(
&mut ctx.core,
&mut ctx.tx,
&mut ctx.rx,
PhysicalPlan::Query(QueryOp::Aggregate {
collection: "chunk_nolimit_col".into(),
input: None,
group_by: vec![GroupKeySpec::column("g")],
aggregates: vec![AggregateSpec {
function: "count".into(),
alias: "count(*)".into(),
user_alias: None,
field: "*".into(),
expr: None,
}],
filters: Vec::new(),
having: Vec::new(),
limit: usize::MAX,
sub_group_by: Vec::new(),
sub_aggregates: Vec::new(),
grouping_sets: Vec::new(),
sort_keys: Vec::new(),
}),
);
let result = payload_value(&payload);
let rows = result
.as_array()
.unwrap_or_else(|| panic!("expected aggregate rows, got {result}"));
assert_eq!(
rows.len(),
GROUPS,
"no-LIMIT GROUP BY with small chunk_size must return all {GROUPS} groups; got {}",
rows.len()
);
}