knut-bifrost 0.1.3

//! 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::{
    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>;
}

/// 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>,
}

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(),
                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 {
            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
                )));
            }
            match kinds.value(row) {
                "INSERT" | "UPDATE_AFTER" => keep_upsert.push(row as u32),
                "DELETE" => keep_delete.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.)
                "UPDATE_BEFORE" => {}
                other => {
                    return Err(BifrostError::Source(format!(
                        "unknown {} '{other}'",
                        SyncMode::CHANGE_TYPE_COLUMN
                    )));
                }
            }
        }

        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: deleting a vertex DETACH-deletes its edges, and an edge-only
/// changelog delete is out of scope.
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(ctx, delete_batches)?;
        for set in mutation.vertices {
            deletes.push(VertexDelete {
                label: set.label,
                id_column: set.id_column,
                batch: set.batch,
            });
        }
    }
    Ok(deletes)
}

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