knut-bifrost 0.3.0

//! **SQL mapping** — a `SELECT` whose result columns say what they become.
//!
//! Feature `sql`. One statement is one step: it runs over each round's rows of
//! the table it names in `FROM`, and every result row is a node or an edge,
//! decided by the reserved columns it returns:
//!
//! | returns | becomes |
//! |---|---|
//! | `_label`, `_id` | a node of label `_label`, keyed on property `id` = `_id` |
//! | `_type`, `_from_label`, `_from`, `_to_label`, `_to` | an edge `(_from_label {id: _from})-[_type]->(_to_label {id: _to})` |
//! | `_edge_id` (edges, optional) | the edge's key: rows with the same endpoints and different `_edge_id` are parallel edges |
//!
//! Every other column is a property. Labels and types are values, so one
//! statement can write several (`CASE … END AS _label`). A row whose label,
//! type, id or endpoint is null writes nothing.
//!
//! ```sql
//! SELECT 'Person' AS _label, user_id AS _id, name, upper(country) AS country
//!   FROM social.users WHERE active;
//! SELECT 'WORKS_AT' AS _type, 'Person' AS _from_label, user_id AS _from,
//!        'Company' AS _to_label, company_id AS _to, since
//!   FROM social.users WHERE company_id IS NOT NULL;
//! ```
//!
//! # One table, row by row
//!
//! A statement reads ONE table (no joins, no `UNION`, no `WITH`): a round
//! carries one table's changed rows, and a join against another table would
//! have to read that table whole every round. Aggregates run, but see only the
//! round's rows.
//!
//! # Deletes
//!
//! In a changelog round a `DELETE` row may carry only its key, every other
//! column null — so the statement's `WHERE` would drop it and the node would
//! never be removed. Delete rows therefore run the statement **without its
//! `WHERE`**: a delete of a node the filter never wrote matches nothing and is
//! a no-op. What the delete needs is the `_label`/`_id` (or the edge's
//! endpoints) — so compute those from key columns, or from constants.
//!
//! The engine is DataFusion, on the rows as they are (a `MemTable`, one
//! partition, no copy).

use std::sync::Arc;

use arrow::compute::{cast, take};
use arrow_array::{ArrayRef, RecordBatch, UInt32Array};
use arrow_schema::{DataType, Field, Schema};
use datafusion::datasource::MemTable;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion::sql::sqlparser::ast::{Ident, ObjectName, ObjectNamePart, Query, SetExpr, Statement, TableFactor};
use datafusion::sql::sqlparser::dialect::GenericDialect;
use datafusion::sql::sqlparser::parser::Parser;

use crate::error::{BifrostError, Result};
use crate::mapping::{MappingSpec, PipelineStep, SyncMode, TableMapping};
use crate::mutation::{EdgeSet, GraphMutation, VertexSet};
use crate::transform::{TableContext, Transform};

/// Node label.
pub const LABEL: &str = "_label";
/// Node key — written as property [`ID_PROPERTY`].
pub const ID: &str = "_id";
/// Edge type.
pub const TYPE: &str = "_type";
/// Edge source label and key.
pub const FROM_LABEL: &str = "_from_label";
pub const FROM: &str = "_from";
/// Edge target label and key.
pub const TO_LABEL: &str = "_to_label";
pub const TO: &str = "_to";
/// Optional edge key (parallel edges) — written as property [`ID_PROPERTY`].
pub const EDGE_ID: &str = "_edge_id";
/// The property a SQL-made node is keyed on, and edge endpoints are matched on.
pub const ID_PROPERTY: &str = "id";

/// The name a round's rows are registered under; every statement's `FROM` is
/// rewritten to it.
const ROWS: &str = "src";

/// One SQL statement as a [`Transform`].
#[derive(Debug, Clone)]
pub struct SqlTransform {
    /// The statement over [`ROWS`].
    query: String,
    /// The same, without its `WHERE` — for delete rows.
    deletes_query: String,
    /// The table its `FROM` named.
    table: String,
}

impl SqlTransform {
    pub fn new(sql: &str) -> Result<Self> {
        let (mut query, table) = parse_select(sql)?;
        let full = query.to_string();
        if let SetExpr::Select(select) = query.body.as_mut() {
            select.selection = None;
        }
        Ok(Self { query: full, deletes_query: query.to_string(), table })
    }

    /// The table the statement's `FROM` names (`ns.name`).
    pub fn table(&self) -> &str {
        &self.table
    }

    fn run(&self, sql: &str, batches: &[RecordBatch]) -> Result<GraphMutation> {
        let Some(first) = batches.first() else { return Ok(GraphMutation::default()) };
        let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
        let rows = MemTable::try_new(first.schema(), vec![batches.to_vec()]).map_err(df)?;
        ctx.register_table(ROWS, Arc::new(rows)).map_err(df)?;
        // One partition, so the plan runs on this thread and needs no spawned
        // task — `Transform::apply` is synchronous and may be called from a
        // current-thread runtime, which a spawn would deadlock.
        let out = futures::executor::block_on(async { ctx.sql(sql).await?.collect().await }).map_err(df)?;
        to_mutation(&out)
    }
}

impl Transform for SqlTransform {
    fn apply(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
        self.run(&self.query, batches)
    }

    fn apply_deletes(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
        self.run(&self.deletes_query, batches)
    }
}

/// A `.sql` mapping file: `;`-separated statements, each a step of the table
/// its `FROM` names, every table in `mode`.
pub fn mapping_from_sql(text: &str, mode: SyncMode) -> Result<MappingSpec> {
    let statements = Parser::parse_sql(&GenericDialect {}, text).map_err(parse_err)?;
    if statements.is_empty() {
        return Err(BifrostError::Other("the SQL mapping holds no statement".into()));
    }
    let mut tables: Vec<TableMapping> = Vec::new();
    for statement in statements {
        let sql = statement.to_string();
        let (_, table) = parse_select(&sql)?;
        let step = PipelineStep::Sql { sql };
        match tables.iter_mut().find(|t| t.table == table) {
            Some(t) => t.pipeline.push(step),
            None => tables.push(TableMapping { table, mode, pipeline: vec![step], columns: Vec::new(), filter: None }),
        }
    }
    Ok(MappingSpec { version: 1, tables })
}

/// Parse exactly one `SELECT … FROM <one table>`, rewrite its `FROM` to
/// [`ROWS`], and return it with the table it named.
fn parse_select(sql: &str) -> Result<(Query, String)> {
    let bad = |why: &str| BifrostError::Other(format!("sql step: {why}: {sql}"));
    let mut statements = Parser::parse_sql(&GenericDialect {}, sql).map_err(parse_err)?;
    if statements.len() != 1 {
        return Err(bad("one statement per step"));
    }
    let Statement::Query(mut query) = statements.remove(0) else {
        return Err(bad("only SELECT"));
    };
    if query.with.is_some() {
        return Err(bad("no WITH - one table, row by row"));
    }
    let SetExpr::Select(select) = query.body.as_mut() else {
        return Err(bad("one SELECT - no UNION/INTERSECT/EXCEPT/VALUES"));
    };
    if select.from.len() != 1 || !select.from[0].joins.is_empty() {
        return Err(bad("exactly one table in FROM, no joins - a round carries one table's rows"));
    }
    let TableFactor::Table { name, .. } = &mut select.from[0].relation else {
        return Err(bad("FROM must name a table"));
    };
    let parts: Vec<String> = name
        .0
        .iter()
        .map(|p| match p {
            ObjectNamePart::Identifier(i) => Ok(i.value.clone()),
            _ => Err(bad("FROM must name a table")),
        })
        .collect::<Result<_>>()?;
    *name = ObjectName::from(vec![Ident::new(ROWS)]);
    Ok((*query, parts.join(".")))
}

fn parse_err(e: datafusion::sql::sqlparser::parser::ParserError) -> BifrostError {
    BifrostError::Other(format!("sql: {e}"))
}

fn df(e: datafusion::error::DataFusionError) -> BifrostError {
    BifrostError::Other(format!("sql: {e}"))
}

/// The statement's result as graph changes, by its reserved columns.
fn to_mutation(out: &[RecordBatch]) -> Result<GraphMutation> {
    let mut m = GraphMutation::default();
    for batch in out.iter().filter(|b| b.num_rows() > 0) {
        let has = |c: &str| batch.schema().column_with_name(c).is_some();
        if has(LABEL) && has(ID) {
            m.vertices.extend(vertices(batch)?);
        } else if [TYPE, FROM_LABEL, FROM, TO_LABEL, TO].iter().all(|c| has(c)) {
            m.edges.extend(edges(batch)?);
        } else {
            return Err(BifrostError::Other(format!(
                "sql step: the SELECT must return `{LABEL}`, `{ID}` (nodes) or `{TYPE}`, `{FROM_LABEL}`, `{FROM}`, \
                 `{TO_LABEL}`, `{TO}` (edges); it returned {:?}",
                batch.schema().fields().iter().map(|f| f.name()).collect::<Vec<_>>()
            )));
        }
    }
    Ok(m)
}

fn vertices(batch: &RecordBatch) -> Result<Vec<VertexSet>> {
    let label = text_column(batch, LABEL)?;
    let id = column(batch, ID)?;
    let keys: Vec<Option<String>> = (0..batch.num_rows())
        .map(|r| if label.is_null(r) || id.is_null(r) { None } else { Some(text_at(&label, r)) })
        .collect();
    let props = properties(batch, &[LABEL, ID])?;
    group(&keys)
        .into_iter()
        .map(|(label, rows)| {
            let mut cols = vec![(ID_PROPERTY.to_string(), id.clone())];
            cols.extend(props.iter().cloned());
            Ok(VertexSet { label, id_column: ID_PROPERTY.into(), batch: taken(&cols, &rows)? })
        })
        .collect()
}

fn edges(batch: &RecordBatch) -> Result<Vec<EdgeSet>> {
    let (ty, from_label, to_label) = (text_column(batch, TYPE)?, text_column(batch, FROM_LABEL)?, text_column(batch, TO_LABEL)?);
    let (from, to) = (column(batch, FROM)?, column(batch, TO)?);
    let edge_id = batch.schema().column_with_name(EDGE_ID).map(|(i, _)| batch.column(i).clone());
    let keys: Vec<Option<String>> = (0..batch.num_rows())
        .map(|r| {
            if [&ty, &from_label, &to_label, &from, &to].iter().any(|c| c.is_null(r)) {
                return None;
            }
            // \u{1f} cannot appear in a label; it only separates the three.
            Some(format!("{}\u{1f}{}\u{1f}{}", text_at(&ty, r), text_at(&from_label, r), text_at(&to_label, r)))
        })
        .collect();
    let props = properties(batch, &[TYPE, FROM_LABEL, FROM, TO_LABEL, TO, EDGE_ID])?;
    group(&keys)
        .into_iter()
        .map(|(key, rows)| {
            let mut parts = key.split('\u{1f}').map(str::to_string);
            let (edge_type, src_label, dst_label) = (parts.next().unwrap(), parts.next().unwrap(), parts.next().unwrap());
            let mut cols = vec![("src".to_string(), from.clone()), ("dst".to_string(), to.clone())];
            if let Some(e) = &edge_id {
                cols.push((ID_PROPERTY.to_string(), e.clone()));
            }
            cols.extend(props.iter().cloned());
            Ok(EdgeSet {
                src_label,
                edge_type,
                dst_label,
                src_id_column: ID_PROPERTY.into(),
                dst_id_column: ID_PROPERTY.into(),
                edge_key_props: if edge_id.is_some() { vec![ID_PROPERTY.into()] } else { Vec::new() },
                replace_outgoing: false,
                batch: taken(&cols, &rows)?,
            })
        })
        .collect()
}

/// Every non-reserved column, by name. One named like a reserved column's
/// output (`id`, `src`, `dst`) would be written twice — refused.
fn properties(batch: &RecordBatch, reserved: &[&str]) -> Result<Vec<(String, ArrayRef)>> {
    let schema = batch.schema();
    let mut out = Vec::new();
    for (i, f) in schema.fields().iter().enumerate() {
        let name = f.name().as_str();
        if reserved.contains(&name) {
            continue;
        }
        if [ID_PROPERTY, "src", "dst"].contains(&name) {
            return Err(BifrostError::DuplicateOutputColumn(format!(
                "{name} (a sql step writes `{ID}`/`{EDGE_ID}` as `{ID_PROPERTY}`; alias this column)"
            )));
        }
        out.push((name.to_string(), batch.column(i).clone()));
    }
    Ok(out)
}

fn column(batch: &RecordBatch, name: &str) -> Result<ArrayRef> {
    let (i, _) = batch.schema().column_with_name(name).ok_or_else(|| BifrostError::ColumnNotFound(name.into()))?;
    Ok(batch.column(i).clone())
}

/// A label/type column as text, whatever the SELECT typed it as.
fn text_column(batch: &RecordBatch, name: &str) -> Result<ArrayRef> {
    cast(&column(batch, name)?, &DataType::Utf8).map_err(|e| BifrostError::Other(format!("sql: `{name}` as text: {e}")))
}

fn text_at(col: &ArrayRef, row: usize) -> String {
    use arrow_array::cast::AsArray;
    col.as_string::<i32>().value(row).to_string()
}

/// Row indices per key, keys in first-seen order; `None` rows are dropped.
fn group(keys: &[Option<String>]) -> Vec<(String, Vec<u32>)> {
    let mut out: Vec<(String, Vec<u32>)> = Vec::new();
    for (row, key) in keys.iter().enumerate() {
        let Some(key) = key else { continue };
        match out.iter_mut().find(|(k, _)| k == key) {
            Some((_, rows)) => rows.push(row as u32),
            None => out.push((key.clone(), vec![row as u32])),
        }
    }
    out
}

/// `cols` at `rows`, as one batch.
fn taken(cols: &[(String, ArrayRef)], rows: &[u32]) -> Result<RecordBatch> {
    let idx = UInt32Array::from(rows.to_vec());
    let mut fields = Vec::with_capacity(cols.len());
    let mut arrays = Vec::with_capacity(cols.len());
    for (name, col) in cols {
        let a = take(col.as_ref(), &idx, None)?;
        fields.push(Field::new(name, a.data_type().clone(), true));
        arrays.push(a);
    }
    Ok(RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?)
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::{BooleanArray, Int64Array, StringArray};

    fn users(ids: &[i64], names: &[Option<&str>], company: &[Option<i64>], active: &[Option<bool>]) -> RecordBatch {
        RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new("user_id", DataType::Int64, false),
                Field::new("name", DataType::Utf8, true),
                Field::new("company_id", DataType::Int64, true),
                Field::new("active", DataType::Boolean, true),
            ])),
            vec![
                Arc::new(Int64Array::from(ids.to_vec())),
                Arc::new(StringArray::from(names.to_vec())),
                Arc::new(Int64Array::from(company.to_vec())),
                Arc::new(BooleanArray::from(active.to_vec())),
            ],
        )
        .unwrap()
    }

    fn ctx() -> TableContext {
        TableContext { table: "social.users".into() }
    }

    #[test]
    fn a_vertex_statement_filters_computes_and_splits_labels() {
        let t = SqlTransform::new(
            "SELECT CASE WHEN user_id < 3 THEN 'Staff' ELSE 'Person' END AS _label, user_id AS _id, \
             upper(name) AS name FROM social.users WHERE active",
        )
        .unwrap();
        assert_eq!(t.table(), "social.users");
        let b = users(&[1, 2, 3, 4], &[Some("a"), Some("b"), Some("c"), Some("d")], &[None; 4], &[Some(true), Some(false), Some(true), Some(true)]);
        let m = t.apply(&ctx(), &[b]).unwrap();
        let got: Vec<(String, usize)> = m.vertices.iter().map(|v| (v.label.clone(), v.batch.num_rows())).collect();
        assert_eq!(got, vec![("Staff".into(), 1), ("Person".into(), 2)], "user 2 is filtered out");
        let person = &m.vertices[1];
        assert_eq!(person.id_column, "id");
        let names: Vec<&str> = person.batch.column_by_name("name").unwrap().as_any().downcast_ref::<StringArray>().unwrap().iter().flatten().collect();
        assert_eq!(names, vec!["C", "D"]);
    }

    #[test]
    fn deletes_ignore_the_where_so_a_key_only_row_still_deletes() {
        let t = SqlTransform::new("SELECT 'Person' AS _label, user_id AS _id, name FROM users WHERE active").unwrap();
        // An equality delete: only the key is set.
        let b = users(&[7], &[None], &[None], &[None]);
        assert!(t.apply(&ctx(), std::slice::from_ref(&b)).unwrap().vertices.is_empty(), "the WHERE drops it as an upsert");
        let d = t.apply_deletes(&ctx(), &[b]).unwrap();
        assert_eq!(d.vertices.len(), 1);
        assert_eq!(d.vertices[0].batch.num_rows(), 1);
    }

    #[test]
    fn an_edge_statement_names_its_endpoints_and_skips_null_ones() {
        let t = SqlTransform::new(
            "SELECT 'WORKS_AT' AS _type, 'Person' AS _from_label, user_id AS _from, 'Company' AS _to_label, \
             company_id AS _to, user_id * 10 AS _edge_id, name AS who FROM users",
        )
        .unwrap();
        let b = users(&[1, 2], &[Some("a"), Some("b")], &[Some(9), None], &[Some(true), Some(true)]);
        let m = t.apply(&ctx(), &[b]).unwrap();
        assert_eq!(m.edges.len(), 1);
        let e = &m.edges[0];
        assert_eq!((e.src_label.as_str(), e.edge_type.as_str(), e.dst_label.as_str()), ("Person", "WORKS_AT", "Company"));
        assert_eq!(e.edge_key_props, vec!["id".to_string()]);
        assert_eq!(e.batch.num_rows(), 1, "the null company writes no edge");
        let names: Vec<String> = e.batch.schema().fields().iter().map(|f| f.name().clone()).collect();
        assert_eq!(names, vec!["src", "dst", "id", "who"]);
    }

    #[test]
    fn one_table_row_by_row_is_enforced() {
        for sql in [
            "SELECT 1 AS _label, 1 AS _id FROM a JOIN b ON a.x = b.x",
            "SELECT 1 AS _label, 1 AS _id FROM a UNION SELECT 1, 1 FROM b",
            "WITH c AS (SELECT * FROM a) SELECT 1 AS _label, 1 AS _id FROM c",
            "DELETE FROM a",
        ] {
            assert!(SqlTransform::new(sql).is_err(), "{sql}");
        }
    }

    #[test]
    fn a_select_without_reserved_columns_says_what_it_needs() {
        let t = SqlTransform::new("SELECT user_id, name FROM users").unwrap();
        let b = users(&[1], &[Some("a")], &[None], &[Some(true)]);
        let e = t.apply(&ctx(), &[b]).unwrap_err().to_string();
        assert!(e.contains("_label") && e.contains("_from_label"), "{e}");
    }

    #[test]
    fn a_sql_file_groups_statements_by_table() {
        let m = mapping_from_sql(
            "SELECT 'Person' AS _label, user_id AS _id FROM social.users;\n\
             SELECT 'Company' AS _label, id AS _id FROM social.companies;\n\
             SELECT 'WORKS_AT' AS _type, 'Person' AS _from_label, user_id AS _from, 'Company' AS _to_label, company_id AS _to FROM social.users;",
            SyncMode::Cdc,
        )
        .unwrap();
        let got: Vec<(&str, usize)> = m.tables.iter().map(|t| (t.table.as_str(), t.pipeline.len())).collect();
        assert_eq!(got, vec![("social.users", 2), ("social.companies", 1)]);
        assert!(m.tables.iter().all(|t| t.mode == SyncMode::Cdc));
    }
}