crates/knut-bifrost/tests/spark_rowlevel_changelog.rs
//! **Spark's UPDATE / DELETE / MERGE, read back by skade as rows** — against a
//! real Spark 4.1.2 + Iceberg 1.11.0, in both write modes.
//!
//! Spark writes the table (copy-on-write, merge-on-read on format v2 — position
//! delete files — and merge-on-read on format **v3**, where Iceberg 1.11 writes
//! **deletion vectors**: roaring bitmaps in Puffin files); skade reads each
//! commit with `read_changelog`; the answer is held to **Spark's own time travel**:
//! `SELECT … VERSION AS OF <before>` vs `… VERSION AS OF <after>` — the rows that
//! left are the DELETEs, the rows that arrived the INSERTs. That oracle is
//! independent of skade and of Iceberg's changelog view (which is not asked).
//!
//! The Spark container comes up through draupnir (host network, the warehouse
//! bind-mounted at the SAME path inside and out, so the metadata's file paths
//! resolve on the host). `#[ignore]`d: needs a podman socket and the image.
//!
//! ```text
//! KNUT_SPARK_IMAGE=localhost/nordisk-spark-iceberg:4.1.2 \
//! cargo test -p knut-bifrost --features spark,skade --test spark_rowlevel_changelog -- --ignored --nocapture
//! ```
#![cfg(all(feature = "spark", feature = "skade"))]
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use arrow_array::{Array, Int64Array, RecordBatch, StringArray};
use draupnir::container::{ContainerBoot, ContainerControl};
use draupnir::{BootSpec, Machine, NetMode};
use knut_bifrost::source::SparkConnectSource;
use skade::iceberg::io::{FileIOBuilder, LocalFsStorageFactory};
use skade::iceberg::table::StaticTable;
use skade::iceberg::TableIdent;
const PORT: u16 = 15055;
const NAME: &str = "knut-spark-rowlevel";
fn image() -> String {
std::env::var("KNUT_SPARK_IMAGE").unwrap_or_else(|_| "localhost/nordisk-spark-iceberg:4.1.2".into())
}
fn warehouse() -> PathBuf {
std::env::var_os("SPARK_WH").map(PathBuf::from).unwrap_or_else(|| {
PathBuf::from(std::env::var("HOME").unwrap()).join("t14s-runs/knut-doris/spark-wh")
})
}
/// Boot Spark Connect with the `lake` hadoop catalog on `wh` (same path in and out).
async fn boot_spark(wh: &Path) -> Machine {
std::fs::create_dir_all(wh).unwrap();
// The image runs as `spark` (uid 185 → a subuid on the host): the bind root
// must be writable by it.
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(wh, std::fs::Permissions::from_mode(0o777)).unwrap();
let whs = wh.display().to_string();
let mut s = BootSpec::container(NAME, image()).with_cmd(
[
"/opt/spark/sbin/start-connect-server.sh",
"--conf",
"spark.sql.extensions=org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions",
"--conf",
"spark.sql.catalog.lake=org.apache.iceberg.spark.SparkCatalog",
"--conf",
"spark.sql.catalog.lake.type=hadoop",
"--conf",
&format!("spark.sql.catalog.lake.warehouse={whs}"),
"--conf",
&format!("spark.connect.grpc.binding.port={PORT}"),
]
.map(String::from),
);
s.env.insert("SPARK_NO_DAEMONIZE".into(), "1".into());
s.net = NetMode::Host;
let binds = vec![format!("{whs}:{whs}")];
tokio::task::spawn_blocking(move || ContainerBoot::new().boot_with_binds(&s, &binds))
.await
.unwrap()
.expect("boot spark connect")
}
async fn connect() -> SparkConnectSource {
let url = format!("sc://127.0.0.1:{PORT}");
for _ in 0..80 {
if let Ok(s) = SparkConnectSource::connect(&url).await {
if s.sql("SELECT 1").await.is_ok() {
return s;
}
}
tokio::time::sleep(Duration::from_secs(3)).await;
}
panic!("Spark Connect at {url} never answered");
}
/// `(id, name)` multiset of a query's rows.
fn bag(batches: &[RecordBatch]) -> BTreeMap<(i64, Option<String>), i64> {
let mut m = BTreeMap::new();
for b in batches {
let id = b.column_by_name("id").unwrap().as_any().downcast_ref::<Int64Array>().unwrap();
let name = b.column_by_name("name").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
for i in 0..b.num_rows() {
*m.entry((id.value(i), (!name.is_null(i)).then(|| name.value(i).to_string()))).or_insert(0) += 1;
}
}
m
}
/// `a - b` as a multiset.
fn minus(a: &BTreeMap<(i64, Option<String>), i64>, b: &BTreeMap<(i64, Option<String>), i64>) -> BTreeMap<(i64, Option<String>), i64> {
let mut out = BTreeMap::new();
for (k, n) in a {
let left = n - b.get(k).copied().unwrap_or(0);
if left > 0 {
out.insert(k.clone(), left);
}
}
out
}
/// skade's changelog of `(from, to]` split into DELETE and INSERT multisets.
fn split(cl: &skade::ChangelogBatch) -> (BTreeMap<(i64, Option<String>), i64>, BTreeMap<(i64, Option<String>), i64>) {
let (mut del, mut ins) = (BTreeMap::new(), BTreeMap::new());
for b in &cl.rows {
let t = b.column_by_name(skade::change_type::COLUMN).unwrap().as_any().downcast_ref::<StringArray>().unwrap();
let id = b.column_by_name("id").unwrap().as_any().downcast_ref::<Int64Array>().unwrap();
let name = b.column_by_name("name").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
for i in 0..b.num_rows() {
let k = (id.value(i), (!name.is_null(i)).then(|| name.value(i).to_string()));
let side = if t.value(i) == "DELETE" { &mut del } else { &mut ins };
*side.entry(k).or_insert(0) += 1;
}
}
(del, ins)
}
/// The hadoop-catalog table at `<wh>/cdc/<name>`, opened off its newest metadata.
async fn open_local(wh: &Path, name: &str) -> skade::iceberg::table::Table {
let meta = wh.join("cdc").join(name).join("metadata");
let hint = std::fs::read_to_string(meta.join("version-hint.text")).unwrap();
let file = meta.join(format!("v{}.metadata.json", hint.trim()));
let io = FileIOBuilder::new(Arc::new(LocalFsStorageFactory)).build();
StaticTable::from_metadata_file(&file.display().to_string(), TableIdent::from_strs(["cdc", name]).unwrap(), io)
.await
.expect("open the Spark-written table off disk")
.into_table()
}
async fn one_mode(spark: &SparkConnectSource, wh: &Path, mode: &str, format: u8, tag: &str) {
let name = format!("t_{}_v{format}_{tag}", mode.replace('-', "_"));
let t = format!("lake.cdc.{name}");
for q in [
"CREATE NAMESPACE IF NOT EXISTS lake.cdc".to_string(),
format!(
"CREATE TABLE {t} (id BIGINT, name STRING) USING iceberg TBLPROPERTIES \
('format-version'='{format}','write.delete.mode'='{mode}','write.update.mode'='{mode}','write.merge.mode'='{mode}')"
),
format!("INSERT INTO {t} SELECT id, concat('v', id) FROM range(0, 50)"),
format!("UPDATE {t} SET name = concat('u', id) WHERE id % 10 = 3"),
format!("DELETE FROM {t} WHERE id % 10 = 7"),
format!(
"MERGE INTO {t} x USING (SELECT * FROM VALUES (1L, 'm1'), (100L, 'm100') AS s(id, name)) s ON x.id = s.id \
WHEN MATCHED THEN UPDATE SET name = s.name WHEN NOT MATCHED THEN INSERT *"
),
] {
spark.sql(&q).await.unwrap_or_else(|e| panic!("spark: {q}: {e}"));
}
let snaps: Vec<(i64, String)> = spark
.sql(&format!("SELECT snapshot_id, operation FROM {t}.snapshots ORDER BY committed_at"))
.await
.unwrap()
.iter()
.flat_map(|b| {
let id = b.column(0).as_any().downcast_ref::<Int64Array>().unwrap().clone();
let op = b.column(1).as_any().downcast_ref::<StringArray>().unwrap().clone();
(0..b.num_rows()).map(move |i| (id.value(i), op.value(i).to_string())).collect::<Vec<_>>()
})
.collect();
eprintln!("{mode}: {} snapshots: {snaps:?}", snaps.len());
assert_eq!(snaps.len(), 4, "insert + update + delete + merge");
// v3 merge-on-read must really have written deletion vectors, or the case
// below tests nothing new.
if format == 3 && mode == "merge-on-read" {
let formats: Vec<String> = spark
.sql(&format!("SELECT DISTINCT file_format FROM {t}.all_delete_files"))
.await
.unwrap()
.iter()
.flat_map(|b| {
let c = b.column(0).as_any().downcast_ref::<StringArray>().unwrap().clone();
(0..b.num_rows()).map(move |i| c.value(i).to_string()).collect::<Vec<_>>()
})
.collect();
eprintln!("v3 {mode}: delete file formats {formats:?}");
assert!(formats.iter().any(|f| f.eq_ignore_ascii_case("puffin")), "Spark wrote no deletion vector: {formats:?}");
}
let table = open_local(wh, &name).await;
// The plain scan applies the deletes too: the table as skade reads it is the
// table as Spark reads it.
let now = bag(&skade::read_all(&table).await.unwrap());
let spark_now = bag(&spark.sql(&format!("SELECT id, name FROM {t}")).await.unwrap());
assert_eq!(now, spark_now, "{mode} v{format}: skade's scan = Spark's SELECT");
for w in snaps.windows(2) {
let ((from, _), (to, op)) = (&w[0], &w[1]);
let before = bag(&spark.sql(&format!("SELECT id, name FROM {t} VERSION AS OF {from}")).await.unwrap());
let after = bag(&spark.sql(&format!("SELECT id, name FROM {t} VERSION AS OF {to}")).await.unwrap());
let (want_del, want_ins) = (minus(&before, &after), minus(&after, &before));
let plan = skade::plan_changelog(&table, Some(*from), *to).await.unwrap();
assert!(!plan.needs_full_reload, "{mode} v{format} {op}: Spark's commit is readable as rows");
let cl = skade::read_changelog(&table, Some(*from), *to).await.unwrap();
let (del, ins) = split(&cl);
eprintln!("{mode} v{format} {op} {from}->{to}: spark -{} +{} · skade -{} +{}", want_del.len(), want_ins.len(), del.len(), ins.len());
assert_eq!(del, want_del, "{mode} v{format} {op}: DELETE rows = rows that left (Spark time travel)");
assert_eq!(ins, want_ins, "{mode} v{format} {op}: INSERT rows = rows that arrived (Spark time travel)");
}
}
#[tokio::test]
#[ignore = "live: needs podman + the Spark 4.1.2 + Iceberg 1.11 image"]
async fn spark_update_delete_merge_read_back_as_rows_in_both_modes() {
let wh = warehouse();
let machine = boot_spark(&wh).await;
let spark = connect().await;
let tag = format!("{}", std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs());
one_mode(&spark, &wh, "copy-on-write", 2, &tag).await;
one_mode(&spark, &wh, "merge-on-read", 2, &tag).await;
// Deletion vectors: merge-on-read on format v3.
one_mode(&spark, &wh, "merge-on-read", 3, &tag).await;
one_mode(&spark, &wh, "copy-on-write", 3, &tag).await;
tokio::task::spawn_blocking(move || ContainerBoot::new().stop(&machine)).await.unwrap();
}