use serde::{Deserialize, Serialize};
use std::sync::Arc;
use crate::{
error::{BifrostError, Result},
registry::TransformRegistry,
transform::{EdgeTransform, Transform, VertexTransform},
};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MappingSpec {
pub version: u32,
pub tables: Vec<TableMapping>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableMapping {
pub table: String,
#[serde(default)]
pub mode: SyncMode,
pub pipeline: Vec<PipelineStep>,
#[serde(default)]
pub columns: Vec<String>,
#[serde(default)]
pub filter: Option<RowFilter>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum FilterValue {
Bool(bool),
Int(i64),
Float(f64),
Str(String),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RowFilter {
pub column: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub eq: Option<FilterValue>,
#[serde(rename = "in", default, skip_serializing_if = "Option::is_none")]
pub is_in: Option<Vec<FilterValue>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ge: Option<FilterValue>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub le: Option<FilterValue>,
}
impl RowFilter {
pub fn eq(column: impl Into<String>, value: impl Into<FilterValue>) -> Self {
Self {
column: column.into(),
eq: Some(value.into()),
is_in: None,
ge: None,
le: None,
}
}
pub fn is_in<V: Into<FilterValue>>(
column: impl Into<String>,
values: impl IntoIterator<Item = V>,
) -> Self {
Self {
column: column.into(),
eq: None,
is_in: Some(values.into_iter().map(Into::into).collect()),
ge: None,
le: None,
}
}
pub fn range(
column: impl Into<String>,
lo: Option<impl Into<FilterValue>>,
hi: Option<impl Into<FilterValue>>,
) -> Self {
Self {
column: column.into(),
eq: None,
is_in: None,
ge: lo.map(Into::into),
le: hi.map(Into::into),
}
}
pub fn validate(&self) -> Result<()> {
let forms = usize::from(self.eq.is_some())
+ usize::from(self.is_in.is_some())
+ usize::from(self.ge.is_some() || self.le.is_some());
match forms {
1 => Ok(()),
0 => Err(BifrostError::Other(format!(
"filter on `{}` sets none of `eq` / `in` / `ge`+`le` — it would prune nothing",
self.column
))),
_ => Err(BifrostError::Other(format!(
"filter on `{}` sets more than one of `eq` / `in` / `ge`+`le`",
self.column
))),
}
}
}
impl From<&str> for FilterValue {
fn from(s: &str) -> Self {
FilterValue::Str(s.to_string())
}
}
impl From<String> for FilterValue {
fn from(s: String) -> Self {
FilterValue::Str(s)
}
}
impl From<i64> for FilterValue {
fn from(v: i64) -> Self {
FilterValue::Int(v)
}
}
impl From<f64> for FilterValue {
fn from(v: f64) -> Self {
FilterValue::Float(v)
}
}
impl From<bool> for FilterValue {
fn from(v: bool) -> Self {
FilterValue::Bool(v)
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SyncMode {
#[default]
Full,
Incremental,
Cdc,
}
impl SyncMode {
pub const CHANGE_TYPE_COLUMN: &'static str = "_change_type";
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PipelineStep {
Vertex {
vertex: VertexTransform,
},
Edge {
edge: EdgeTransform,
},
FkEdges {
fk_edges: EdgeTransform,
},
Custom {
custom: String,
},
Sql {
sql: String,
},
}
impl MappingSpec {
pub fn from_yaml(yaml: &str) -> Result<Self> {
Ok(serde_yaml::from_str(yaml)?)
}
pub fn load(path: impl AsRef<std::path::Path>) -> Result<Self> {
Self::from_yaml(&std::fs::read_to_string(path)?)
}
pub fn from_ron(ron: &str) -> Result<Self> {
Ok(ron::from_str(ron)?)
}
pub fn to_ron(&self) -> Result<String> {
Ok(ron::ser::to_string_pretty(
self,
ron::ser::PrettyConfig::default(),
)?)
}
pub fn load_ron(path: impl AsRef<std::path::Path>) -> Result<Self> {
Self::from_ron(&std::fs::read_to_string(path)?)
}
pub fn from_arrow_schema(table: &str, schema: &arrow_schema::Schema) -> SchemaMappingCandidate {
let label = label_for_table(table);
let mut unmapped: Vec<UnmappedColumn> = Vec::new();
let mut guesses: Vec<String> = Vec::new();
let mut scalar_cols: Vec<String> = Vec::new();
for f in schema.fields() {
if is_scalar(f.data_type()) {
scalar_cols.push(f.name().clone());
} else {
unmapped.push(UnmappedColumn {
name: f.name().clone(),
arrow_type: format!("{}", f.data_type()),
reason: nested_reason(f.data_type()),
});
}
}
let table_fk = format!("{}_id", singularize(&last_segment(table)));
let id_column = scalar_cols
.iter()
.find(|c| c.as_str() == "id")
.or_else(|| scalar_cols.iter().find(|c| **c == table_fk))
.or_else(|| scalar_cols.iter().find(|c| c.ends_with("_id")))
.or_else(|| scalar_cols.first())
.cloned();
let Some(id_column) = id_column else {
guesses.push(format!(
"no scalar column to use as `{label}` id — every column is nested (see `unmapped`)"
));
return SchemaMappingCandidate {
spec: MappingSpec {
version: 1,
tables: vec![TableMapping {
table: table.to_string(),
mode: SyncMode::default(),
pipeline: Vec::new(),
columns: Vec::new(),
filter: None,
}],
},
unmapped,
guesses,
};
};
guesses.push(format!("id column for `{label}` guessed as `{id_column}`"));
let fk_cols: Vec<String> = scalar_cols
.iter()
.filter(|c| **c != id_column && c.ends_with("_id"))
.cloned()
.collect();
let properties: Vec<String> = scalar_cols
.iter()
.filter(|c| **c != id_column && !fk_cols.contains(c))
.cloned()
.collect();
let mut pipeline = vec![PipelineStep::Vertex {
vertex: VertexTransform {
label: label.clone(),
id_column: id_column.clone(),
properties,
rename: Default::default(),
},
}];
for fk in &fk_cols {
let dst_label = label_for_fk(fk);
let edge_type = format!("HAS_{}", fk.trim_end_matches("_id").to_uppercase());
guesses.push(format!(
"edge `{edge_type}` guessed from FK column `{fk}` → `{dst_label}` \
(verify the target label + endpoint id column)"
));
pipeline.push(PipelineStep::FkEdges {
fk_edges: EdgeTransform {
edge_type,
src: EndpointSpec {
label: label.clone(),
column: id_column.clone(),
id_column: None,
},
dst: EndpointSpec {
label: dst_label,
column: fk.clone(),
id_column: None,
},
properties: Vec::new(),
rename: Default::default(),
key: Vec::new(),
replace: None,
},
});
}
SchemaMappingCandidate {
spec: MappingSpec {
version: 1,
tables: vec![TableMapping {
table: table.to_string(),
mode: SyncMode::default(),
pipeline,
columns: Vec::new(),
filter: None,
}],
},
unmapped,
guesses,
}
}
}
use crate::transform::EndpointSpec;
#[derive(Debug, Clone)]
pub struct SchemaMappingCandidate {
pub spec: MappingSpec,
pub unmapped: Vec<UnmappedColumn>,
pub guesses: Vec<String>,
}
impl SchemaMappingCandidate {
pub fn is_fully_mapped(&self) -> bool {
self.unmapped.is_empty()
}
pub fn to_ron(&self) -> Result<String> {
self.spec.to_ron()
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct UnmappedColumn {
pub name: String,
pub arrow_type: String,
pub reason: String,
}
fn is_scalar(dt: &arrow_schema::DataType) -> bool {
use arrow_schema::DataType as D;
matches!(
dt,
D::Boolean
| D::Int8
| D::Int16
| D::Int32
| D::Int64
| D::UInt8
| D::UInt16
| D::UInt32
| D::UInt64
| D::Float16
| D::Float32
| D::Float64
| D::Utf8
| D::LargeUtf8
| D::Utf8View
| D::Date32
| D::Date64
| D::Timestamp(_, _)
| D::Decimal128(_, _)
| D::Decimal256(_, _)
)
}
fn nested_reason(dt: &arrow_schema::DataType) -> String {
use arrow_schema::DataType as D;
let kind = match dt {
D::Struct(_) => "a struct",
D::List(_)
| D::LargeList(_)
| D::FixedSizeList(_, _)
| D::ListView(_)
| D::LargeListView(_) => "a list",
D::Map(_, _) => "a map",
D::Union(_, _) => "a union",
_ => "a non-scalar",
};
format!(
"{kind} type is not a scalar graph property — map it by hand (flatten, or a custom transform)"
)
}
fn last_segment(table: &str) -> String {
table.rsplit('.').next().unwrap_or(table).to_string()
}
fn singularize(word: &str) -> String {
if let Some(stem) = word.strip_suffix("ies") {
format!("{stem}y")
} else if word.ends_with("ss") {
word.to_string()
} else if let Some(stem) = word.strip_suffix('s') {
stem.to_string()
} else {
word.to_string()
}
}
fn pascal_case(word: &str) -> String {
word.split('_')
.filter(|s| !s.is_empty())
.map(|s| {
let mut c = s.chars();
match c.next() {
Some(first) => first.to_uppercase().collect::<String>() + c.as_str(),
None => String::new(),
}
})
.collect()
}
fn label_for_table(table: &str) -> String {
pascal_case(&singularize(&last_segment(table)))
}
fn label_for_fk(fk: &str) -> String {
pascal_case(fk.trim_end_matches("_id"))
}
impl TableMapping {
pub fn projected_columns(&self) -> Option<Vec<String>> {
if !self.columns.is_empty() {
return Some(dedup_preserving_order(self.columns.iter().cloned()));
}
let mut cols: Vec<String> = Vec::new();
for step in &self.pipeline {
match step {
PipelineStep::Vertex { vertex } => {
cols.push(vertex.id_column.clone());
cols.extend(vertex.properties.iter().cloned());
}
PipelineStep::Edge { edge: t } | PipelineStep::FkEdges { fk_edges: t } => {
cols.push(t.src.column.clone());
cols.push(t.dst.column.clone());
cols.extend(t.properties.iter().cloned());
cols.extend(t.key.iter().cloned());
}
PipelineStep::Custom { .. } | PipelineStep::Sql { .. } => return None,
}
}
let cols = dedup_preserving_order(cols);
if cols.is_empty() { None } else { Some(cols) }
}
pub fn compile(&self, registry: &TransformRegistry) -> Result<Vec<Arc<dyn Transform>>> {
if self.pipeline.is_empty() {
return Err(BifrostError::EmptyPipeline(self.table.clone()));
}
self.pipeline
.iter()
.map(|step| -> Result<Arc<dyn Transform>> {
Ok(match step {
PipelineStep::Vertex { vertex } => Arc::new(vertex.clone()),
PipelineStep::Edge { edge: t } => Arc::new(t.clone()),
PipelineStep::FkEdges { fk_edges: t } => {
let entity_row = self.pipeline.iter().any(|s| {
matches!(s, PipelineStep::Vertex { vertex } if vertex.label == t.src.label && vertex.id_column == t.src.column)
});
Arc::new(EdgeTransform { replace: Some(t.replace.unwrap_or(entity_row)), ..t.clone() })
}
PipelineStep::Custom { custom } => registry.get(custom)?,
#[cfg(feature = "sql")]
PipelineStep::Sql { sql } => Arc::new(crate::sql::SqlTransform::new(sql)?),
#[cfg(not(feature = "sql"))]
PipelineStep::Sql { .. } => {
return Err(BifrostError::Other(format!(
"table `{}` has a `sql:` step and this build has no SQL engine - rebuild with --features sql",
self.table
)));
}
})
})
.collect()
}
}
fn dedup_preserving_order(items: impl IntoIterator<Item = String>) -> Vec<String> {
let mut seen = std::collections::HashSet::new();
items
.into_iter()
.filter(|c| seen.insert(c.clone()))
.collect()
}
pub fn apply_row_filter(
batches: &[arrow_array::RecordBatch],
filter: &RowFilter,
) -> Result<Vec<arrow_array::RecordBatch>> {
apply_row_filter_with(batches, filter, FilterNulls::Drop)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FilterNulls {
Drop,
Keep,
}
pub fn apply_row_filter_with(
batches: &[arrow_array::RecordBatch],
filter: &RowFilter,
nulls: FilterNulls,
) -> Result<Vec<arrow_array::RecordBatch>> {
filter.validate()?;
batches
.iter()
.map(|b| filter_batch(b, filter, nulls))
.collect()
}
fn filter_batch(
batch: &arrow_array::RecordBatch,
filter: &RowFilter,
nulls: FilterNulls,
) -> Result<arrow_array::RecordBatch> {
use arrow::compute::kernels::{boolean, cmp};
let idx = batch.schema().index_of(&filter.column).map_err(|_| {
BifrostError::Other(format!(
"filter column `{}` is not in the batch schema ({:?})",
filter.column,
batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect::<Vec<_>>()
))
})?;
let col = batch.column(idx).clone();
let mask = if let Some(v) = &filter.eq {
cmp::eq(&col, &scalar_datum(v, &col)?).map_err(arrow_err)?
} else if let Some(vs) = &filter.is_in {
let mut acc: Option<arrow_array::BooleanArray> = None;
for v in vs {
let m = cmp::eq(&col, &scalar_datum(v, &col)?).map_err(arrow_err)?;
acc = Some(match acc {
None => m,
Some(a) => boolean::or(&a, &m).map_err(arrow_err)?,
});
}
acc.unwrap_or_else(|| arrow_array::BooleanArray::from(vec![false; batch.num_rows()]))
} else {
let mut acc: Option<arrow_array::BooleanArray> = None;
if let Some(lo) = &filter.ge {
acc = Some(cmp::gt_eq(&col, &scalar_datum(lo, &col)?).map_err(arrow_err)?);
}
if let Some(hi) = &filter.le {
let m = cmp::lt_eq(&col, &scalar_datum(hi, &col)?).map_err(arrow_err)?;
acc = Some(match acc {
None => m,
Some(a) => boolean::and(&a, &m).map_err(arrow_err)?,
});
}
acc.expect("validate() guarantees a range filter sets ge and/or le")
};
let mask = match nulls {
FilterNulls::Drop => mask,
FilterNulls::Keep => {
use arrow_array::Array;
(0..batch.num_rows())
.map(|r| Some(col.is_null(r) || (mask.is_valid(r) && mask.value(r))))
.collect::<arrow_array::BooleanArray>()
}
};
arrow::compute::filter_record_batch(batch, &mask).map_err(arrow_err)
}
fn scalar_datum(
v: &FilterValue,
col: &arrow_array::ArrayRef,
) -> Result<arrow_array::Scalar<arrow_array::ArrayRef>> {
use arrow_array::{
ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, Scalar, StringArray,
};
use arrow_schema::DataType as D;
let dt = col.data_type();
let arr: ArrayRef = match (v, dt) {
(FilterValue::Str(s), D::Utf8 | D::LargeUtf8) => {
Arc::new(StringArray::from(vec![s.clone()]))
}
(FilterValue::Int(i), D::Int64) => Arc::new(Int64Array::from(vec![*i])),
(FilterValue::Int(i), D::Int32) => Arc::new(Int32Array::from(vec![*i as i32])),
(FilterValue::Float(f), D::Float64) => Arc::new(Float64Array::from(vec![*f])),
(FilterValue::Int(i), D::Float64) => Arc::new(Float64Array::from(vec![*i as f64])),
(FilterValue::Bool(b), D::Boolean) => Arc::new(BooleanArray::from(vec![*b])),
_ => {
return Err(BifrostError::Other(format!(
"filter literal {v:?} is not comparable with column type {dt:?} — \
refusing to coerce (a coerced predicate selects the wrong rows)"
)));
}
};
Ok(Scalar::new(arr))
}
fn arrow_err(e: arrow_schema::ArrowError) -> BifrostError {
BifrostError::Other(format!("residual row filter: {e}"))
}
pub trait GraphMapping {
fn compile(&self) -> Result<MappingSpec>;
}
impl GraphMapping for MappingSpec {
fn compile(&self) -> Result<MappingSpec> {
Ok(self.clone())
}
}
#[derive(Debug, Clone)]
pub struct PuppyGraphMapping {
pub json: String,
pub mode: SyncMode,
}
impl GraphMapping for PuppyGraphMapping {
fn compile(&self) -> Result<MappingSpec> {
crate::puppygraph::mapping_from_puppygraph(&self.json, self.mode)
}
}
#[derive(Debug, Clone)]
pub struct SqlMapping {
pub sql: String,
pub mode: SyncMode,
}
impl GraphMapping for SqlMapping {
#[cfg(feature = "sql")]
fn compile(&self) -> Result<MappingSpec> {
crate::sql::mapping_from_sql(&self.sql, self.mode)
}
#[cfg(not(feature = "sql"))]
fn compile(&self) -> Result<MappingSpec> {
Err(BifrostError::Other("a .sql mapping needs the SQL engine - rebuild with --features sql".into()))
}
}
pub fn load_mapping(path: impl AsRef<std::path::Path>, mode: SyncMode) -> Result<MappingSpec> {
let path = path.as_ref();
let text = std::fs::read_to_string(path)?;
let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("").to_ascii_lowercase();
match ext.as_str() {
"json" => PuppyGraphMapping { json: text, mode }.compile(),
"sql" => SqlMapping { sql: text, mode }.compile(),
"ron" => MappingSpec::from_ron(&text),
_ => MappingSpec::from_yaml(&text),
}
}
#[cfg(test)]
mod ron_seam_tests {
use super::*;
use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit};
#[test]
fn ron_round_trips_and_matches_yaml() {
let yaml = r#"
version: 1
tables:
- table: lake.social.users
mode: incremental
pipeline:
- vertex: { label: Person, id_column: user_id, properties: [name] }
- fk_edges:
type: WORKS_AT
src: { label: Person, column: user_id }
dst: { label: Company, column: company_id }
properties: [since]
rename: { since: joined_at }
"#;
let from_yaml = MappingSpec::from_yaml(yaml).expect("yaml parses");
let ron = from_yaml.to_ron().expect("to_ron");
let from_ron = MappingSpec::from_ron(&ron).expect("from_ron parses our own RON");
assert_eq!(
from_ron.to_ron().unwrap(),
from_yaml.to_ron().unwrap(),
"YAML→RON→spec equals the YAML spec"
);
assert_eq!(from_ron.tables.len(), 1);
assert_eq!(from_ron.tables[0].table, "lake.social.users");
assert!(matches!(from_ron.tables[0].mode, SyncMode::Incremental));
assert_eq!(from_ron.tables[0].pipeline.len(), 2);
}
#[test]
fn schema_infers_candidate_mapping_and_emits_ron() {
let schema = Schema::new(vec![
Field::new("user_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
Field::new("age", DataType::Int32, true),
Field::new("company_id", DataType::Int64, true), Field::new(
"created_at",
DataType::Timestamp(TimeUnit::Microsecond, None),
true,
),
]);
let cand = MappingSpec::from_arrow_schema("lake.social.users", &schema);
assert!(
cand.is_fully_mapped(),
"all scalar columns mapped, none unmapped"
);
let tm = &cand.spec.tables[0];
assert_eq!(tm.table, "lake.social.users");
let PipelineStep::Vertex { vertex } = &tm.pipeline[0] else {
panic!("first step is the vertex");
};
assert_eq!(vertex.label, "User");
assert_eq!(vertex.id_column, "user_id");
assert_eq!(vertex.properties, vec!["name", "age", "created_at"]);
let PipelineStep::FkEdges { fk_edges } = &tm.pipeline[1] else {
panic!("second step is the FK edge");
};
assert_eq!(fk_edges.edge_type, "HAS_COMPANY");
assert_eq!(fk_edges.src.label, "User");
assert_eq!(fk_edges.dst.label, "Company");
assert_eq!(fk_edges.dst.column, "company_id");
assert!(cand.guesses.iter().any(|g| g.contains("id column")));
assert!(cand.guesses.iter().any(|g| g.contains("HAS_COMPANY")));
let ron = cand.to_ron().expect("candidate → RON");
let reparsed = MappingSpec::from_ron(&ron).expect("candidate RON re-parses");
assert_eq!(reparsed.tables[0].pipeline.len(), 2);
}
#[test]
fn nested_column_is_surfaced_not_dropped() {
let addr = Fields::from(vec![
Field::new("city", DataType::Utf8, true),
Field::new("zip", DataType::Utf8, true),
]);
let schema = Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
Field::new("address", DataType::Struct(addr), true), Field::new(
"tags",
DataType::List(Field::new("item", DataType::Utf8, true).into()),
true,
),
]);
let cand = MappingSpec::from_arrow_schema("lake.crm.accounts", &schema);
assert!(!cand.is_fully_mapped(), "there are unmapped nested columns");
let names: Vec<&str> = cand.unmapped.iter().map(|u| u.name.as_str()).collect();
assert!(
names.contains(&"address"),
"the struct column is surfaced: {names:?}"
);
assert!(
names.contains(&"tags"),
"the list column is surfaced: {names:?}"
);
let PipelineStep::Vertex { vertex } = &cand.spec.tables[0].pipeline[0] else {
panic!("vertex step");
};
let mut accounted: Vec<String> = vec![vertex.id_column.clone()];
accounted.extend(vertex.properties.iter().cloned());
accounted.extend(cand.unmapped.iter().map(|u| u.name.clone()));
accounted.sort();
assert_eq!(
accounted,
vec!["address", "id", "name", "tags"],
"no column silently dropped — mapped or surfaced, every one"
);
assert!(cand.unmapped.iter().any(|u| u.reason.contains("struct")));
}
fn replaces(yaml: &str) -> bool {
use arrow_array::{Int64Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
let spec = MappingSpec::from_yaml(yaml).unwrap();
let steps = spec.tables[0].compile(&TransformRegistry::new()).unwrap();
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false), Field::new("b", DataType::Int64, false)])),
vec![Arc::new(Int64Array::from(vec![1])), Arc::new(Int64Array::from(vec![2]))],
)
.unwrap();
let ctx = crate::transform::TableContext { table: "t".into() };
steps.iter().flat_map(|s| s.apply(&ctx, std::slice::from_ref(&batch)).unwrap().edges).any(|e| e.replace_outgoing)
}
#[test]
fn fk_edges_replace_only_on_an_entity_row() {
let entity = "version: 1\ntables:\n - table: t\n pipeline:\n - vertex: { label: A, id_column: a }\n \
- fk_edges: { type: T, src: { label: A, column: a }, dst: { label: B, column: b } }\n";
assert!(replaces(entity), "the row is its source node: one edge per source");
let junction = "version: 1\ntables:\n - table: t\n pipeline:\n \
- fk_edges: { type: T, src: { label: A, column: a }, dst: { label: B, column: b } }\n";
assert!(!replaces(junction), "a junction table keeps every edge");
let opted_out = entity.replace("dst: { label: B, column: b } }", "dst: { label: B, column: b }, replace: false }");
assert!(!replaces(&opted_out));
let edge_opt_in = junction.replace("fk_edges:", "edge:").replace("dst: { label: B, column: b } }", "dst: { label: B, column: b }, replace: true }");
assert!(replaces(&edge_opt_in), "an edge step may opt in");
}
}