crates/knut-bifrost/tests/arrow_source.rs
//! **GATE 5 — an Arrow source really reads a non-Spark warehouse.**
//!
//! ## What already existed, and what did not
//!
//! bifröst was never Spark-only. Measured before writing a line:
//!
//! | source | reaches | already there |
//! |---|---|---|
//! | `SparkConnectSource` (`spark`) | Iceberg via a Spark Connect endpoint | yes |
//! | `IcebergSource` (`skade`) | a LOCAL embedded skade warehouse **directory** | yes |
//! | `FileSource` | a watched directory of parquet/orc/csv/json/arrow drops | yes |
//! | `rel2graph::ArrowTables` + `build_mutation` | Arrow batches you already hold — **offline**, no listener, no sink, no state ledger | yes |
//!
//! So "only the two Spark ones can feed bifröst" was false, and korp's
//! `Cargo.toml` note about `knut-bifrost/skade` being "the read-half" is about
//! `IcebergSource`, a genuinely non-Spark Iceberg reader.
//!
//! **The two things that were missing** are what this file covers:
//!
//! 1. **No LIVE Arrow source.** `ArrowTables` fed the offline `build_mutation`, but
//! nothing let an Arrow producer drive the `Listener` round loop — the sink, the
//! state ledger, the changelog split, the asset events. A warehouse reached any
//! way other than "a Spark endpoint" or "a local warehouse directory" — a remote
//! skade over its Iceberg-REST front, an Iceberg REST catalog, a client living in
//! the calling process — could not drive a sync at all. `ArrowSource` is that
//! source, and it is built ON `ArrowTables`, not beside it.
//! 2. **An empty read looked exactly like a successful build.** `build_mutation`
//! silently `continue`s past a spec table the source has no rows for, so pointing
//! it at the wrong namespace returns `Ok` with an empty `GraphMutation`.
//! `BuildReport::require_non_empty` and `EmptyRead::Refuse` are the
//! discrimination.
//!
//! ## What this does NOT give you
//!
//! **Change-data-capture.** `SparkIcebergCdc`, `SparkNativeCdc`,
//! `watch_iceberg_delta` and `IcebergSource`'s incremental path all rest on a
//! snapshot lineage the engine can diff. A bag of Arrow batches has none, so
//! `SyncMode::Cdc` over an `ArrowSource` is a typed refusal, asserted below. skade
//! reached this way and iceberg-rest feed bifröst for a **full build**, not for
//! streaming.
use std::sync::{Arc, Mutex};
use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use async_trait::async_trait;
use knut_bifrost::mapping::SyncMode;
use knut_bifrost::source::{ArrowBatches, ArrowSource, EmptyRead, TableSource, TableSpec};
use knut_bifrost::state::TableState;
use knut_bifrost::{
ArrowTables, GraphMutation, Listener, MappingSpec, SinkReport, SyncState, TransformRegistry,
build_mutation, build_mutation_reported, sink::GraphSink,
};
const MAPPING: &str = r#"
version: 1
tables:
- table: lake.hr.people
mode: full
pipeline:
- vertex:
label: Person
id_column: id
properties: [name]
"#;
/// `id` + `name`, the shape the mapping reads.
fn people(ids: std::ops::RangeInclusive<i64>) -> RecordBatch {
let ids: Vec<i64> = ids.collect();
let names: Vec<String> = ids.iter().map(|i| format!("person_{i}")).collect();
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(ids)),
Arc::new(StringArray::from(names)),
],
)
.expect("people batch")
}
/// An in-memory graph sink — the "ingested" side, no FalkorDB needed.
#[derive(Clone, Default)]
struct RecordingSink {
vertices: Arc<Mutex<usize>>,
labels: Arc<Mutex<Vec<String>>>,
}
#[async_trait]
impl GraphSink for RecordingSink {
async fn apply(&mut self, mutation: &GraphMutation) -> knut_bifrost::Result<SinkReport> {
let v: usize = mutation.vertices.iter().map(|s| s.batch.num_rows()).sum();
*self.vertices.lock().unwrap() += v;
self.labels
.lock()
.unwrap()
.extend(mutation.vertices.iter().map(|s| s.label.clone()));
Ok(SinkReport {
vertices_written: v,
edges_written: mutation.edges.iter().map(|s| s.batch.num_rows()).sum(),
deleted: mutation.deletes.iter().map(|s| s.batch.num_rows()).sum(),
})
}
}
fn listener_over(source: ArrowSource, sink: RecordingSink, dir: &std::path::Path) -> Listener<ArrowSource, RecordingSink> {
Listener::new(
MappingSpec::from_yaml(MAPPING).expect("mapping"),
TransformRegistry::new(),
source,
sink,
SyncState::open(dir.join("state.json")).expect("state"),
)
}
// ═══════════════════════════════════════════════════════════════════════════
// 1 — Arrow batches drive the LIVE listener
// ═══════════════════════════════════════════════════════════════════════════
/// Batches a caller already holds run the whole round loop: mapping → transform →
/// sink → state ledger. No Spark endpoint, no warehouse directory, no files.
///
/// RED when broken: return `Ok(None)` from `ArrowSource::fetch_new` for a non-empty
/// read and `round.synced` is empty while `round.failed` stays empty too — the
/// row-count assertion is what catches a source that "succeeds" without shipping.
#[tokio::test]
async fn arrow_batches_drive_the_live_listener() {
const N: i64 = 137; // a known, non-round count asserted end to end.
let sink = RecordingSink::default();
let dir = tempfile::tempdir().unwrap();
let source = ArrowSource::from_batches("in-memory", "lake.hr.people", vec![people(1..=N)]);
assert_eq!(source.producer(), "in-memory");
assert_eq!(source.empty_policy(), EmptyRead::Refuse, "refusing is the default");
let mut listener = listener_over(source, sink.clone(), dir.path());
let round = listener.sync_once().await.expect("a round over Arrow batches");
assert_eq!(round.failed, vec![], "no sync failures");
assert_eq!(round.synced.len(), 1, "the one mapped table synced");
assert_eq!(round.synced[0].0, "lake.hr.people");
assert_eq!(round.synced[0].1, N as usize, "every injected row shipped");
assert_eq!(
*sink.vertices.lock().unwrap(),
N as usize,
"the sink received one vertex row per input row"
);
assert_eq!(sink.labels.lock().unwrap().as_slice(), ["Person"]);
}
/// A `SyncMode::Full` Arrow source re-reads every round — the idempotent-MERGE
/// contract every other full-mode source has.
#[tokio::test]
async fn full_mode_re_reads_every_round() {
let sink = RecordingSink::default();
let dir = tempfile::tempdir().unwrap();
let source = ArrowSource::from_batches("in-memory", "lake.hr.people", vec![people(1..=10)]);
let mut listener = listener_over(source, sink.clone(), dir.path());
listener.sync_once().await.expect("round 1");
listener.sync_once().await.expect("round 2");
assert_eq!(
*sink.vertices.lock().unwrap(),
20,
"full mode ships the same 10 rows twice; the sink's MERGE makes that safe"
);
}
// ═══════════════════════════════════════════════════════════════════════════
// 2 — THE GATE: an empty read is not a built graph
// ═══════════════════════════════════════════════════════════════════════════
/// **An empty batch must not silently produce an empty graph that looks like
/// success.**
///
/// Three ways it used to look like success, all closed here:
///
/// * zero batches from the producer;
/// * a batch set whose batches carry zero ROWS (the shape a schema-only read gives);
/// * the wrong table name, which reads as "empty" from the producer's side.
///
/// RED when broken: change `EmptyRead::Refuse` to yield `Ok(None)` and every
/// `expect_err` below flips — the round then reports `Ok`, `synced` is empty, and
/// nothing anywhere says the graph was not built.
#[tokio::test]
async fn an_empty_arrow_read_is_refused_not_a_built_graph() {
let spec = TableSpec::new("lake.hr.people", SyncMode::Full);
let fresh = TableState::default();
// (a) zero batches.
let none = ArrowSource::from_batches("skade-remote", "lake.hr.people", vec![]);
let e = none
.fetch_new(&spec, &fresh)
.await
.expect_err("a zero-batch read must not look like a successful empty build");
let msg = e.to_string();
assert!(msg.contains("lake.hr.people"), "the refusal names the table: {msg}");
assert!(msg.contains("skade-remote"), "and the producer: {msg}");
assert!(msg.contains("0 rows"), "and what it actually got: {msg}");
// (b) batches present, zero rows in them — the schema-only read.
let empty_rows =
ArrowSource::from_batches("iceberg-rest", "lake.hr.people", vec![people(1..=0)]);
let e = empty_rows
.fetch_new(&spec, &fresh)
.await
.expect_err("a zero-ROW batch is still an empty read");
assert!(e.to_string().contains("1 batch(es) and 0 rows"), "{e}");
// (c) the wrong table — and the error says what the producer DOES hold, which is
// the difference between a five-minute and a five-second diagnosis.
let wrong = ArrowSource::new(
"skade-remote",
ArrowTables::new().with_table("lake.hr.People", vec![people(1..=3)]),
);
let e = wrong
.fetch_new(&spec, &fresh)
.await
.expect_err("a table the producer does not hold is not an empty table");
assert!(
e.to_string().contains("lake.hr.People"),
"the refusal lists the tables the producer holds: {e}"
);
// And the opt-out is explicit, never the default.
let allowed = ArrowSource::from_batches("skade-remote", "lake.hr.people", vec![])
.allowing_empty();
assert_eq!(allowed.empty_policy(), EmptyRead::Allow);
assert!(
allowed.fetch_new(&spec, &fresh).await.unwrap().is_none(),
"a declared-empty table is Ok(None) — nothing to build, and nobody claimed there was"
);
}
/// The same hazard at the OFFLINE seam: `build_mutation` returns `Ok` with an empty
/// mutation when the source carries nothing, and `BuildReport` is what tells them
/// apart.
///
/// RED when broken: make `require_non_empty` return `Ok(self)` unconditionally and
/// the `expect_err` fires; drop the `rows == 0` arm from `build_mutation_reported`
/// and `tables_empty` no longer names the table.
#[test]
fn an_empty_offline_build_is_reported_not_silently_ok() {
let spec = MappingSpec::from_yaml(MAPPING).unwrap();
// The source carries a DIFFERENT table — the classic wrong-namespace mistake.
let wrong = ArrowTables::new().with_table("lake.hr.persons", vec![people(1..=50)]);
let mutation = build_mutation(&spec, &wrong).expect("the old API still returns Ok");
assert!(
mutation.vertices.is_empty(),
"this is the shape the gate exists for: Ok, and an empty graph"
);
let (_, report) = build_mutation_reported(&spec, &wrong).unwrap();
assert_eq!(report.tables_mapped, Vec::<String>::new());
assert_eq!(report.tables_empty, vec!["lake.hr.people".to_string()]);
assert_eq!((report.rows_read, report.vertices, report.edges), (0, 0, 0));
let e = report
.require_non_empty()
.expect_err("an empty build must be refusable in one line");
assert!(e.to_string().contains("lake.hr.people"), "{e}");
assert!(e.to_string().contains("EMPTY graph"), "{e}");
// The right table: mapped, counted, and it passes the same check.
let right = ArrowTables::new().with_table("lake.hr.people", vec![people(1..=50)]);
let (m, report) = build_mutation_reported(&spec, &right).unwrap();
assert_eq!(report.tables_mapped, vec!["lake.hr.people".to_string()]);
assert_eq!(report.tables_empty, Vec::<String>::new());
assert_eq!(report.rows_read, 50);
assert_eq!(report.vertices, 50);
assert_eq!(m.vertices[0].label, "Person");
report.require_non_empty().expect("a real build passes");
}
// ═══════════════════════════════════════════════════════════════════════════
// 3 — DO NOT OVERCLAIM: this is a batch read, not CDC
// ═══════════════════════════════════════════════════════════════════════════
/// `SyncMode::Cdc` over an Arrow source is a typed refusal naming the reason and the
/// sources that CAN do it — never an insert-only batch the listener would treat as
/// the whole changelog.
///
/// RED when broken: delete the guard and the ask succeeds, quietly turning a
/// changelog round into a full re-read whose deletes never happen.
#[tokio::test]
async fn cdc_over_an_arrow_source_is_refused_naming_why() {
let source = ArrowSource::from_batches("iceberg-rest", "lake.hr.people", vec![people(1..=5)]);
let e = source
.fetch_new(
&TableSpec::new("lake.hr.people", SyncMode::Cdc),
&TableState::default(),
)
.await
.expect_err("an Arrow batch source has no change feed");
let msg = e.to_string();
assert!(msg.contains("no change feed"), "{msg}");
assert!(msg.contains("SparkIcebergCdc"), "it names what can: {msg}");
assert!(msg.contains("IcebergSource"), "and the pure-Rust one: {msg}");
// And it advertises no CDC capability at all — the same honest `None` the file
// source gives, so an operator reading `state_json` is not misled either.
assert!(source.cdc_capability("lake.hr.people").is_none());
// Incremental is allowed but ships ONCE: there is no delta to compute.
let spec = TableSpec::new("lake.hr.people", SyncMode::Incremental);
let first = source.fetch_new(&spec, &TableState::default()).await.unwrap();
assert!(first.is_some(), "the first incremental round ships the batches");
let synced = TableState {
synced_at: Some("2026-09-01T00:00:00Z".into()),
..Default::default()
};
assert!(
source.fetch_new(&spec, &synced).await.unwrap().is_none(),
"a second incremental round has nothing new — a batch source is not a stream"
);
}
// ═══════════════════════════════════════════════════════════════════════════
// 4 — a REAL non-Spark warehouse: skade
// ═══════════════════════════════════════════════════════════════════════════
/// **The gate proper: feed bifröst from skade and build a graph.**
///
/// A real embedded skade Iceberg warehouse in a tempdir — no Spark, no JVM, no
/// container. Its rows are read out as Arrow, handed to `ArrowSource`, and driven
/// through the live listener to a sink. This is the path every korp `Warehouse` impl
/// can now take, because Arrow is what they all already produce.
///
/// RED when broken: seed the warehouse with zero rows and the round fails with the
/// empty-read refusal instead of writing an empty graph — asserted in the second half
/// of this test, against the SAME warehouse, so the two answers are told apart on the
/// live path and not only in a unit test.
#[cfg(feature = "skade")]
#[tokio::test]
async fn a_skade_warehouse_feeds_bifrost_through_the_arrow_source() {
const N: i64 = 91;
let wh_dir = tempfile::tempdir().unwrap();
// ── seed a real Iceberg table through skade ──────────────────────────────
{
let wh = skade::Warehouse::open(wh_dir.path()).await.expect("open warehouse");
let batch = people(1..=N);
let mut table = wh
.table_or_create("people", &batch.schema())
.await
.expect("create table");
table.append(&[batch]).await.expect("append");
}
// ── read it back as Arrow — this is the half every Warehouse impl has ────
let batches = {
let wh = skade::Warehouse::open(wh_dir.path()).await.expect("reopen");
let table = wh.table("people").await.expect("table handle");
table.read().await.expect("read the table as Arrow")
};
let rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
assert_eq!(rows, N as usize, "skade handed back every seeded row");
// ── feed bifröst ─────────────────────────────────────────────────────────
let sink = RecordingSink::default();
let dir = tempfile::tempdir().unwrap();
let source = ArrowSource::new(
"skade",
ArrowTables::new().with_table("lake.hr.people", batches),
);
let mut listener = listener_over(source, sink.clone(), dir.path());
let round = listener.sync_once().await.expect("a round over a skade warehouse");
assert_eq!(round.failed, vec![], "no sync failures");
assert_eq!(round.synced.len(), 1);
assert_eq!(round.synced[0].1, N as usize, "every skade row reached the mapping");
assert_eq!(
*sink.vertices.lock().unwrap(),
N as usize,
"a graph was built from a NON-Spark warehouse: {N} Person vertices"
);
assert_eq!(sink.labels.lock().unwrap().as_slice(), ["Person"]);
// ── the red half, on the same live path: an EMPTY warehouse read ─────────
let empty_dir = tempfile::tempdir().unwrap();
let empty_batches = {
let wh = skade::Warehouse::open(empty_dir.path()).await.expect("open");
let schema = people(1..=0).schema();
let table = wh.table_or_create("people", &schema).await.expect("create");
table.read().await.expect("read an empty table")
};
assert_eq!(
empty_batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
0,
"the second warehouse really is empty"
);
let empty_sink = RecordingSink::default();
let empty_state = tempfile::tempdir().unwrap();
let mut empty_listener = listener_over(
ArrowSource::new(
"skade",
ArrowTables::new().with_table("lake.hr.people", empty_batches),
),
empty_sink.clone(),
empty_state.path(),
);
let round = empty_listener.sync_once().await.expect("the round itself returns");
assert_eq!(
round.synced,
vec![],
"an empty warehouse must not report a synced table"
);
assert_eq!(round.failed.len(), 1, "it reports a FAILURE, not a silent success");
assert!(
round.failed[0].1.contains("0 rows"),
"and the failure says what happened: {}",
round.failed[0].1
);
assert_eq!(
*empty_sink.vertices.lock().unwrap(),
0,
"nothing was written, and nothing claimed it was"
);
}