knut-bifrost 0.1.4

//! **The canonical offline rel2graph engine seam.**
//!
//! bifröst *is* the rel2graph engine: a [`MappingSpec`](crate::MappingSpec) is
//! compiled to a [`Transform`](crate::Transform) fabric and applied to a table's
//! Arrow batches to produce a [`GraphMutation`](crate::GraphMutation) (the
//! `VertexSet`s + `EdgeSet`s that ARE the graph). The live path runs this per
//! round inside the [`Listener`](crate::Listener), reading from a
//! [`Source`](crate::source) and writing to a [`GraphSink`](crate::sink::GraphSink).
//!
//! This module exposes the SAME work with **no source and no sink** — you hand it
//! the per-table Arrow batches, it hands you the `GraphMutation`. It is the one
//! entrypoint every DB-free consumer of the engine shares:
//!
//! * the [correctness + timing bench](../../knut_rel2graph_bench/index.html) drives
//!   it over a seeded synthetic database;
//! * [`knut-flight`](../../knut_flight/index.html) drives it over an Iceberg /
//!   parquet export to serve a REAL graph as Arrow Flight;
//! * the networkx file-export drives it to dump a node-list + edge-list + oracle;
//! * **korp's `iceberg2graph` demo step** calls it to turn the tables it just read
//!   from the lakehouse into the graph it hands the viewer — the same graph the
//!   PySpark projection ([`GraphPlan::to_pyspark`](crate::GraphPlan::to_pyspark))
//!   and the thund projection (`knut_thund::rel2graph::to_thund_pipeline`) would
//!   build, since all three derive from ONE reviewed [`MappingSpec`].
//!
//! The point of the seam is that swapping the data behind it never touches the
//! engine: whatever yields per-table [`RecordBatch`]es — the bench generator, a
//! parquet directory, a Spark-Connect Iceberg read — implements [`TableSource`]
//! and feeds the identical [`build_mutation`].

use std::collections::BTreeMap;

use arrow_array::RecordBatch;

use crate::{
    error::Result, mapping::MappingSpec, mutation::GraphMutation, registry::TransformRegistry,
    transform::TableContext,
};

/// A source of Arrow batches keyed by fully-qualified table name — the seam the
/// offline rel2graph engine reads from.
///
/// Implemented by anything that can produce a table's rows as Arrow: the bench's
/// `GeneratedDb`, the in-memory [`ArrowTables`] (a parquet-directory load or an
/// Iceberg scan collected into memory), or a bespoke adapter. [`build_mutation`]
/// asks only for `batches(table)`; an unknown table returns an empty slice and is
/// skipped, so a source may legitimately carry a superset of the spec's tables.
pub trait TableSource {
    /// The batches for `table`, or an empty slice if the source has no such table.
    fn batches(&self, table: &str) -> &[RecordBatch];
}

/// **Run the bifröst rel2graph engine offline over `source`.**
///
/// Compiles every table's declarative pipeline once, applies the resulting
/// [`Transform`](crate::Transform)s to that table's batches, and merges every
/// step's output into one [`GraphMutation`]. This mirrors exactly what the live
/// [`Listener`](crate::Listener) does per round, minus the source read and the
/// sink write — so the mutation this returns is the mutation a live load would
/// MERGE, at the same seam the correctness oracle asserts on.
///
/// Tables the `source` does not carry are skipped (empty batches), so a partial
/// source yields a partial-but-consistent graph.
///
/// ```
/// use knut_bifrost::rel2graph::{build_mutation, ArrowTables};
/// use knut_bifrost::MappingSpec;
/// # use arrow_array::{RecordBatch, Int64Array};
/// # use arrow_schema::{Schema, Field, DataType};
/// # use std::sync::Arc;
/// let spec = MappingSpec::from_yaml(r#"
/// version: 1
/// tables:
///   - table: lake.hr.employees
///     pipeline:
///       - vertex: { label: Employee, id_column: employee_id, properties: [] }
/// "#).unwrap();
/// # let schema = Arc::new(Schema::new(vec![Field::new("employee_id", DataType::Int64, false)]));
/// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1i64, 2, 3]))]).unwrap();
/// let mut src = ArrowTables::new();
/// src.insert("lake.hr.employees", vec![batch]);
/// let mutation = build_mutation(&spec, &src).unwrap();
/// assert_eq!(mutation.vertices.len(), 1);
/// ```
pub fn build_mutation(spec: &MappingSpec, source: &dyn TableSource) -> Result<GraphMutation> {
    let registry = TransformRegistry::new();
    let mut mutation = GraphMutation::default();

    for tm in &spec.tables {
        let batches = source.batches(&tm.table);
        if batches.is_empty() {
            continue;
        }
        let ctx = TableContext {
            table: tm.table.clone(),
        };
        let transforms = tm.compile(&registry)?;
        for transform in &transforms {
            let m = transform.apply(&ctx, batches)?;
            mutation.merge(m);
        }
    }

    Ok(mutation)
}

/// An in-memory [`TableSource`]: a map of fully-qualified table name → Arrow
/// batches.
///
/// This is the adapter that turns "some Arrow I already have" into engine input.
/// Fill it from a parquet directory, an Iceberg scan, a Spark-Connect read, or by
/// hand, then drive [`build_mutation`] over real data with no bench and no live
/// FalkorDB.
#[derive(Debug, Clone, Default)]
pub struct ArrowTables {
    tables: BTreeMap<String, Vec<RecordBatch>>,
}

impl ArrowTables {
    /// An empty source.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add (or replace) a table's batches. Returns nothing — see
    /// [`with_table`](Self::with_table) for the builder form.
    pub fn insert(&mut self, table: impl Into<String>, batches: Vec<RecordBatch>) {
        self.tables.insert(table.into(), batches);
    }

    /// Builder form of [`insert`](Self::insert): add a table, return `self`.
    pub fn with_table(mut self, table: impl Into<String>, batches: Vec<RecordBatch>) -> Self {
        self.insert(table, batches);
        self
    }

    /// The table names this source carries, in sorted order.
    pub fn table_names(&self) -> impl Iterator<Item = &str> {
        self.tables.keys().map(String::as_str)
    }

    /// Whether the source carries no tables.
    pub fn is_empty(&self) -> bool {
        self.tables.is_empty()
    }

    /// The number of tables carried.
    pub fn len(&self) -> usize {
        self.tables.len()
    }
}

impl TableSource for ArrowTables {
    fn batches(&self, table: &str) -> &[RecordBatch] {
        self.tables.get(table).map(Vec::as_slice).unwrap_or(&[])
    }
}

/// A shared reference to a [`TableSource`] is itself a [`TableSource`], so
/// `&db`/`&tables` can be passed where an owned source is not convenient.
impl<T: TableSource + ?Sized> TableSource for &T {
    fn batches(&self, table: &str) -> &[RecordBatch] {
        (**self).batches(table)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::Int64Array;
    use arrow_schema::{DataType, Field, Schema};
    use std::sync::Arc;

    fn id_batch(col: &str, ids: &[i64]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![Field::new(col, DataType::Int64, false)]));
        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ids.to_vec()))]).unwrap()
    }

    #[test]
    fn arrow_tables_source_drives_the_engine() {
        let spec = MappingSpec::from_yaml(
            r#"
version: 1
tables:
  - table: lake.hr.employees
    pipeline:
      - vertex: { label: Employee, id_column: employee_id, properties: [] }
"#,
        )
        .unwrap();

        let src = ArrowTables::new().with_table(
            "lake.hr.employees",
            vec![id_batch("employee_id", &[1, 2, 3])],
        );
        assert_eq!(src.len(), 1);

        let mutation = build_mutation(&spec, &src).unwrap();
        assert_eq!(
            mutation.vertices.len(),
            1,
            "one VertexSet for the one table"
        );
        assert_eq!(mutation.vertices[0].label, "Employee");
        assert_eq!(mutation.vertices[0].batch.num_rows(), 3);
    }

    #[test]
    fn missing_tables_are_skipped_not_errors() {
        let spec = MappingSpec::from_yaml(
            r#"
version: 1
tables:
  - table: present.table
    pipeline:
      - vertex: { label: A, id_column: id, properties: [] }
  - table: absent.table
    pipeline:
      - vertex: { label: B, id_column: id, properties: [] }
"#,
        )
        .unwrap();

        // Source carries only `present.table`; `absent.table` yields no batches.
        let src = ArrowTables::new().with_table("present.table", vec![id_batch("id", &[10, 20])]);
        let mutation = build_mutation(&spec, &src).unwrap();
        assert_eq!(
            mutation.vertices.len(),
            1,
            "absent table skipped, present table mapped"
        );
        assert_eq!(mutation.vertices[0].label, "A");
    }
}