fv_streams_engine/

steps.rs

//! INLINE STEPS ON THE BATCH: `select` / `rename` / `drop` / `filter` / `applyExpression` applied to
//! an Arrow `RecordBatch`, with the value dialect compiled through `fv-value-datafusion` into
//! DataFusion physical expressions (vectorised over the columns; a sub-expression the translator
//! cannot reproduce exactly runs as a per-row UDF inside the same expression, so semantics never
//! change — only speed).
//!
//! Compilation depends on the batch's schema (column types decide what is exact), so a chain
//! compiles lazily on the first batch and again whenever the input schema changes. A step that
//! cannot be compiled against the schema at all (a column of a type the dialect has no kind for)
//! makes the whole chain fall back to the row path — today's behaviour, logged once.
//!
//! Errors are isolated per row exactly as before: a batch that fails to evaluate (the dialect's
//! run-time errors, raised through the translator's error UDF) is re-applied one row at a time,
//! the survivors kept and the poison rows counted.

use std::sync::Arc;

use arrow::array::{Array, ArrayRef, BooleanArray, RecordBatch};
use arrow::compute::{concat_batches, filter_record_batch};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::common::DFSchema;
use datafusion::execution::context::ExecutionProps;
use datafusion::physical_expr::{create_physical_expr, PhysicalExpr};
use fv_plan::inline::Step;
use fv_value_datafusion::{data_type_for, Kind, Options, Translator};

/// Why a chain could not be compiled for a schema (the caller falls back to rows).
#[derive(Debug, Clone)]
pub struct NotCompilable(pub String);

/// Columns with this prefix are the engine's routing metadata (offset, key): a `select` keeps
/// them, and they are stripped before the rows reach an operator or an output.
pub use fv_streams_types::decode::META_PREFIX;

enum Compiled {
    /// Project to these input column indexes (`None` = a missing column, which lands null).
    Select {
        columns: Vec<(String, Option<usize>)>,
    },
    Rename {
        mapping: Vec<(String, String)>,
    },
    Drop {
        columns: Vec<String>,
    },
    Filter {
        expr: Arc<dyn PhysicalExpr>,
    },
    Apply {
        column: String,
        expr: Arc<dyn PhysicalExpr>,
        data_type: DataType,
    },
}

/// A chain of inline steps compiled for one input schema.
pub struct BatchSteps {
    steps: Vec<Step>,
    compiled: Option<(SchemaRef, Vec<Compiled>)>,
    /// Rows dropped as poison (the per-row isolation path) since construction.
    pub dropped: u64,
    /// Whether every expression in the chain translated exactly (no per-row UDF fallback).
    pub exact: bool,
    /// The expressions that run per row are named once, at the first compile: a vectorised
    /// chain costs a few ns per row per expression, a per-row fallback tens to hundreds — and a
    /// regex compiled per call, thousands. Nobody should learn that from a benchmark.
    reported: bool,
}

impl BatchSteps {
    pub fn new(steps: Vec<Step>) -> Self {
        BatchSteps {
            steps,
            compiled: None,
            dropped: 0,
            exact: true,
            reported: false,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.steps.is_empty()
    }

    /// Apply the chain to a batch. `Err(NotCompilable)` means this schema cannot be handled on the
    /// batch path at all (fall back to rows); poison rows are isolated and counted, never an error.
    pub fn apply(&mut self, batch: &RecordBatch) -> Result<RecordBatch, NotCompilable> {
        if self.steps.is_empty() {
            return Ok(batch.clone());
        }
        let schema = batch.schema();
        let needs_compile = match &self.compiled {
            Some((s, _)) => s.as_ref() != schema.as_ref(),
            None => true,
        };
        if needs_compile {
            let (plan, exact, per_row) = compile(&self.steps, &schema)?;
            self.exact = exact;
            if !per_row.is_empty() && !self.reported {
                self.reported = true;
                eprintln!(
                    "steps: {} expression(s) run per row, not vectorised (the translator has no exact form for them): {}",
                    per_row.len(),
                    per_row.join("; ")
                );
            }
            self.compiled = Some((Arc::clone(&schema), plan));
        }
        let plan = &self.compiled.as_ref().expect("compiled above").1;
        match apply_plan(plan, batch) {
            Ok(b) => Ok(b),
            Err(_) => {
                // per-row isolation: the batch has at least one poison row.
                let mut good: Vec<RecordBatch> = Vec::new();
                let mut out_schema: Option<SchemaRef> = None;
                for i in 0..batch.num_rows() {
                    match apply_plan(plan, &batch.slice(i, 1)) {
                        Ok(b) => {
                            out_schema.get_or_insert_with(|| b.schema());
                            good.push(b);
                        }
                        Err(_) => self.dropped += 1,
                    }
                }
                match out_schema {
                    Some(s) => concat_batches(&s, &good).map_err(|e| NotCompilable(e.to_string())),
                    // every row was poison: an empty batch of the plan's output shape.
                    None => apply_plan(plan, &batch.slice(0, 0)).map_err(|e| NotCompilable(e.to_string())),
                }
            }
        }
    }
}

fn compile(steps: &[Step], input: &SchemaRef) -> Result<(Vec<Compiled>, bool, Vec<String>), NotCompilable> {
    let mut schema = Arc::clone(input);
    let mut plan = Vec::with_capacity(steps.len());
    let mut exact = true;
    let mut per_row: Vec<String> = Vec::new();
    for step in steps {
        let c = match step {
            Step::Select { columns } => {
                let cols: Vec<(String, Option<usize>)> =
                    columns.iter().map(|c| (c.clone(), schema.index_of(c).ok())).collect();
                let fields: Vec<Field> = cols
                    .iter()
                    .map(|(name, idx)| match idx {
                        Some(i) => schema.field(*i).clone().with_name(name),
                        None => Field::new(name, DataType::Null, true),
                    })
                    .collect();
                schema = Arc::new(Schema::new(fields));
                Compiled::Select { columns: cols }
            }
            Step::Rename { mapping } => {
                let fields: Vec<Field> = schema
                    .fields()
                    .iter()
                    .map(|f| {
                        let name = mapping
                            .iter()
                            .find(|(from, _)| from == f.name())
                            .map(|(_, to)| to.as_str())
                            .unwrap_or(f.name());
                        f.as_ref().clone().with_name(name)
                    })
                    .collect();
                schema = Arc::new(Schema::new(fields));
                Compiled::Rename {
                    mapping: mapping.clone(),
                }
            }
            Step::Drop { columns } => {
                let fields: Vec<Field> = schema
                    .fields()
                    .iter()
                    .filter(|f| !columns.contains(f.name()))
                    .map(|f| f.as_ref().clone())
                    .collect();
                schema = Arc::new(Schema::new(fields));
                Compiled::Drop {
                    columns: columns.clone(),
                }
            }
            Step::Filter { expression } => {
                let ast = fv_value::compile(expression).map_err(|e| NotCompilable(format!("filter: {e}")))?;
                let t = Translator::new(Arc::clone(&schema))
                    .with_options(Options { assume_no_nan: true })
                    .translate_predicate_or_fallback(&ast)
                    .map_err(|e| NotCompilable(format!("filter '{expression}': {e}")))?;
                exact &= t.exact;
                if !t.exact {
                    per_row.push(format!("filter `{expression}`"));
                }
                Compiled::Filter {
                    expr: physical(&t.expr, &schema)?,
                }
            }
            Step::ApplyExpression { column, expression } => {
                let ast = fv_value::compile(expression).map_err(|e| NotCompilable(format!("applyExpression: {e}")))?;
                let t = Translator::new(Arc::clone(&schema))
                    .with_options(Options { assume_no_nan: true })
                    .translate_or_fallback(&ast)
                    .map_err(|e| NotCompilable(format!("applyExpression '{expression}': {e}")))?;
                exact &= t.exact;
                if !t.exact {
                    per_row.push(format!("{column} = `{expression}`"));
                }
                let data_type = data_type_for(t.kind);
                let expr = physical(&t.expr, &schema)?;
                schema = Arc::new(schema_with(&schema, column, &data_type));
                Compiled::Apply {
                    column: column.clone(),
                    expr,
                    data_type,
                }
            }
        };
        plan.push(c);
    }
    Ok((plan, exact, per_row))
}

fn physical(expr: &datafusion::logical_expr::Expr, schema: &SchemaRef) -> Result<Arc<dyn PhysicalExpr>, NotCompilable> {
    let df_schema = DFSchema::try_from(Arc::clone(schema)).map_err(|e| NotCompilable(e.to_string()))?;
    create_physical_expr(expr, &df_schema, &ExecutionProps::new()).map_err(|e| NotCompilable(e.to_string()))
}

/// The schema after `column` is set to `data_type`: overwritten in place if present, else appended.
fn schema_with(schema: &Schema, column: &str, data_type: &DataType) -> Schema {
    let mut fields: Vec<Field> = schema.fields().iter().map(|f| f.as_ref().clone()).collect();
    match fields.iter_mut().find(|f| f.name() == column) {
        Some(f) => *f = Field::new(column, data_type.clone(), true),
        None => fields.push(Field::new(column, data_type.clone(), true)),
    }
    Schema::new(fields)
}

fn apply_plan(plan: &[Compiled], batch: &RecordBatch) -> Result<RecordBatch, arrow::error::ArrowError> {
    let mut b = batch.clone();
    for step in plan {
        b = match step {
            Compiled::Select { columns } => {
                let n = b.num_rows();
                let (mut fields, mut arrays): (Vec<Field>, Vec<ArrayRef>) = columns
                    .iter()
                    .map(|(name, idx)| match idx {
                        Some(i) => (b.schema().field(*i).clone().with_name(name), Arc::clone(b.column(*i))),
                        None => (
                            Field::new(name, DataType::Null, true),
                            arrow::array::new_null_array(&DataType::Null, n),
                        ),
                    })
                    .unzip();
                // the engine's own routing columns ride along through a projection.
                for (i, f) in b.schema().fields().iter().enumerate() {
                    if f.name().starts_with(META_PREFIX) {
                        fields.push(f.as_ref().clone());
                        arrays.push(Arc::clone(b.column(i)));
                    }
                }
                RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?
            }
            Compiled::Rename { mapping } => {
                let fields: Vec<Field> = b
                    .schema()
                    .fields()
                    .iter()
                    .map(|f| {
                        let name = mapping
                            .iter()
                            .find(|(from, _)| from == f.name())
                            .map(|(_, to)| to.as_str())
                            .unwrap_or(f.name());
                        f.as_ref().clone().with_name(name)
                    })
                    .collect();
                RecordBatch::try_new(Arc::new(Schema::new(fields)), b.columns().to_vec())?
            }
            Compiled::Drop { columns } => {
                let keep: Vec<usize> = b
                    .schema()
                    .fields()
                    .iter()
                    .enumerate()
                    .filter(|(_, f)| !columns.contains(f.name()))
                    .map(|(i, _)| i)
                    .collect();
                b.project(&keep)?
            }
            Compiled::Filter { expr } => {
                let mask = expr
                    .evaluate(&b)
                    .and_then(|v| v.into_array(b.num_rows()))
                    .map_err(|e| arrow::error::ArrowError::ComputeError(e.to_string()))?;
                let mask = mask
                    .as_any()
                    .downcast_ref::<BooleanArray>()
                    .ok_or_else(|| arrow::error::ArrowError::ComputeError("filter did not yield booleans".into()))?;
                // a NULL predicate drops the row, which is the dialect's answer for a comparison with null.
                filter_record_batch(&b, mask)?
            }
            Compiled::Apply {
                column,
                expr,
                data_type,
            } => {
                let value = expr
                    .evaluate(&b)
                    .and_then(|v| v.into_array(b.num_rows()))
                    .map_err(|e| arrow::error::ArrowError::ComputeError(e.to_string()))?;
                let value = if value.data_type() == data_type {
                    value
                } else {
                    arrow::compute::cast(&value, data_type)?
                };
                let schema = Arc::new(schema_with(&b.schema(), column, data_type));
                let mut arrays: Vec<ArrayRef> = b.columns().to_vec();
                match b.schema().index_of(column) {
                    Ok(i) => arrays[i] = value,
                    Err(_) => arrays.push(value),
                }
                RecordBatch::try_new(schema, arrays)?
            }
        };
    }
    Ok(b)
}

/// The Arrow type a dialect kind lands in; re-exported so callers can size declared columns.
pub fn kind_type(kind: Kind) -> DataType {
    data_type_for(kind)
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Float64Array, StringArray};

    fn batch() -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new("type", DataType::Utf8, true),
            Field::new("price", DataType::Float64, true),
            Field::new("auction", DataType::Float64, true),
        ]));
        let ty: StringArray = vec!["bid", "auction", "bid"].into();
        let price: Float64Array = vec![Some(10.0), None, Some(0.0)].into();
        let auction: Float64Array = vec![1.0, 2.0, 3.0].into();
        RecordBatch::try_new(schema, vec![Arc::new(ty), Arc::new(price), Arc::new(auction)]).unwrap()
    }

    fn col_f64(b: &RecordBatch, name: &str) -> Vec<Option<f64>> {
        let a = b.column_by_name(name).unwrap();
        let a = a.as_any().downcast_ref::<Float64Array>().unwrap();
        (0..a.len())
            .map(|i| if a.is_null(i) { None } else { Some(a.value(i)) })
            .collect()
    }

    #[test]
    fn filter_keeps_only_true_rows_and_drops_null_comparisons() {
        let mut s = BatchSteps::new(vec![Step::Filter {
            expression: "price > 5".into(),
        }]);
        let out = s.apply(&batch()).unwrap();
        assert_eq!(out.num_rows(), 1);
        assert_eq!(col_f64(&out, "auction"), vec![Some(1.0)]);
        assert!(s.exact, "a plain comparison translates exactly");
    }

    #[test]
    fn apply_adds_or_overwrites_a_column_and_the_chain_sees_it() {
        let mut s = BatchSteps::new(vec![
            Step::ApplyExpression {
                column: "eur".into(),
                expression: "price * 2".into(),
            },
            Step::Filter {
                expression: "eur >= 20".into(),
            },
            Step::ApplyExpression {
                column: "auction".into(),
                expression: "auction + 100".into(),
            },
        ]);
        let out = s.apply(&batch()).unwrap();
        assert_eq!(out.num_rows(), 1);
        assert_eq!(col_f64(&out, "eur"), vec![Some(20.0)]);
        assert_eq!(col_f64(&out, "auction"), vec![Some(101.0)]);
        assert_eq!(out.schema().fields().len(), 4);
    }

    #[test]
    fn select_rename_drop_shape_the_schema_in_order() {
        let mut s = BatchSteps::new(vec![
            Step::Rename {
                mapping: vec![("auction".into(), "a".into())],
            },
            Step::Select {
                columns: vec!["a".into(), "missing".into(), "price".into()],
            },
            Step::Drop {
                columns: vec!["price".into()],
            },
        ]);
        let out = s.apply(&batch()).unwrap();
        let names: Vec<String> = out.schema().fields().iter().map(|f| f.name().clone()).collect();
        assert_eq!(names, vec!["a", "missing"]);
        assert_eq!(out.column(1).logical_null_count(), 3, "a missing column lands null");
        assert_eq!(out.column(1).data_type(), &DataType::Null);
        assert_eq!(col_f64(&out, "a"), vec![Some(1.0), Some(2.0), Some(3.0)]);
    }

    #[test]
    fn select_keeps_the_engines_meta_columns() {
        let b = batch();
        let off: arrow::array::Int64Array = vec![7, 8, 9].into();
        let with_meta = RecordBatch::try_new(
            Arc::new(Schema::new(
                b.schema()
                    .fields()
                    .iter()
                    .map(|f| f.as_ref().clone())
                    .chain([Field::new("__fv_offset", DataType::Int64, true)])
                    .collect::<Vec<_>>(),
            )),
            b.columns().iter().cloned().chain([Arc::new(off) as ArrayRef]).collect(),
        )
        .unwrap();
        let mut s = BatchSteps::new(vec![Step::Select {
            columns: vec!["price".into()],
        }]);
        let out = s.apply(&with_meta).unwrap();
        let names: Vec<String> = out.schema().fields().iter().map(|f| f.name().clone()).collect();
        assert_eq!(names, vec!["price", "__fv_offset"]);
    }

    #[test]
    fn a_poison_row_is_isolated_and_counted() {
        // `/ 0` is a dialect run-time error, raised per row through the error UDF.
        let mut s = BatchSteps::new(vec![Step::ApplyExpression {
            column: "r".into(),
            expression: "10 / price".into(),
        }]);
        let out = s.apply(&batch()).unwrap();
        assert_eq!(out.num_rows(), 2, "the price=0 row is dropped");
        assert_eq!(s.dropped, 1);
        assert_eq!(col_f64(&out, "r"), vec![Some(1.0), None]);
    }

    #[test]
    fn recompiles_when_the_schema_changes() {
        let mut s = BatchSteps::new(vec![Step::Filter {
            expression: "price > 5".into(),
        }]);
        assert_eq!(s.apply(&batch()).unwrap().num_rows(), 1);
        let other = batch().project(&[1]).unwrap(); // just `price`
        assert_eq!(s.apply(&other).unwrap().num_rows(), 1);
    }

    #[test]
    fn a_column_the_dialect_has_no_kind_for_runs_through_the_row_fallback() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "s",
            DataType::Struct(vec![Field::new("x", DataType::Int64, true)].into()),
            true,
        )]));
        let inner = arrow::array::StructArray::from(vec![(
            Arc::new(Field::new("x", DataType::Int64, true)),
            Arc::new(arrow::array::Int64Array::from(vec![1])) as ArrayRef,
        )]);
        let b = RecordBatch::try_new(schema, vec![Arc::new(inner)]).unwrap();
        let mut s = BatchSteps::new(vec![Step::Filter {
            expression: "s == 1".into(),
        }]);
        // the translator cannot type a struct column, so the comparison becomes a per-row UDF
        // (the dialect's own evaluator): the chain still runs, and reports itself inexact.
        let out = s.apply(&b).unwrap();
        assert_eq!(out.num_rows(), 0, "a struct never equals 1");
        assert!(!s.exact);
    }
}