crates/knut-bifrost/src/sql.rs
//! **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));
}
}