crates/knut-thund/tests/native_projection_pushdown.rs
//! Integration proof for the native backend's **typed projection / filter
//! pushdown** — the introspectable query-optimization node.
//!
//! Today a flow's transform is an opaque SQL `query` string the engine cannot
//! prune columns or push filters from. The typed [`Flow::with_projection`] /
//! [`Flow::with_filter`] nodes are *structured*: the native backend lowers them
//! to a DataFusion scan that reads **only the selected columns** (column
//! pruning pushed into the source) and **only the matching rows** (filter
//! predicate pushed toward the source). This test proves both against a real
//! parquet file on disk, asserting the pushdown from the **physical plan** (the
//! scan node's `projection=[…]`) — not just the output rows.
//!
//! Gated on `--features native` (the DataFusion tree); a no-op otherwise.
#![cfg(feature = "native")]
use knut_thund::backend::{ExecBackend, native::NativeBackend};
use knut_thund::ir::{Dataset, Flow, OutputType, Pipeline};
use datafusion::arrow::array::{Array, Int64Array, RecordBatch, StringArray, StringViewArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::parquet::arrow::ArrowWriter;
use std::sync::Arc;
/// Read a `Utf8`/`Utf8View` string column by name into owned `String`s.
fn str_col(b: &RecordBatch, name: &str) -> Vec<String> {
let col = b
.column_by_name(name)
.unwrap_or_else(|| panic!("column `{name}` present"))
.as_any();
if let Some(a) = col.downcast_ref::<StringArray>() {
(0..a.len()).map(|i| a.value(i).to_string()).collect()
} else if let Some(a) = col.downcast_ref::<StringViewArray>() {
(0..a.len()).map(|i| a.value(i).to_string()).collect()
} else {
panic!("column `{name}` is neither Utf8 nor Utf8View");
}
}
/// Read an `Int64` column by name.
fn i64_col(b: &RecordBatch, name: &str) -> Vec<i64> {
let a = b
.column_by_name(name)
.unwrap_or_else(|| panic!("column `{name}` present"))
.as_any()
.downcast_ref::<Int64Array>()
.unwrap_or_else(|| panic!("column `{name}` is Int64"));
(0..a.len()).map(|i| a.value(i)).collect()
}
/// Write a 4-column `orders.parquet` fixture: `id, customer, amount, note`.
/// Amounts are chosen so a `amount > 100` filter keeps exactly two rows.
fn write_orders_parquet(path: &std::path::Path) {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("customer", DataType::Utf8, false),
Field::new("amount", DataType::Int64, false),
Field::new("note", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5, 6])),
Arc::new(StringArray::from(vec![
"alice", "bob", "carol", "dave", "erin", "frank",
])),
Arc::new(Int64Array::from(vec![10, 5, 200, 7, 350, 3])),
Arc::new(StringArray::from(vec!["n1", "n2", "n3", "n4", "n5", "n6"])),
],
)
.expect("build orders batch");
let file = std::fs::File::create(path).expect("create parquet");
let mut w = ArrowWriter::try_new(file, schema, None).expect("arrow writer");
w.write(&batch).expect("write parquet batch");
w.close().expect("close parquet");
}
/// A **typed projection** flow (no SQL) reads ONLY the two selected columns off
/// disk. RED-when-broken on two axes:
/// 1. output schema carries exactly `[customer, amount]` (the pruned `id` /
/// `note` never surface) — drop `select_columns` in the backend and the
/// full 4-column schema fails this;
/// 2. the optimized **physical plan** shows `projection=[customer, amount]`
/// on the scan and never mentions `note`/`id` — proving the pruning was
/// pushed into the source, not applied after a full read.
#[test]
fn projection_prunes_columns_at_scan() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("orders.parquet");
write_orders_parquet(&path);
let pipeline = Pipeline::new("prune")
.with_dataset(Dataset::new("slim", OutputType::MaterializedView))
.with_flow(
// No `.with_query(...)` — the typed projection IS the transform.
Flow::batch("proj", "slim", ["orders"]).with_projection(["customer", "amount"]),
);
let run = NativeBackend::new()
.with_file_input("orders", path.to_str().unwrap(), "parquet")
.run(&pipeline)
.expect("native run over the on-disk parquet completes");
// (1) Output schema is pruned to exactly the selected columns, in order.
let out = run.output("slim").expect("materialised output");
let names: Vec<String> = out[0]
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
assert_eq!(
names,
vec!["customer".to_string(), "amount".to_string()],
"the output carries ONLY the projected columns (id/note pruned)"
);
// (2) The pushdown is proven from the physical plan: the scan lists exactly
// the two projected columns and the pruned columns never appear.
let pd = run
.pushdown("proj")
.expect("a pushdown was recorded for the projection flow");
assert_eq!(
pd.projected_columns,
vec!["customer".to_string(), "amount".to_string()]
);
assert!(
pd.physical_plan.contains("projection=[customer, amount]"),
"the scan node shows the pushed-down projection; plan was:\n{}",
pd.physical_plan
);
assert!(
!pd.physical_plan.contains("note"),
"the pruned `note` column is never read; plan was:\n{}",
pd.physical_plan
);
knut_thund::functional_status(
"knut-thund/native_projection_pushdown",
"projection_prunes_columns",
names == vec!["customer".to_string(), "amount".to_string()]
&& pd.physical_plan.contains("projection=[customer, amount]"),
"typed projection -> parquet scan reads only [customer, amount]",
);
}
/// A **typed filter** flow drops exactly the non-matching rows. RED-when-broken:
/// drop the filter application in the backend and all six rows survive, failing
/// the exact-row assertion. The predicate also reaches the plan (proving it is
/// carried toward the scan, not discarded).
#[test]
fn filter_drops_nonmatching_rows() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("orders.parquet");
write_orders_parquet(&path);
let pipeline = Pipeline::new("big_orders")
.with_dataset(Dataset::new("big", OutputType::MaterializedView))
.with_flow(Flow::batch("flt", "big", ["orders"]).with_filter("amount > 100"));
let run = NativeBackend::new()
.with_file_input("orders", path.to_str().unwrap(), "parquet")
.run(&pipeline)
.expect("native run completes");
let out = run.output("big").expect("materialised output");
// Exactly the two rows with amount > 100: carol=200, erin=350.
let mut got: Vec<(String, i64)> = Vec::new();
for b in out {
let cust = str_col(b, "customer");
let amt = i64_col(b, "amount");
for i in 0..b.num_rows() {
got.push((cust[i].clone(), amt[i]));
}
}
got.sort();
assert_eq!(
got,
vec![("carol".to_string(), 200), ("erin".to_string(), 350)],
"exactly the rows matching `amount > 100` survive"
);
let pd = run
.pushdown("flt")
.expect("a pushdown was recorded for the filter flow");
assert_eq!(pd.filter.as_deref(), Some("amount > 100"));
// The predicate is PUSHED INTO the parquet scan (row-group pruning), not
// merely applied as a post-scan filter over a full read.
assert!(
pd.physical_plan.contains("predicate=") && pd.physical_plan.contains("amount"),
"the filter predicate reached the source scan; plan was:\n{}",
pd.physical_plan
);
knut_thund::functional_status(
"knut-thund/native_projection_pushdown",
"filter_drops_rows",
got == vec![("carol".to_string(), 200), ("erin".to_string(), 350)],
"typed filter `amount > 100` -> only matching rows materialise",
);
}
/// Projection **and** filter compose on one flow: only the selected columns of
/// only the matching rows. Proves the two typed nodes stack (filter first over
/// all columns, then the projection prunes).
#[test]
fn projection_and_filter_compose() {
let dir = tempfile::tempdir().expect("tempdir");
let path = dir.path().join("orders.parquet");
write_orders_parquet(&path);
let pipeline = Pipeline::new("slim_big")
.with_dataset(Dataset::new("slim_big", OutputType::MaterializedView))
.with_flow(
Flow::batch("pf", "slim_big", ["orders"])
.with_filter("amount > 100")
.with_projection(["customer"]),
);
let run = NativeBackend::new()
.with_file_input("orders", path.to_str().unwrap(), "parquet")
.run(&pipeline)
.expect("native run completes");
let out = run.output("slim_big").expect("materialised output");
let names: Vec<String> = out[0]
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
assert_eq!(
names,
vec!["customer".to_string()],
"projection pruned to `customer` only"
);
let mut custs: Vec<String> = out.iter().flat_map(|b| str_col(b, "customer")).collect();
custs.sort();
assert_eq!(
custs,
vec!["carol".to_string(), "erin".to_string()],
"only the `amount > 100` customers survive, projected to `customer`"
);
// The scan reads exactly the columns it needs — `customer` (projected out)
// and `amount` (needed by the pushed-down predicate) — while `id`/`note`
// are pruned at the source. The final projection to just `customer` happens
// above the scan; the point proven here is the source prunes id/note.
let pd = run.pushdown("pf").expect("pushdown recorded");
assert!(
pd.physical_plan.contains("projection=[customer, amount]")
&& pd.physical_plan.contains("predicate="),
"scan pruned to only the needed columns with the filter pushed down; plan was:\n{}",
pd.physical_plan
);
assert!(
!pd.physical_plan.contains("note") && !pd.physical_plan.contains(" id"),
"the unused `id`/`note` columns are pruned at the source; plan was:\n{}",
pd.physical_plan
);
knut_thund::functional_status(
"knut-thund/native_projection_pushdown",
"projection_and_filter_compose",
names == vec!["customer".to_string()]
&& custs == vec!["carol".to_string(), "erin".to_string()],
"typed filter + projection compose into one pruned + filtered scan",
);
}
/// A typed projection/filter flow with no explicit query **lowers to Spark** via
/// the synthesized `SELECT … WHERE …` (so Catalyst does its own pruning +
/// pushdown). This proves the Spark half of the node without a live Spark.
/// RED-when-broken: if `to_sdp` dropped the synthesized query the SDP flow would
/// carry `None` and the assert fails.
#[test]
fn typed_pushdown_lowers_to_spark_select() {
let flow = Flow::batch("pf", "slim_big", ["orders"])
.with_filter("amount > 100")
.with_projection(["customer", "amount"]);
// The IR-level synthesis is what to_sdp lowers with.
assert_eq!(
flow.synth_query().as_deref(),
Some("SELECT customer, amount FROM orders WHERE amount > 100"),
);
let sdp = flow.to_sdp();
assert_eq!(
sdp.query.as_deref(),
Some("SELECT customer, amount FROM orders WHERE amount > 100"),
"the typed projection/filter lowers to a pushdown SELECT for Spark",
);
knut_thund::functional_status(
"knut-thund/native_projection_pushdown",
"lowers_to_spark_select",
sdp.query.as_deref() == Some("SELECT customer, amount FROM orders WHERE amount > 100"),
"typed projection/filter -> synthesized SELECT for the Spark backend",
);
}