knut_bifrost/

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)?)
}