crates/knut-bifrost/tests/iceberg_pushdown.rs
//! **Iceberg read pushdown** — projection + predicate, end to end.
//!
//! knut's Iceberg source used to call `skade::Table::read()` for every
//! `SyncMode::Full` round: **every column of every data file**, no matter how
//! narrow the mapping was. skade has carried the pruning surface all along
//! (`read_columns` pushes the projection into the row-group readers,
//! `read_filtered` pushes a predicate into the scan planner, `plan_stats` costs a
//! plan without reading it) and knut called none of it.
//!
//! These tests assert the three halves of the fix:
//!
//! 1. the projection is DERIVED from the mapping pipeline, so an existing mapping
//! gets the narrowing with no config change;
//! 2. a projected read returns only those columns and **byte-identical values**
//! for them — pruning must not change data, only cost;
//! 3. a filtered read SKIPS data files, proven at plan time (`plan_stats`), not
//! inferred from wall time.
//!
//! Everything runs against a local embedded warehouse in a tempdir — no
//! container, no network.
#![cfg(feature = "skade")]
use std::sync::Arc;
use knut_bifrost::mapping::{MappingSpec, RowFilter, SyncMode};
use knut_bifrost::source::{IcebergSource, TableSource, TableSpec};
use knut_bifrost::state::TableState;
use skade::arrow_array::{Array, Int64Array, RecordBatch as SkadeBatch, StringArray};
use skade::arrow_schema::{DataType, Field, Schema};
/// A WIDE row shape: the id + name a graph mapping actually needs, plus eight
/// payload columns it never looks at. That width is the whole point — it is what
/// the unprojected read was paying for.
fn wide_batch(country: &str, ids: std::ops::Range<i64>) -> SkadeBatch {
let ids: Vec<i64> = ids.collect();
let n = ids.len();
let names: Vec<String> = ids.iter().map(|i| format!("person_{i}")).collect();
let countries: Vec<String> = std::iter::repeat_n(country.to_string(), n).collect();
// Eight fat, never-read payload columns.
let payloads: Vec<Vec<String>> = (0..8)
.map(|c| {
ids.iter()
.map(|i| format!("payload_{c}_{i}_{}", "x".repeat(64)))
.collect()
})
.collect();
let mut fields = vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("country", DataType::Utf8, false),
];
let mut cols: Vec<skade::arrow_array::ArrayRef> = vec![
Arc::new(Int64Array::from(ids)),
Arc::new(StringArray::from(names)),
Arc::new(StringArray::from(countries)),
];
for (c, p) in payloads.into_iter().enumerate() {
fields.push(Field::new(format!("payload_{c}"), DataType::Utf8, false));
cols.push(Arc::new(StringArray::from(p)));
}
SkadeBatch::try_new(Arc::new(Schema::new(fields)), cols).expect("wide batch")
}
/// A warehouse whose `people` table holds ONE data file per country, so a
/// `country == …` predicate prunes at file granularity.
async fn seed(dir: &std::path::Path, countries: &[&str], per_country: i64) {
let wh = skade::Warehouse::open(dir).await.expect("open warehouse");
let first = wide_batch(countries[0], 0..per_country);
let mut table = wh
.table_or_create("people", &first.schema())
.await
.expect("create table");
for (i, c) in countries.iter().enumerate() {
let lo = i as i64 * per_country;
let batch = wide_batch(c, lo..lo + per_country);
// One append per country → one data file per country.
table.append(&[batch]).await.expect("append");
}
}
const MAPPING: &str = r#"
version: 1
tables:
- table: people
mode: full
pipeline:
- vertex:
label: Person
id_column: id
properties: [name]
"#;
// ═══════════════════════════════════════════════════════════════════════════
// 1 — the projection is DERIVED from the mapping
// ═══════════════════════════════════════════════════════════════════════════
/// A `vertex` step names its inputs exactly (`id_column` + `properties`), so the
/// projection falls out of the spec with no new config. RED-when-broken: widen
/// the derivation (or lose a column) and the assert names it.
#[test]
fn projection_is_derived_from_the_pipeline() {
let spec = MappingSpec::from_yaml(MAPPING).expect("parse mapping");
let cols = spec.tables[0]
.projected_columns()
.expect("a declarative pipeline projects");
assert_eq!(cols, vec!["id".to_string(), "name".to_string()]);
// Edges contribute both endpoint columns plus their properties, de-duplicated
// against the vertex columns, first-seen order preserved.
let with_edges = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: t
pipeline:
- vertex: { label: Person, id_column: id, properties: [name] }
- fk_edges:
type: WORKS_AT
src: { label: Person, column: id }
dst: { label: Company, column: company_id }
properties: [since]
"#,
)
.expect("parse");
assert_eq!(
with_edges.tables[0].projected_columns().expect("projects"),
vec![
"id".to_string(),
"name".to_string(),
"company_id".to_string(),
"since".to_string()
],
"`id` appears once even though vertex and edge both name it"
);
// A `custom` step is opaque — it may read anything, so projecting is refused
// and the source falls back to a full-width read. Narrowing an opaque
// transform's input would starve it; that must stay impossible.
let opaque = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: t
pipeline:
- vertex: { label: Person, id_column: id, properties: [name] }
- custom: my_transform
"#,
)
.expect("parse");
assert!(
opaque.tables[0].projected_columns().is_none(),
"a custom transform must disable the projection, not narrow its input"
);
// An explicit override wins over the derivation.
let override_spec = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: t
columns: [id, name, country]
pipeline:
- vertex: { label: Person, id_column: id, properties: [name] }
"#,
)
.expect("parse");
assert_eq!(
override_spec.tables[0]
.projected_columns()
.expect("override projects"),
vec!["id".to_string(), "name".to_string(), "country".to_string()],
);
}
// ═══════════════════════════════════════════════════════════════════════════
// 2 — a projected read is NARROWER and byte-identical on the columns it keeps
// ═══════════════════════════════════════════════════════════════════════════
/// The projected read must return exactly the asked-for columns, and their values
/// must equal what the full-width read produced. This is the correctness gate on
/// the whole optimisation: pruning changes cost, never data.
#[tokio::test]
async fn projected_read_is_narrower_and_byte_identical_on_the_kept_columns() {
let dir = tempfile::tempdir().expect("tempdir");
seed(dir.path(), &["SE", "NO", "DK", "FI"], 500).await;
let src = IcebergSource::new(dir.path());
let latest = src
.latest_snapshot("people")
.await
.expect("snapshot")
.expect("has snapshots");
// Baseline: the pre-fix behaviour — no projection, no filter.
let full_spec = TableSpec::new("people", SyncMode::Full);
let (full, _) = src
.fetch_new(&full_spec, &TableState::default())
.await
.expect("full read")
.expect("news");
// The fix: the same round, projected to what the mapping reads.
let proj_spec =
TableSpec::new("people", SyncMode::Full).with_columns(vec!["id".into(), "name".into()]);
let (proj, state) = src
.fetch_new(&proj_spec, &TableState::default())
.await
.expect("projected read")
.expect("news");
assert_eq!(state.snapshot_id, Some(latest), "the cursor still advances");
let full_rows: usize = full.iter().map(|b| b.num_rows()).sum();
let proj_rows: usize = proj.iter().map(|b| b.num_rows()).sum();
assert_eq!(full_rows, 2000, "4 countries x 500");
assert_eq!(
proj_rows, full_rows,
"the projection drops COLUMNS, never rows"
);
let full_width = full[0].num_columns();
let proj_width = proj[0].num_columns();
assert_eq!(full_width, 11, "id + name + country + 8 payloads");
assert_eq!(proj_width, 2, "only the two columns the mapping reads");
assert!(
proj_width < full_width,
"the projected read really is narrower"
);
// Same schema names, in the asked order.
let proj_schema = proj[0].schema();
let names: Vec<&str> = proj_schema
.fields()
.iter()
.map(|f| f.name().as_str())
.collect();
assert_eq!(names, vec!["id", "name"]);
// BYTE-IDENTICAL values: concatenate both reads' `id`/`name` and compare.
let take =
|batches: &[skade::arrow_array::RecordBatch], col: &str| -> (Vec<i64>, Vec<String>) {
let _ = col;
let mut ids = Vec::new();
let mut nms = Vec::new();
for b in batches {
let i = b
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let n = b
.column_by_name("name")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
for k in 0..b.num_rows() {
ids.push(i.value(k));
nms.push(n.value(k).to_string());
}
}
(ids, nms)
};
let (mut fid, mut fnm) = take(&full, "id");
let (mut pid, mut pnm) = take(&proj, "id");
// Scan order across data files is not part of the contract; compare as sets.
fid.sort_unstable();
pid.sort_unstable();
fnm.sort();
pnm.sort();
assert_eq!(
pid, fid,
"projected `id` values are identical to the full read's"
);
assert_eq!(
pnm, fnm,
"projected `name` values are identical to the full read's"
);
}
// ═══════════════════════════════════════════════════════════════════════════
// 3 — a filtered read SKIPS data files, proven at plan time
// ═══════════════════════════════════════════════════════════════════════════
/// `plan_stats` costs a scan WITHOUT reading it, so the data-skipping win is a
/// measured number rather than a wall-clock inference. RED-when-broken: unwire
/// the filter and `files_planned` equals the baseline — no skipping — and this
/// fails.
#[tokio::test]
async fn filtered_plan_skips_data_files_and_rows() {
let dir = tempfile::tempdir().expect("tempdir");
seed(dir.path(), &["SE", "NO", "DK", "FI"], 500).await;
let src = IcebergSource::new(dir.path());
let baseline = src.plan_stats("people", None).await.expect("baseline plan");
assert_eq!(baseline.data_files, 4, "one data file per country append");
assert_eq!(baseline.rows_planned, 2000);
let one = RowFilter::eq("country", "SE");
let pruned = src
.plan_stats("people", Some(&one))
.await
.expect("filtered plan");
assert!(
pruned.data_files < baseline.data_files,
"the predicate must prune data files: {} planned vs {} baseline",
pruned.data_files,
baseline.data_files
);
assert_eq!(
pruned.data_files, 1,
"only the SE file survives `country == SE`"
);
assert_eq!(pruned.rows_planned, 500);
// A set predicate prunes to the union.
let two = RowFilter::is_in("country", ["SE", "NO"]);
let pruned2 = src.plan_stats("people", Some(&two)).await.expect("IN plan");
assert_eq!(
pruned2.data_files, 2,
"`country IN (SE, NO)` keeps two files"
);
assert_eq!(pruned2.rows_planned, 1000);
}
/// The filtered READ returns only matching rows, and combines with the
/// projection: `read_filtered(f, cols)` is one call doing both.
#[tokio::test]
async fn filtered_read_returns_only_matching_rows_and_projects() {
let dir = tempfile::tempdir().expect("tempdir");
seed(dir.path(), &["SE", "NO", "DK", "FI"], 500).await;
let src = IcebergSource::new(dir.path());
let spec = TableSpec::new("people", SyncMode::Full)
.with_columns(vec!["id".into(), "name".into(), "country".into()])
.with_filter(Some(RowFilter::eq("country", "SE")));
let (batches, _) = src
.fetch_new(&spec, &TableState::default())
.await
.expect("filtered read")
.expect("news");
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, 500, "file pruning left exactly the SE data file");
assert_eq!(
batches[0].num_columns(),
3,
"projection applied alongside the filter"
);
// Every surviving row really is SE (this layout prunes exactly; skade
// documents pruning as file-granular, so a mixed file could carry extras).
for b in &batches {
let c = b
.column_by_name("country")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
for k in 0..b.num_rows() {
assert_eq!(c.value(k), "SE");
}
}
}
/// An empty projection and no filter must reproduce the OLD behaviour exactly —
/// the pushdown is strictly additive. RED-when-broken: if the source started
/// projecting by default, this width assert fails.
#[tokio::test]
async fn no_projection_and_no_filter_is_the_unchanged_full_read() {
let dir = tempfile::tempdir().expect("tempdir");
seed(dir.path(), &["SE", "NO"], 100).await;
let src = IcebergSource::new(dir.path());
let (batches, _) = src
.fetch_new(
&TableSpec::new("people", SyncMode::Full),
&TableState::default(),
)
.await
.expect("read")
.expect("news");
assert_eq!(batches[0].num_columns(), 11, "every column, as before");
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 200);
}
/// A filter that names no form would push down a predicate that prunes nothing —
/// a silent no-op. It must be rejected, not quietly ignored.
#[tokio::test]
async fn a_formless_filter_is_rejected_not_silently_ignored() {
let dir = tempfile::tempdir().expect("tempdir");
seed(dir.path(), &["SE"], 10).await;
let formless = RowFilter {
column: "country".into(),
eq: None,
is_in: None,
ge: None,
le: None,
};
assert!(formless.validate().is_err(), "a formless filter is invalid");
let src = IcebergSource::new(dir.path());
let err = src
.plan_stats("people", Some(&formless))
.await
.expect_err("a formless filter must not plan");
assert!(
err.to_string().contains("prune nothing"),
"the error says why: {err}"
);
}
/// The YAML surface round-trips: a mapping can declare both halves of the
/// pushdown and they reach `TableSpec` unchanged.
#[test]
fn mapping_yaml_carries_columns_and_filter() {
let spec = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: people
mode: full
columns: [id, name]
filter: { column: country, eq: SE }
pipeline:
- vertex: { label: Person, id_column: id, properties: [name] }
"#,
)
.expect("parse mapping with pushdown");
let t = &spec.tables[0];
assert_eq!(t.columns, vec!["id".to_string(), "name".to_string()]);
let f = t.filter.clone().expect("filter parsed");
assert_eq!(f, RowFilter::eq("country", "SE"));
f.validate().expect("a declared eq filter is valid");
// Range form.
let ranged = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: t
filter: { column: year, ge: 2020, le: 2026 }
pipeline:
- vertex: { label: X, id_column: id }
"#,
)
.expect("parse range filter");
let f = ranged.tables[0].filter.clone().expect("filter");
f.validate().expect("a range filter is valid");
assert_eq!(f, RowFilter::range("year", Some(2020i64), Some(2026i64)));
}
// ═══════════════════════════════════════════════════════════════════════════
// 4 — the predicate must not be silently dropped by the OTHER sync modes
// ═══════════════════════════════════════════════════════════════════════════
/// **A mapping's `filter` is a row PREDICATE, and it must mean the same thing in
/// every `SyncMode`.**
///
/// The pushdown wiring routes only `SyncMode::Full` through `read_snapshot` (the
/// function that reads `spec.filter`). `Incremental` and `Cdc` go to
/// `delta_from_table`, which takes neither `spec.columns` nor `spec.filter` —
/// so the predicate is silently a no-op there and rows the operator excluded are
/// ingested anyway.
///
/// The projection being dropped is only a cost regression (a superset of columns
/// is harmless — every transform projects what it names). The FILTER being
/// dropped is a CORRECTNESS regression: nothing downstream re-applies it. The
/// `Listener` hands `compiled.filter` to the source and never filters a row
/// itself (`listener.rs:244`), so whatever the source returns is what lands in
/// the graph.
///
/// The first round of an `Incremental` table has no cursor and reads the whole
/// table — exactly the shape asserted here, and exactly when a `filter` matters
/// most, because it is the round that decides the bulk of the graph.
///
/// RED-when-broken: this is red today. It goes green when the delta path honours
/// `spec.filter` (or when `fetch_new` refuses a filter it cannot push, which is
/// also acceptable — a loud refusal beats a silent superset).
#[tokio::test]
async fn a_filter_is_honoured_in_incremental_mode_too_not_only_in_full() {
let dir = tempfile::tempdir().unwrap();
seed(dir.path(), &["SE", "NO", "DK", "FI"], 250).await;
let src = IcebergSource::new(dir.path());
let only_se = RowFilter::eq("country", "SE");
// The control: in `Full` mode the predicate IS honoured — 250 of 1000 rows.
let full = TableSpec::new("people", SyncMode::Full).with_filter(Some(only_se.clone()));
let (batches, _) = src
.fetch_new(&full, &TableState::default())
.await
.expect("full read")
.expect("a fresh table has a snapshot");
let full_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(
full_rows, 250,
"Full mode honours the predicate (the control)"
);
// The same mapping, the same predicate, the only difference being the mode.
let incr = TableSpec::new("people", SyncMode::Incremental).with_filter(Some(only_se));
let (batches, _) = src
.fetch_new(&incr, &TableState::default())
.await
.expect("incremental read")
.expect("a fresh table has a snapshot");
let incr_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(
incr_rows,
full_rows,
"the SAME `filter` must select the SAME rows in Incremental as in Full — got \
{incr_rows} rows under Incremental vs {full_rows} under Full, i.e. the predicate \
was silently dropped and {} extra rows would land in the graph",
incr_rows.saturating_sub(full_rows)
);
}