crates/knut-bifrost/src/transform.rs
//! The transform fabric: built-in declarative transforms plus the trait
//! custom Rust transforms implement.
use std::{collections::HashMap, sync::Arc};
use arrow_array::{Array, RecordBatch, StringArray};
use arrow_schema::{Field, Schema};
use crate::{
cdc::ChangeOp,
error::{BifrostError, Result},
mapping::SyncMode,
mutation::{EdgeSet, GraphMutation, VertexDelete, VertexSet},
};
/// Information about the table a batch came from, available to transforms.
#[derive(Debug, Clone)]
pub struct TableContext {
/// Fully qualified table name, e.g. `lake.social.users`.
pub table: String,
}
/// One step of a pipeline: table rows in, graph elements out.
///
/// Implement this for cases the declarative mappings can't express, and
/// register it under a name with [`crate::registry::TransformRegistry`].
pub trait Transform: Send + Sync {
fn apply(&self, ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation>;
/// The same step over a changelog round's `DELETE` rows, whose nodes and
/// edges are then removed. Such a row may carry only its key — every other
/// column null — so a step that filters or computes on other columns must
/// not let that drop the delete. The default is [`Self::apply`], right for
/// every step that only picks columns.
fn apply_deletes(&self, ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
self.apply(ctx, batches)
}
}
/// Endpoint reference used by edge transforms.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct EndpointSpec {
pub label: String,
/// Column in *this* table holding the endpoint's id value.
pub column: String,
/// Id column on the endpoint node itself; defaults to `column`.
#[serde(default)]
pub id_column: Option<String>,
}
impl EndpointSpec {
fn node_id_column(&self) -> &str {
self.id_column.as_deref().unwrap_or(&self.column)
}
}
// ── vertex ──────────────────────────────────────────────────────────────────
/// Rows → nodes. Selects `id_column` + `properties`, applies `rename`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct VertexTransform {
pub label: String,
pub id_column: String,
#[serde(default)]
pub properties: Vec<String>,
#[serde(default)]
pub rename: HashMap<String, String>,
}
impl Transform for VertexTransform {
fn apply(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
let mut out = GraphMutation::default();
for batch in batches {
let mut keep = vec![self.id_column.clone()];
keep.extend(self.properties.iter().cloned());
let projected = project(batch, &keep, &self.rename)?;
out.vertices.push(VertexSet {
label: self.label.clone(),
id_column: renamed(&self.id_column, &self.rename),
batch: projected,
});
}
Ok(out)
}
}
// ── edge ────────────────────────────────────────────────────────────────────
/// Rows → relationships. `src.column`/`dst.column` become the `src`/`dst`
/// columns of the output batch; the `properties` columns ride along and the
/// sink folds them into the relationship's `props` (`SET e += r.props`).
/// `rename` maps a source column to the relationship-property name it lands
/// under, mirroring [`VertexTransform`]; the endpoint columns always become
/// `src`/`dst` regardless of `rename`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct EdgeTransform {
#[serde(rename = "type")]
pub edge_type: String,
pub src: EndpointSpec,
pub dst: EndpointSpec,
#[serde(default)]
pub properties: Vec<String>,
/// Optional source-column → property-name remapping for the property
/// columns (the endpoint columns are renamed to `src`/`dst` for free).
#[serde(default)]
pub rename: HashMap<String, String>,
/// Source columns that tell **parallel** edges apart: they join the
/// relationship's MERGE key (`-[e:type {k: …}]->`), so two rows with the same
/// endpoints and different keys are two edges. Empty (the default) keeps one
/// edge per `(src, type, dst)`. A key column is written as a property too.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub key: Vec<String>,
/// One edge of this type per source node: an upsert first deletes the
/// source's other outgoing edges of this type (see
/// [`EdgeSet::replace_outgoing`]). Unset means: yes for an `fk_edges` step
/// on an entity row — the table's vertex step IS the edge's source (label and
/// id column) — no for a junction table and for an `edge` step (their rows are
/// independent edges). Never with `key`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub replace: Option<bool>,
}
impl Transform for EdgeTransform {
fn apply(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
let mut out = GraphMutation::default();
for batch in batches {
let projected = self.project_edge(batch)?;
out.edges.push(EdgeSet {
src_label: self.src.label.clone(),
edge_type: self.edge_type.clone(),
dst_label: self.dst.label.clone(),
src_id_column: self.src.node_id_column().to_string(),
dst_id_column: self.dst.node_id_column().to_string(),
// One relationship per endpoint pair unless `key` names the
// columns that tell parallel edges apart.
edge_key_props: self.key.iter().map(|k| renamed(k, &self.rename)).collect(),
replace_outgoing: self.replace.unwrap_or(false) && self.key.is_empty(),
batch: projected,
});
}
Ok(out)
}
}
impl EdgeTransform {
/// Build the edge output batch: the `src.column` becomes `src`, the
/// `dst.column` becomes `dst`, and the `properties` ride along (renamed).
///
/// Done explicitly rather than via [`project`] because a **self-key** edge
/// (`src.column == dst.column`, e.g. a "same-as"/sibling link keyed on one
/// column) must emit *both* `src` and `dst` from that single input column —
/// a name-keyed rename map can only hold one mapping for the column, so the
/// generic projection would drop one endpoint. Property columns that clash
/// with `src`/`dst` or with each other are reported as
/// [`BifrostError::DuplicateOutputColumn`].
fn project_edge(&self, batch: &RecordBatch) -> Result<RecordBatch> {
let schema = batch.schema();
let mut fields = Vec::new();
let mut arrays = Vec::new();
let mut seen_output = std::collections::HashSet::new();
let mut push = |out_name: &str, src_col: &str| -> Result<()> {
if !seen_output.insert(out_name.to_string()) {
return Err(BifrostError::DuplicateOutputColumn(out_name.to_string()));
}
let idx = schema
.index_of(src_col)
.map_err(|_| BifrostError::ColumnNotFound(src_col.to_string()))?;
let field = schema.field(idx);
fields.push(Field::new(
out_name,
field.data_type().clone(),
field.is_nullable(),
));
arrays.push(batch.column(idx).clone());
Ok(())
};
// Endpoints first — both emitted even when they share one input column.
push("src", &self.src.column)?;
push("dst", &self.dst.column)?;
// Properties (de-duplicated by source column), renamed; an output-name
// clash with an endpoint or another property is a hard error.
let mut seen_prop_source = std::collections::HashSet::new();
for prop in self.properties.iter().chain(&self.key) {
if !seen_prop_source.insert(prop.clone()) {
continue;
}
let out_name = renamed(prop, &self.rename);
push(&out_name, prop)?;
}
Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?)
}
}
/// `fk_edges` is an `edge` whose endpoints live on other tables/steps — the
/// mechanics are identical, the alias exists so mappings read naturally.
pub type FkEdgesTransform = EdgeTransform;
// ── changelog (CDC) routing ───────────────────────────────────────────────────
/// A changelog batch split by its [`SyncMode::CHANGE_TYPE_COLUMN`] into the
/// rows that upsert vs. the rows that delete, both with the `_change_type`
/// column stripped so they render like ordinary table rows.
///
/// `INSERT` / `UPDATE_AFTER` → `upserts`; `DELETE` → `deletes`. `UPDATE_BEFORE`
/// is dropped: it is the pre-image of an in-place update whose `UPDATE_AFTER`
/// upsert overwrites the node, so emitting a delete for it would destroy the
/// updated row (deletes are applied last in the sink). An unknown change type is
/// a hard error (better than silently dropping a row).
#[derive(Debug, Default)]
pub struct ChangeRows {
/// Rows to feed the upsert pipeline (vertices + edges).
pub upserts: Vec<RecordBatch>,
/// Rows that name nodes to remove.
pub deletes: Vec<RecordBatch>,
}
impl ChangeRows {
pub fn is_empty(&self) -> bool {
self.upserts.is_empty() && self.deletes.is_empty()
}
}
/// Does this batch carry a changelog `_change_type` column?
pub fn is_changelog_batch(batch: &RecordBatch) -> bool {
batch
.schema()
.index_of(SyncMode::CHANGE_TYPE_COLUMN)
.is_ok()
}
/// Partition changelog `batches` into upsert vs. delete rows by `_change_type`,
/// stripping that column from the output. Batches without a `_change_type`
/// column are treated as all-upsert (and passed through unchanged), so a
/// non-CDC source flowing through CDC routing is a no-op.
pub fn partition_change_rows(batches: &[RecordBatch]) -> Result<ChangeRows> {
let mut out = ChangeRows::default();
for batch in batches {
let Ok(ct_idx) = batch.schema().index_of(SyncMode::CHANGE_TYPE_COLUMN) else {
out.upserts.push(batch.clone());
continue;
};
let kinds = batch
.column(ct_idx)
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| {
BifrostError::Source(format!(
"{} must be a string column",
SyncMode::CHANGE_TYPE_COLUMN
))
})?;
let mut keep_upsert = Vec::with_capacity(batch.num_rows());
let mut keep_delete = Vec::with_capacity(batch.num_rows());
for row in 0..batch.num_rows() {
// A null change type is ambiguous; refuse it loudly.
if kinds.is_null(row) {
return Err(BifrostError::Source(format!(
"null {} at row {row}",
SyncMode::CHANGE_TYPE_COLUMN
)));
}
// Parsed through the ONE shared vocabulary (`cdc::ChangeOp`), so
// this routes Iceberg's `INSERT`/`UPDATE_AFTER`/… and Spark 4.2's
// `insert`/`update_postimage`/… identically. Before that, a Spark
// 4.2 native changelog would have hit the `other =>` arm and failed
// the whole round — the two engines genuinely disagree on spelling.
match ChangeOp::parse(kinds.value(row)) {
Some(ChangeOp::Insert) | Some(ChangeOp::UpdateAfter) => {
keep_upsert.push(row as u32)
}
Some(ChangeOp::Delete) => keep_delete.push(row as u32),
// A folded `update` row (see `cdc::fold_updates`) carries the
// post-image, so it upserts like an `UPDATE_AFTER`.
Some(ChangeOp::Update) => keep_upsert.push(row as u32),
// UPDATE_BEFORE is the pre-image of an in-place update. Its
// paired UPDATE_AFTER (an upsert) fully overwrites the node via
// MERGE+SET, so the pre-image carries no net state change. It
// must NOT become a delete: deletes run last in the sink, so
// routing UPDATE_BEFORE to `deletes` would DETACH DELETE the
// node the UPDATE_AFTER just upserted — destroying an updated
// row. Drop it. (A true row removal arrives as DELETE.)
Some(ChangeOp::UpdateBefore) => {}
None => {
return Err(BifrostError::Source(format!(
"unknown {} '{}'",
SyncMode::CHANGE_TYPE_COLUMN,
kinds.value(row)
)));
}
}
}
let stripped = drop_column(batch, ct_idx)?;
if !keep_upsert.is_empty() {
out.upserts.push(take_rows(&stripped, &keep_upsert)?);
}
if !keep_delete.is_empty() {
out.deletes.push(take_rows(&stripped, &keep_delete)?);
}
}
Ok(out)
}
/// Run `transforms` over the delete rows and reinterpret every emitted
/// [`VertexSet`] as a [`VertexDelete`] (same label + id column). Edge sets are
/// dropped here: deleting a vertex DETACH-deletes its edges. An edge-only
/// table's deletes are [`edge_deletes_from_change_rows`].
pub fn deletes_from_change_rows(
transforms: &[Arc<dyn Transform>],
ctx: &TableContext,
delete_batches: &[RecordBatch],
) -> Result<Vec<VertexDelete>> {
if delete_batches.is_empty() {
return Ok(Vec::new());
}
let mut deletes = Vec::new();
for transform in transforms {
let mutation = transform.apply_deletes(ctx, delete_batches)?;
for set in mutation.vertices {
deletes.push(VertexDelete {
label: set.label,
id_column: set.id_column,
batch: set.batch,
});
}
}
Ok(deletes)
}
/// Run `transforms` over the delete rows and keep every emitted [`EdgeSet`]:
/// the relationships those rows made, to remove.
///
/// For every edge step, not only an edge-only table's: when a row changes its
/// foreign key (an UPDATE is a DELETE of the old row plus an INSERT of the new
/// one), the old row's edge points at the old endpoint and must go, while its
/// vertex delete is suppressed by the upsert. The sink suppresses an edge
/// delete the same way — an edge also upserted this round survives.
pub fn edge_deletes_from_change_rows(
transforms: &[Arc<dyn Transform>],
ctx: &TableContext,
delete_batches: &[RecordBatch],
) -> Result<Vec<EdgeSet>> {
if delete_batches.is_empty() {
return Ok(Vec::new());
}
let mut out = Vec::new();
for transform in transforms {
out.extend(transform.apply_deletes(ctx, delete_batches)?.edges);
}
Ok(out)
}
// ── helpers ─────────────────────────────────────────────────────────────────
fn renamed(name: &str, rename: &HashMap<String, String>) -> String {
rename
.get(name)
.cloned()
.unwrap_or_else(|| name.to_string())
}
/// Return a copy of `batch` with column `idx` removed.
fn drop_column(batch: &RecordBatch, idx: usize) -> Result<RecordBatch> {
let schema = batch.schema();
let fields: Vec<Field> = schema
.fields()
.iter()
.enumerate()
.filter(|(i, _)| *i != idx)
.map(|(_, f)| f.as_ref().clone())
.collect();
let columns = batch
.columns()
.iter()
.enumerate()
.filter(|(i, _)| *i != idx)
.map(|(_, c)| c.clone())
.collect();
Ok(RecordBatch::try_new(
Arc::new(Schema::new(fields)),
columns,
)?)
}
/// Select `rows` (by index) from every column of `batch`.
fn take_rows(batch: &RecordBatch, rows: &[u32]) -> Result<RecordBatch> {
let indices = arrow_array::UInt32Array::from(rows.to_vec());
let columns = batch
.columns()
.iter()
.map(|c| arrow::compute::take(c.as_ref(), &indices, None))
.collect::<std::result::Result<Vec<_>, _>>()?;
Ok(RecordBatch::try_new(batch.schema(), columns)?)
}
/// Select `columns` (by name, in order) from `batch`, applying `rename`.
///
/// `columns` is de-duplicated by **source** name first (a column listed twice —
/// e.g. an `id_column` that also appears in a vertex's `properties` — is taken
/// once), then any remaining clash on the *output* name (two different source
/// columns renamed to the same name) is reported as
/// [`BifrostError::DuplicateOutputColumn`] rather than surfacing as an opaque
/// Arrow duplicate-field error from `RecordBatch::try_new`.
fn project(
batch: &RecordBatch,
columns: &[String],
rename: &HashMap<String, String>,
) -> Result<RecordBatch> {
let schema = batch.schema();
let mut fields = Vec::with_capacity(columns.len());
let mut arrays = Vec::with_capacity(columns.len());
let mut seen_source = std::collections::HashSet::new();
let mut seen_output = std::collections::HashSet::new();
for name in columns {
// Same source column requested twice → keep the first, skip the rest.
if !seen_source.insert(name.clone()) {
continue;
}
let out_name = renamed(name, rename);
if !seen_output.insert(out_name.clone()) {
return Err(BifrostError::DuplicateOutputColumn(out_name));
}
let idx = schema
.index_of(name)
.map_err(|_| BifrostError::ColumnNotFound(name.clone()))?;
let field = schema.field(idx);
fields.push(Field::new(
out_name,
field.data_type().clone(),
field.is_nullable(),
));
arrays.push(batch.column(idx).clone());
}
Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?)
}