knut_bifrost/

rel2graph.rs

//! **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> {
    build_mutation_reported(spec, source).map(|(m, _)| m)
}

/// **What a `build_mutation` run actually read.**
///
/// `build_mutation` skips a spec table the source does not carry, which is the right
/// default (a partial source yields a partial-but-consistent graph) and is also
/// exactly how an **empty graph comes back looking like a successful build**: point
/// the source at the wrong namespace and every table is "skipped", the mutation is
/// empty, and the call returns `Ok`.
///
/// This report is the discrimination. It is carried BESIDE the mutation rather than
/// folded into it, because the mutation is what a sink MERGEs and a report is not.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct BuildReport {
    /// Spec tables the source carried rows for, in spec order.
    pub tables_mapped: Vec<String>,
    /// Spec tables the source carried NO rows for — absent, or present with zero
    /// rows. These are the ones that silently vanish from an unreported build.
    pub tables_empty: Vec<String>,
    /// Total input rows read across every mapped table.
    pub rows_read: usize,
    /// Vertex rows in the resulting mutation.
    pub vertices: usize,
    /// Edge rows in the resulting mutation.
    pub edges: usize,
}

impl BuildReport {
    /// **The honest one-liner a caller uses instead of trusting `Ok`.**
    ///
    /// `Ok(self)` when at least one spec table contributed rows AND the mutation is
    /// non-empty; otherwise a typed [`crate::error::BifrostError::Source`] naming
    /// every table that came back empty. A build that mapped nothing is not a graph.
    pub fn require_non_empty(self) -> Result<Self> {
        if self.rows_read == 0 || (self.vertices == 0 && self.edges == 0) {
            return Err(crate::error::BifrostError::Source(format!(
                "the rel2graph build read {} row(s) from {} table(s) and produced {} vertex \
                 row(s) + {} edge row(s) — that is an EMPTY graph, not a built one. Tables \
                 the source carried nothing for: {:?}. Check the catalog, the namespace and \
                 the table names in the mapping spec.",
                self.rows_read,
                self.tables_mapped.len(),
                self.vertices,
                self.edges,
                self.tables_empty,
            )));
        }
        Ok(self)
    }
}

/// [`build_mutation`] plus a [`BuildReport`] saying what it read. Same work, same
/// result — this is the one implementation and `build_mutation` delegates to it, so
/// the counts can never describe a different run than the mutation.
pub fn build_mutation_reported(
    spec: &MappingSpec,
    source: &dyn TableSource,
) -> Result<(GraphMutation, BuildReport)> {
    let registry = TransformRegistry::new();
    let mut mutation = GraphMutation::default();
    let mut report = BuildReport::default();

    for tm in &spec.tables {
        let batches = source.batches(&tm.table);
        let rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
        if batches.is_empty() || rows == 0 {
            report.tables_empty.push(tm.table.clone());
            continue;
        }
        report.tables_mapped.push(tm.table.clone());
        report.rows_read += rows;
        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);
        }
    }

    report.vertices = mutation.vertices.iter().map(|s| s.batch.num_rows()).sum();
    report.edges = mutation.edges.iter().map(|s| s.batch.num_rows()).sum();
    Ok((mutation, report))
}

/// 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");
    }
}