crates/knut-bifrost/src/mapping.rs
//! Declarative YAML mapping spec — the fabric's blueprint.
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use crate::{
error::{BifrostError, Result},
registry::TransformRegistry,
transform::{EdgeTransform, Transform, VertexTransform},
};
/// Top-level mapping file.
///
/// ```yaml
/// 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] # edge property columns
/// rename: { since: joined_at } # optional column → prop renaming
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MappingSpec {
pub version: u32,
pub tables: Vec<TableMapping>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableMapping {
/// Fully qualified table name as Spark sees it.
pub table: String,
#[serde(default)]
pub mode: SyncMode,
pub pipeline: Vec<PipelineStep>,
/// **Projection override.** The columns the source is allowed to read. Empty
/// (the default) means *derive it from the pipeline* — see
/// [`projected_columns`](TableMapping::projected_columns) — which is what you
/// want almost always; set it only to read a column the pipeline does not
/// name (a custom transform's input, say).
///
/// ```yaml
/// columns: [id, name, country]
/// ```
#[serde(default)]
pub columns: Vec<String>,
/// **Predicate pushdown.** An optional row filter pushed into the Iceberg
/// scan planner, so partitions / data files / row groups that cannot match
/// are never opened.
///
/// ```yaml
/// filter: { column: country, eq: SE }
/// ```
#[serde(default)]
pub filter: Option<RowFilter>,
}
/// A literal in a [`RowFilter`]. Untagged, so YAML writes the natural scalar
/// (`eq: SE`, `eq: 42`, `eq: true`) with no type tag.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum FilterValue {
Bool(bool),
Int(i64),
Float(f64),
Str(String),
}
/// A **backend-neutral row filter** on one column — the YAML face of the scan
/// predicate the Iceberg source pushes down.
///
/// Backend-neutral on purpose: `bifrost`'s core carries no engine dependency (the
/// `skade` source is an optional feature), so the spec type lives here and is
/// lowered to `skade::ScanFilter` only inside the Iceberg source. Exactly ONE of
/// `eq` / `in` / (`ge`,`le`) must be set.
///
/// ```yaml
/// filter: { column: country, eq: SE } # equality (partition pruning)
/// filter: { column: country, in: [SE, NO, DK] } # set membership
/// filter: { column: year, ge: 2020, le: 2026 } # inclusive range
/// ```
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RowFilter {
/// The column the predicate is on.
pub column: String,
/// `column == value`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub eq: Option<FilterValue>,
/// `column IN (values…)`.
#[serde(rename = "in", default, skip_serializing_if = "Option::is_none")]
pub is_in: Option<Vec<FilterValue>>,
/// Inclusive lower bound (`value <= column`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ge: Option<FilterValue>,
/// Inclusive upper bound (`column <= value`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub le: Option<FilterValue>,
}
impl RowFilter {
/// `column == value`.
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,
}
}
/// `column IN (values…)`.
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,
}
}
/// Inclusive range `lo <= column <= hi`; either bound may be `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),
}
}
/// Reject a filter that names no form (or more than one) — a mapping that
/// says `filter: { column: x }` silently matching everything would be a
/// pushdown that quietly does nothing.
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 {
/// Re-read the whole table when its snapshot changes (idempotent MERGE).
#[default]
Full,
/// Read only what changed since the last processed snapshot. Over Spark
/// Connect it is an `EXCEPT ALL` diff — append-only, inserts alone. Over an
/// Iceberg table read by skade (`watch`) it is the same per-commit changelog
/// as [`SyncMode::Cdc`]: updates and deletes are applied too.
Incremental,
/// Changelog (CDC) mode: read Iceberg's `system.create_changelog_view`
/// between the last and the current snapshot, exposing a `_change_type`
/// column (`INSERT` / `UPDATE_AFTER` / `UPDATE_BEFORE` / `DELETE`).
/// `INSERT`/`UPDATE_AFTER` rows become upserts; `DELETE` rows become
/// `DETACH DELETE`s. `UPDATE_BEFORE` (the update pre-image) is dropped — its
/// `UPDATE_AFTER` upsert already carries the new state. The first round (no
/// prior snapshot) reads the whole table as inserts, like `Full`.
Cdc,
}
impl SyncMode {
/// The sentinel column the Iceberg changelog view adds in [`Cdc`] mode,
/// and that the listener partitions on before stripping it from values.
///
/// [`Cdc`]: SyncMode::Cdc
pub const CHANGE_TYPE_COLUMN: &'static str = "_change_type";
}
/// One `- vertex: …` / `- edge: …` / `- fk_edges: …` / `- custom: name`
/// entry. Untagged so the YAML reads as a single-key map per step.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PipelineStep {
Vertex {
vertex: VertexTransform,
},
Edge {
edge: EdgeTransform,
},
/// Same shape as `edge`; named for FK-derived relationships.
FkEdges {
fk_edges: EdgeTransform,
},
/// Reference to a transform registered in [`TransformRegistry`].
Custom {
custom: String,
},
/// A SQL `SELECT` over the table's rows whose result names what it becomes
/// (`_label` + `_id` → nodes, `_type` + `_from_label`/`_from` +
/// `_to_label`/`_to` → edges). See [`crate::sql`]. Feature `sql`.
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)?)
}
/// Parse a [`MappingSpec`] from **RON** — the constellation's
/// RON-as-description form, sibling to [`from_yaml`](Self::from_yaml).
///
/// RON parity closes the seam the findings doc named: the Iceberg→graph
/// mapping is now a first-class RON object like every other spec, and RON is
/// what [`from_arrow_schema`](Self::from_arrow_schema) emits. YAML stays
/// accepted — this is additive.
pub fn from_ron(ron: &str) -> Result<Self> {
Ok(ron::from_str(ron)?)
}
/// Render this spec to **RON** text (pretty, stable field order via the
/// struct derive) — the round-trip partner of [`from_ron`](Self::from_ron).
pub fn to_ron(&self) -> Result<String> {
Ok(ron::ser::to_string_pretty(
self,
ron::ser::PrettyConfig::default(),
)?)
}
/// Load a [`MappingSpec`] from a `.ron` file.
pub fn load_ron(path: impl AsRef<std::path::Path>) -> Result<Self> {
Self::from_ron(&std::fs::read_to_string(path)?)
}
/// **Derive a *candidate* `MappingSpec` from an Iceberg table's Arrow
/// schema** — the schema→mapping codegen half of the RON seam.
///
/// The generator emits ONE vertex per table, a best-guess id column, the
/// remaining scalar columns as vertex properties, and one `fk_edges`
/// candidate per foreign-key-shaped (`*_id`) column. It is meant to be
/// emitted as RON, reviewed, and corrected by a human — it is a *starting
/// point*, not an oracle.
///
/// **Honest by construction (the findings-doc law): no column is ever
/// silently dropped.** Every input field is accounted for exactly once —
/// mapped as the id, a property, or an FK edge endpoint, or else SURFACED in
/// [`SchemaMappingCandidate::unmapped`] (a nested/complex type that cannot
/// become a scalar graph property). Every *guessed* role (the id column,
/// each FK edge) is recorded in [`SchemaMappingCandidate::guesses`] so a
/// reviewer knows precisely what to check.
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();
// Partition scalar vs non-scalar columns. Non-scalar (nested) columns
// cannot be a scalar graph property → surfaced, never dropped.
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()),
});
}
}
// Pick the id column: exact `id`, else `<table_singular>_id`, else the
// first `*_id`, else the first scalar column. Always a GUESS.
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 {
// A table with only nested columns: nothing to key a vertex on. Emit
// an empty-pipeline candidate; everything is in `unmapped`.
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}`"));
// FK-shaped columns (other `*_id`) → edge candidates.
let fk_cols: Vec<String> = scalar_cols
.iter()
.filter(|c| **c != id_column && c.ends_with("_id"))
.cloned()
.collect();
// Remaining scalar columns become vertex properties (id + FKs excluded).
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,
}
}
}
// ═══════════════════════════════════════════════════════════════════════════
// Schema → candidate mapping (codegen)
// ═══════════════════════════════════════════════════════════════════════════
use crate::transform::EndpointSpec;
/// The outcome of [`MappingSpec::from_arrow_schema`] — a candidate mapping plus
/// the honest accounting of what the generator guessed and what it could not map.
///
/// Honest by construction: `unmapped` + every field the `spec` references cover
/// **every** input column. A reviewer reads `guesses` to see what to check and
/// `unmapped` to see what to add by hand; nothing was silently dropped.
#[derive(Debug, Clone)]
pub struct SchemaMappingCandidate {
/// The generated candidate mapping — emit it with [`MappingSpec::to_ron`].
pub spec: MappingSpec,
/// Columns the generator could not map to a scalar graph property (nested /
/// complex Arrow types). Surfaced for a human, **never dropped**.
pub unmapped: Vec<UnmappedColumn>,
/// Every role the generator GUESSED (the id column, each FK edge) — the
/// review checklist.
pub guesses: Vec<String>,
}
impl SchemaMappingCandidate {
/// `true` when the generator mapped every column with no unmapped residue
/// (there may still be [`guesses`](Self::guesses) to review).
pub fn is_fully_mapped(&self) -> bool {
self.unmapped.is_empty()
}
/// Render the candidate spec to RON (convenience for `spec.to_ron()`).
pub fn to_ron(&self) -> Result<String> {
self.spec.to_ron()
}
}
/// One column [`MappingSpec::from_arrow_schema`] could not confidently map —
/// surfaced, not dropped.
#[derive(Debug, Clone, PartialEq)]
pub struct UnmappedColumn {
/// The column name.
pub name: String,
/// Its Arrow type, rendered.
pub arrow_type: String,
/// Why it could not be mapped to a scalar graph property.
pub reason: String,
}
/// `true` for Arrow types that can be a scalar graph property.
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(_, _)
)
}
/// Why a non-scalar column can't be a graph property (for the `unmapped` note).
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)"
)
}
/// The last dotted segment of a qualified table name (`lake.social.users` → `users`).
fn last_segment(table: &str) -> String {
table.rsplit('.').next().unwrap_or(table).to_string()
}
/// A crude English singularizer good enough for a *candidate* label
/// (`users` → `user`, `companies` → `company`). A human corrects it.
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()
}
}
/// PascalCase a snake/plain identifier for a node label.
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()
}
/// The candidate node label for a table (`lake.social.users` → `User`).
fn label_for_table(table: &str) -> String {
pascal_case(&singularize(&last_segment(table)))
}
/// The candidate destination label for an FK column (`company_id` → `Company`).
fn label_for_fk(fk: &str) -> String {
pascal_case(fk.trim_end_matches("_id"))
}
impl TableMapping {
/// **The columns this mapping actually touches** — the projection a source
/// may prune to.
///
/// Every declarative step names its inputs exactly: a `vertex` reads its
/// `id_column` plus its `properties`; an `edge`/`fk_edges` reads its two
/// endpoint columns plus its `properties`. A mapping that builds a `Person`
/// node out of `(id, name)` therefore needs 2 of the table's columns — and
/// until the projection was wired, the Iceberg source read **all of them**
/// on every snapshot change.
///
/// Returns `None` when the projection would be UNSOUND:
/// * a `custom` step is an opaque registry transform that may read any
/// column, so pruning could starve it;
/// * an explicit [`columns`](TableMapping::columns) override is absent AND
/// the derived set is empty (nothing to prune to).
///
/// `None` means "read every column", i.e. the previous behaviour — the
/// projection is a strict optimisation that can only ever narrow a read it
/// is provably safe to narrow.
pub fn projected_columns(&self) -> Option<Vec<String>> {
// An explicit override wins — the operator has taken responsibility.
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());
}
// Opaque: bail out of projecting entirely. A SQL step could name
// its columns, but reading them out of an expression is the
// engine's job, not this one's.
PipelineStep::Custom { .. } | PipelineStep::Sql { .. } => return None,
}
}
let cols = dedup_preserving_order(cols);
if cols.is_empty() { None } else { Some(cols) }
}
/// Resolve the pipeline into executable transforms.
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()),
// A foreign key on an ENTITY row is one edge per source node
// — replace — but only when the row IS its source node: the
// table's own vertex step has the edge source's label and id
// column. A junction table (no such vertex; many rows per
// source) keeps every edge.
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()
}
}
/// De-duplicate while preserving first-seen order (the projection must be stable
/// so a source's column list — and any bench measuring it — is deterministic).
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()
}
// ═══════════════════════════════════════════════════════════════════════════
// The RESIDUAL row filter — the same predicate, applied in memory
// ═══════════════════════════════════════════════════════════════════════════
/// Apply a [`RowFilter`] to `batches` in memory, keeping only matching rows.
///
/// **Why this exists.** A `filter` on a [`TableMapping`] is a row PREDICATE, not a
/// hint: nothing downstream re-applies it (the [`Listener`](crate::Listener) hands
/// it to the source and never drops a row itself), so a source that cannot push it
/// down must still honour it or the operator's exclusions land in the graph anyway.
///
/// The Iceberg source pushes the predicate into the scan planner for
/// [`SyncMode::Full`] — file/row-group pruning, which is strictly faster. For every
/// other route (the manifest-level delta scan of `Incremental`/`Cdc`, and any
/// source with no pruning surface at all) this is the fallback that keeps the
/// MEANING identical: same predicate, same rows, only a different cost.
///
/// Pure and engine-free (arrow compute only), so it lives beside the spec type
/// rather than inside the optional `skade` source.
///
/// A filter whose column is absent from the batch schema is an error, not a
/// silently-empty or silently-full result — an operator who names a column that
/// does not exist has a broken mapping and must be told.
pub fn apply_row_filter(
batches: &[arrow_array::RecordBatch],
filter: &RowFilter,
) -> Result<Vec<arrow_array::RecordBatch>> {
apply_row_filter_with(batches, filter, FilterNulls::Drop)
}
/// **What to do with a row the predicate cannot judge** — one whose filter
/// column is NULL.
///
/// The two answers are both right, for different data, and getting them the
/// wrong way round costs correctness rather than speed:
///
/// * [`FilterNulls::Drop`] — a NULL does not satisfy `country = 'SE'`, so the
/// row goes. This is what the Iceberg scan planner does, so it is what an
/// ordinary (`Full` / `Incremental`) read must do, or the same mapping would
/// select different rows depending on the mode.
/// * [`FilterNulls::Keep`] — the NULL means *unknown*, not *absent*. A CDC
/// `DELETE` row is widened to the table schema with **nulls in every
/// non-identity column**, so a `Drop` predicate over a changelog batch
/// silently removes deletions — the one CDC failure that must never happen
/// quietly. Changelog reads therefore keep what they cannot judge.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FilterNulls {
/// A NULL filter column fails the predicate — the row is dropped.
/// Scan-planner semantics; the default.
Drop,
/// A NULL filter column is *unjudgeable* — the row is kept. Changelog
/// semantics (a widened `DELETE` row's columns are all null).
Keep,
}
/// [`apply_row_filter`] with an explicit [`FilterNulls`] policy.
///
/// One predicate implementation, two null policies — rather than a second copy
/// of the kernel lowering living in the CDC module (L5).
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()
}
/// One batch through the residual predicate.
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 {
// An empty `IN ()` matches nothing — the same answer the scan planner
// gives when the predicate prunes every file.
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 {
// Inclusive range; either bound may be absent (`validate` already
// guaranteed at least one is present).
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")
};
// `filter_record_batch` treats a NULL mask entry as FALSE, i.e. it drops the
// row. That is right for `Drop` and catastrophic for `Keep`, so `Keep`
// resolves the nulls to TRUE first instead of relying on that default.
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)
}
/// Lower one [`FilterValue`] to a single-element arrow `Datum` typed to MATCH the
/// column it is compared against. Comparing across types is refused rather than
/// coerced: a silent coercion is how a predicate starts selecting the wrong rows.
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])),
// An integer literal against a float column is unambiguous and lossless
// for the ranges a mapping writes; anything else is a type error.
(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}"))
}
// ═══════════════════════════════════════════════════════════════════════════
// Tests — RON parity + schema→mapping codegen (the closed RON seam)
// ═══════════════════════════════════════════════════════════════════════════
// ── mapping formats ─────────────────────────────────────────────────────────
/// A **mapping format**: whatever it is written in, it compiles to the one
/// [`MappingSpec`] the listener runs — so the changelog path, the cursors and
/// the sinks are the same whichever format a mapping came from.
///
/// | format | file | impl |
/// |---|---|---|
/// | bifröst's own: per table, a pipeline of steps | `.yml` / `.yaml` / `.ron` | [`MappingSpec`] |
/// | PuppyGraph's `schema.json`: graph-first, `node` + `edge` lists | `.json` | [`PuppyGraphMapping`] |
/// | SQL: one `SELECT` per step, reserved columns say what a row becomes | `.sql` | [`SqlMapping`] (feature `sql`) |
pub trait GraphMapping {
fn compile(&self) -> Result<MappingSpec>;
}
impl GraphMapping for MappingSpec {
fn compile(&self) -> Result<MappingSpec> {
Ok(self.clone())
}
}
/// A PuppyGraph `schema.json`; it names no sync mode, so `mode` is every table's.
/// See [`crate::puppygraph`].
#[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)
}
}
/// A `.sql` mapping: `;`-separated `SELECT`s, each a step of the table its
/// `FROM` names, every table in `mode`. See [`crate::sql`].
#[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()))
}
}
/// Load a mapping file in whichever format its extension names; `mode` is the
/// sync mode of the formats that carry none (PuppyGraph, SQL). bifröst's own
/// format names a mode per table and ignores it.
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};
/// A hand-authored YAML mapping still parses, then round-trips through RON
/// losslessly — YAML and RON are two faces of the SAME `MappingSpec`.
#[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");
// Render to RON and parse it straight back — no field lost.
let ron = from_yaml.to_ron().expect("to_ron");
let from_ron = MappingSpec::from_ron(&ron).expect("from_ron parses our own RON");
// Structural equality across the whole spec (MappingSpec isn't Eq, so
// compare the RON renders — a stable, total serialization).
assert_eq!(
from_ron.to_ron().unwrap(),
from_yaml.to_ron().unwrap(),
"YAML→RON→spec equals the YAML spec"
);
// And a spot-check on the parsed content.
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);
}
/// A candidate mapping is derived from an Arrow schema: one vertex keyed on
/// the guessed id, scalar columns as properties, an FK column as an edge —
/// and it round-trips through RON.
#[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), // FK
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");
// Vertex: label singularized+PascalCased from the table, id guessed as
// `user_id`, properties = the non-id non-FK scalar columns.
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"]);
// FK edge from `company_id` → Company.
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");
// The guesses checklist names the id + the edge.
assert!(cand.guesses.iter().any(|g| g.contains("id column")));
assert!(cand.guesses.iter().any(|g| g.contains("HAS_COMPANY")));
// The candidate emits as RON and re-parses (it's a real MappingSpec).
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);
}
/// **Honest-by-construction: a nested column is SURFACED, never dropped.** A
/// struct column can't be a scalar graph property, so it lands in `unmapped`
/// with its type + reason — not silently omitted from the mapping.
#[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), // nested!
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:?}"
);
// Every input column is accounted for: id + property + the two unmapped.
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"
);
// The reason is human-actionable.
assert!(cand.unmapped.iter().any(|u| u.reason.contains("struct")));
}
/// Whether the compiled edge step of `yaml`'s first table replaces.
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");
}
}