rbt/testing/
mod.rs

1//! `rbt::testing`: Zero-copy, streamable data quality assertion & validation engine for Arrow RecordBatches.
2
3use anyhow::{anyhow, Result};
4use arrow::array::{Array, ArrayRef, StringArray};
5use arrow::record_batch::RecordBatch;
6use std::collections::HashSet;
7
8/// Individual column assertion types for dbt-compatible model testing.
9#[derive(Debug, Clone)]
10pub enum Assertion {
11    NotNull {
12        column: String,
13    },
14    /// Single Utf8 column uniqueness (legacy; prefer UniqueKey for multi-type).
15    Unique {
16        column: String,
17    },
18    /// Composite uniqueness across one or more columns (global over all batches).
19    UniqueKey {
20        columns: Vec<String>,
21    },
22    AcceptedValues {
23        column: String,
24        values: Vec<String>,
25    },
26}
27
28/// Validation result summary for a tested model.
29#[derive(Debug, Clone)]
30pub struct ValidationResult {
31    pub total_rows: usize,
32    pub passed_assertions: usize,
33    pub failed_assertions: usize,
34    pub errors: Vec<String>,
35}
36
37/// Streaming data validator for Apache Arrow record batches.
38pub struct RecordBatchValidator;
39
40impl RecordBatchValidator {
41    /// Validates that a specified column contains no NULL entries in the batch.
42    pub fn assert_not_null(batch: &RecordBatch, column: &str) -> Result<()> {
43        let schema = batch.schema();
44        let column_idx = schema
45            .index_of(column)
46            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
47
48        let array = batch.column(column_idx);
49        let null_count = array.null_count();
50        if null_count > 0 {
51            return Err(anyhow!(
52                "Assertion failed: Column '{}' has {} null values",
53                column,
54                null_count
55            ));
56        }
57
58        Ok(())
59    }
60
61    /// Validates that a string/Utf8 column contains strictly unique entries in the batch.
62    pub fn assert_unique(batch: &RecordBatch, column: &str) -> Result<()> {
63        let schema = batch.schema();
64        let column_idx = schema
65            .index_of(column)
66            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
67
68        let array = batch.column(column_idx);
69        let string_array = array
70            .as_any()
71            .downcast_ref::<StringArray>()
72            .ok_or_else(|| {
73                anyhow!(
74                    "Column '{}' must be Utf8 type for unique validation",
75                    column
76                )
77            })?;
78
79        let mut seen = HashSet::new();
80        for i in 0..string_array.len() {
81            if string_array.is_valid(i) {
82                let val = string_array.value(i);
83                if !seen.insert(val) {
84                    return Err(anyhow!(
85                        "Assertion failed: Duplicate value '{}' found in column '{}'",
86                        val,
87                        column
88                    ));
89                }
90            }
91        }
92
93        Ok(())
94    }
95
96    /// Global composite uniqueness across all batches (stringified cell values).
97    pub fn assert_unique_key(batches: &[RecordBatch], columns: &[String]) -> Result<()> {
98        if columns.is_empty() {
99            return Err(anyhow!("unique_key assertion requires at least one column"));
100        }
101        if batches.is_empty() {
102            return Ok(());
103        }
104        for col in columns {
105            if batches[0].schema().index_of(col).is_err() {
106                return Err(anyhow!(
107                    "Column '{}' not found in RecordBatch schema for unique_key",
108                    col
109                ));
110            }
111        }
112
113        let mut seen: HashSet<String> = HashSet::new();
114        for batch in batches {
115            let idxs: Vec<usize> = columns
116                .iter()
117                .map(|c| batch.schema().index_of(c).unwrap())
118                .collect();
119            for row in 0..batch.num_rows() {
120                let mut key = String::new();
121                for (i, &col_idx) in idxs.iter().enumerate() {
122                    if i > 0 {
123                        key.push('\u{1f}');
124                    }
125                    key.push_str(&array_value_to_key(batch.column(col_idx), row));
126                }
127                if !seen.insert(key.clone()) {
128                    return Err(anyhow!(
129                        "Assertion failed: Duplicate composite key {:?} on columns {:?}",
130                        key.split('\u{1f}').collect::<Vec<_>>(),
131                        columns
132                    ));
133                }
134            }
135        }
136        Ok(())
137    }
138
139    /// Validates that all non-null values in a string-like column belong to the `accepted_values` set.
140    ///
141    /// Accepts Utf8 / LargeUtf8 / Utf8View (Parquet re-read / spill often yields Utf8View).
142    pub fn assert_accepted_values(
143        batch: &RecordBatch,
144        column: &str,
145        accepted_values: &[&str],
146    ) -> Result<()> {
147        let schema = batch.schema();
148        let column_idx = schema
149            .index_of(column)
150            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
151
152        let array = batch.column(column_idx);
153        let valid_set: HashSet<&str> = accepted_values.iter().copied().collect();
154        for i in 0..array.len() {
155            if array.is_null(i) {
156                continue;
157            }
158            let val = array_value_to_key(array, i);
159            if !valid_set.contains(val.as_str()) {
160                return Err(anyhow!(
161                    "Assertion failed: Value '{}' in column '{}' is not in accepted list {:?}",
162                    val,
163                    column,
164                    accepted_values
165                ));
166            }
167        }
168
169        Ok(())
170    }
171
172    /// Runs a suite of assertions over an incoming array of RecordBatches.
173    ///
174    /// `Unique` / per-batch not_null still walk batches; `UniqueKey` is global.
175    pub fn validate_batches(batches: &[RecordBatch], assertions: &[Assertion]) -> ValidationResult {
176        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
177        let mut passed = 0;
178        let mut failed = 0;
179        let mut errors = Vec::new();
180
181        for assertion in assertions {
182            let res = match assertion {
183                Assertion::NotNull { column } => {
184                    let mut ok = Ok(());
185                    for batch in batches {
186                        if let Err(e) = Self::assert_not_null(batch, column) {
187                            ok = Err(e);
188                            break;
189                        }
190                    }
191                    ok
192                }
193                Assertion::Unique { column } => {
194                    // Global unique for single column via UniqueKey path
195                    Self::assert_unique_key(batches, std::slice::from_ref(column))
196                }
197                Assertion::UniqueKey { columns } => Self::assert_unique_key(batches, columns),
198                Assertion::AcceptedValues { column, values } => {
199                    let str_refs: Vec<&str> = values.iter().map(|s| s.as_str()).collect();
200                    let mut ok = Ok(());
201                    for batch in batches {
202                        if let Err(e) = Self::assert_accepted_values(batch, column, &str_refs) {
203                            ok = Err(e);
204                            break;
205                        }
206                    }
207                    ok
208                }
209            };
210
211            match res {
212                Ok(()) => passed += 1,
213                Err(e) => {
214                    errors.push(e.to_string());
215                    failed += 1;
216                }
217            }
218        }
219
220        ValidationResult {
221            total_rows,
222            passed_assertions: passed,
223            failed_assertions: failed,
224            errors,
225        }
226    }
227}
228
229/// Streaming frontmatter assertion runner (per-batch + global unique state).
230///
231/// Drop each batch after `observe_batch` so peak RAM stays O(unique cardinality)
232/// rather than O(full result).
233#[derive(Debug, Default)]
234pub struct StreamingAssertionRunner {
235    total_rows: usize,
236    /// (column) not-null checks applied every batch
237    not_null: Vec<String>,
238    /// (column, accepted values) checks every batch
239    accepted: Vec<(String, Vec<String>)>,
240    /// Global unique key trackers
241    unique_trackers: Vec<UniqueKeyTracker>,
242    errors: Vec<String>,
243    fail_fast: bool,
244}
245
246impl StreamingAssertionRunner {
247    pub fn new(assertions: &[Assertion], fail_fast: bool) -> Self {
248        let mut not_null = Vec::new();
249        let mut accepted = Vec::new();
250        let mut unique_trackers = Vec::new();
251        for a in assertions {
252            match a {
253                Assertion::NotNull { column } => not_null.push(column.clone()),
254                Assertion::AcceptedValues { column, values } => {
255                    accepted.push((column.clone(), values.clone()))
256                }
257                Assertion::Unique { column } => {
258                    unique_trackers.push(UniqueKeyTracker::new(vec![column.clone()]))
259                }
260                Assertion::UniqueKey { columns } => {
261                    unique_trackers.push(UniqueKeyTracker::new(columns.clone()))
262                }
263            }
264        }
265        Self {
266            total_rows: 0,
267            not_null,
268            accepted,
269            unique_trackers,
270            errors: Vec::new(),
271            fail_fast,
272        }
273    }
274
275    pub fn is_empty(&self) -> bool {
276        self.not_null.is_empty() && self.accepted.is_empty() && self.unique_trackers.is_empty()
277    }
278
279    /// Observe one batch. On fail-fast, returns Err on first violation.
280    pub fn observe_batch(&mut self, batch: &RecordBatch) -> Result<()> {
281        self.total_rows += batch.num_rows();
282        for col in &self.not_null {
283            if let Err(e) = RecordBatchValidator::assert_not_null(batch, col) {
284                self.errors.push(e.to_string());
285                if self.fail_fast {
286                    return Err(anyhow!("{}", self.errors.last().unwrap()));
287                }
288            }
289        }
290        for (col, vals) in &self.accepted {
291            let refs: Vec<&str> = vals.iter().map(|s| s.as_str()).collect();
292            if let Err(e) = RecordBatchValidator::assert_accepted_values(batch, col, &refs) {
293                self.errors.push(e.to_string());
294                if self.fail_fast {
295                    return Err(anyhow!("{}", self.errors.last().unwrap()));
296                }
297            }
298        }
299        for tracker in &mut self.unique_trackers {
300            if let Err(e) = tracker.observe_batch(batch) {
301                self.errors.push(e.to_string());
302                if self.fail_fast {
303                    return Err(anyhow!("{}", self.errors.last().unwrap()));
304                }
305            }
306        }
307        Ok(())
308    }
309
310    pub fn finish(self) -> ValidationResult {
311        let assertion_count =
312            self.not_null.len() + self.accepted.len() + self.unique_trackers.len();
313        if self.errors.is_empty() {
314            ValidationResult {
315                total_rows: self.total_rows,
316                passed_assertions: assertion_count,
317                failed_assertions: 0,
318                errors: Vec::new(),
319            }
320        } else {
321            // Count distinct error messages as failed assertion signals (fail-fast often = 1).
322            let failed = self.errors.len().min(assertion_count.max(1));
323            ValidationResult {
324                total_rows: self.total_rows,
325                passed_assertions: assertion_count.saturating_sub(failed),
326                failed_assertions: failed,
327                errors: self.errors,
328            }
329        }
330    }
331}
332
333/// Global composite uniqueness tracker for streaming materialize.
334#[derive(Debug)]
335pub struct UniqueKeyTracker {
336    columns: Vec<String>,
337    seen: HashSet<String>,
338    column_idxs: Option<Vec<usize>>,
339}
340
341impl UniqueKeyTracker {
342    pub fn new(columns: Vec<String>) -> Self {
343        Self {
344            columns,
345            seen: HashSet::new(),
346            column_idxs: None,
347        }
348    }
349
350    pub fn observe_batch(&mut self, batch: &RecordBatch) -> Result<()> {
351        if self.columns.is_empty() {
352            return Err(anyhow!("unique_key assertion requires at least one column"));
353        }
354        if self.column_idxs.is_none() {
355            let mut idxs = Vec::with_capacity(self.columns.len());
356            for col in &self.columns {
357                let idx = batch.schema().index_of(col).map_err(|_| {
358                    anyhow!("Column '{col}' not found in RecordBatch schema for unique_key")
359                })?;
360                idxs.push(idx);
361            }
362            self.column_idxs = Some(idxs);
363        }
364        let idxs = self.column_idxs.as_ref().unwrap();
365        for row in 0..batch.num_rows() {
366            let mut key = String::new();
367            for (i, &col_idx) in idxs.iter().enumerate() {
368                if i > 0 {
369                    key.push('\u{1f}');
370                }
371                key.push_str(&array_value_to_key(batch.column(col_idx), row));
372            }
373            if !self.seen.insert(key.clone()) {
374                return Err(anyhow!(
375                    "Assertion failed: Duplicate composite key {:?} on columns {:?}",
376                    key.split('\u{1f}').collect::<Vec<_>>(),
377                    self.columns
378                ));
379            }
380        }
381        Ok(())
382    }
383
384    pub fn cardinality(&self) -> usize {
385        self.seen.len()
386    }
387}
388
389pub(crate) fn array_value_to_key(array: &ArrayRef, row: usize) -> String {
390    if array.is_null(row) {
391        return "\u{0}".to_string();
392    }
393    // Common primitives (fast paths)
394    if let Some(s) = array.as_any().downcast_ref::<StringArray>() {
395        return s.value(row).to_string();
396    }
397    if let Some(s) = array
398        .as_any()
399        .downcast_ref::<arrow::array::StringViewArray>()
400    {
401        return s.value(row).to_string();
402    }
403    if let Some(s) = array
404        .as_any()
405        .downcast_ref::<arrow::array::LargeStringArray>()
406    {
407        return s.value(row).to_string();
408    }
409    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int64Array>() {
410        return a.value(row).to_string();
411    }
412    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Float64Array>() {
413        return a.value(row).to_string();
414    }
415    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int32Array>() {
416        return a.value(row).to_string();
417    }
418    if let Some(a) = array.as_any().downcast_ref::<arrow::array::UInt64Array>() {
419        return a.value(row).to_string();
420    }
421    if let Some(a) = array.as_any().downcast_ref::<arrow::array::BooleanArray>() {
422        return a.value(row).to_string();
423    }
424    // Dictionary / timestamps / views / nested: Arrow display formatter
425    // (Parquet re-read often yields Utf8View or dictionary-encoded strings.)
426    use arrow::util::display::{ArrayFormatter, FormatOptions};
427    let opts = FormatOptions::default().with_display_error(true);
428    if let Ok(fmt) = ArrayFormatter::try_new(array.as_ref(), &opts) {
429        return fmt.value(row).to_string();
430    }
431    format!("row{row}")
432}
433
434/// Build assertions from frontmatter-style test declarations.
435pub fn assertions_from_model_tests(
436    not_null: Option<&[String]>,
437    unique: Option<&[String]>,
438    accepted_values: Option<&std::collections::HashMap<String, Vec<String>>>,
439) -> Vec<Assertion> {
440    let mut out = Vec::new();
441    if let Some(cols) = not_null {
442        for c in cols {
443            out.push(Assertion::NotNull { column: c.clone() });
444        }
445    }
446    if let Some(cols) = unique {
447        if !cols.is_empty() {
448            out.push(Assertion::UniqueKey {
449                columns: cols.to_vec(),
450            });
451        }
452    }
453    if let Some(map) = accepted_values {
454        for (col, vals) in map {
455            out.push(Assertion::AcceptedValues {
456                column: col.clone(),
457                values: vals.clone(),
458            });
459        }
460    }
461    out
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use arrow::array::{Int64Array, StringArray};
468    use arrow::datatypes::{DataType, Field, Schema};
469    use std::sync::Arc;
470
471    #[test]
472    fn streaming_unique_tracker_cross_batch() -> Result<()> {
473        let schema = Arc::new(Schema::new(vec![
474            Field::new("id", DataType::Int64, false),
475            Field::new("k", DataType::Utf8, false),
476        ]));
477        let b1 = RecordBatch::try_new(
478            schema.clone(),
479            vec![
480                Arc::new(Int64Array::from(vec![1, 2])),
481                Arc::new(StringArray::from(vec!["a", "b"])),
482            ],
483        )?;
484        let b2 = RecordBatch::try_new(
485            schema,
486            vec![
487                Arc::new(Int64Array::from(vec![3, 2])),
488                Arc::new(StringArray::from(vec!["c", "b"])),
489            ],
490        )?;
491        let mut runner = StreamingAssertionRunner::new(
492            &[Assertion::UniqueKey {
493                columns: vec!["id".into()],
494            }],
495            true,
496        );
497        runner.observe_batch(&b1)?;
498        let err = runner.observe_batch(&b2).unwrap_err().to_string();
499        assert!(err.contains("Duplicate"), "got {err}");
500        Ok(())
501    }
502
503    #[test]
504    fn test_record_batch_assertions() -> Result<()> {
505        let schema = Arc::new(Schema::new(vec![
506            Field::new("id", DataType::Int64, false),
507            Field::new("status", DataType::Utf8, false),
508        ]));
509
510        let batch = RecordBatch::try_new(
511            schema,
512            vec![
513                Arc::new(Int64Array::from(vec![1, 2, 3])),
514                Arc::new(StringArray::from(vec!["active", "pending", "active"])),
515            ],
516        )?;
517
518        RecordBatchValidator::assert_not_null(&batch, "id")?;
519        RecordBatchValidator::assert_accepted_values(
520            &batch,
521            "status",
522            &["active", "pending", "completed"],
523        )?;
524
525        let assertions = vec![
526            Assertion::NotNull {
527                column: "id".to_string(),
528            },
529            Assertion::AcceptedValues {
530                column: "status".to_string(),
531                values: vec!["active".to_string(), "pending".to_string()],
532            },
533            Assertion::UniqueKey {
534                columns: vec!["id".to_string()],
535            },
536        ];
537
538        let result = RecordBatchValidator::validate_batches(&[batch], &assertions);
539        assert_eq!(result.passed_assertions, 3);
540        assert_eq!(result.failed_assertions, 0);
541
542        Ok(())
543    }
544
545    #[test]
546    fn test_composite_unique_global() -> Result<()> {
547        let schema = Arc::new(Schema::new(vec![
548            Field::new("symbol", DataType::Utf8, false),
549            Field::new("ts", DataType::Int64, false),
550        ]));
551        let b1 = RecordBatch::try_new(
552            schema.clone(),
553            vec![
554                Arc::new(StringArray::from(vec!["A", "B"])),
555                Arc::new(Int64Array::from(vec![1, 1])),
556            ],
557        )?;
558        let b2 = RecordBatch::try_new(
559            schema,
560            vec![
561                Arc::new(StringArray::from(vec!["A"])),
562                Arc::new(Int64Array::from(vec![1])),
563            ],
564        )?;
565        let res = RecordBatchValidator::validate_batches(
566            &[b1, b2],
567            &[Assertion::UniqueKey {
568                columns: vec!["symbol".into(), "ts".into()],
569            }],
570        );
571        assert_eq!(res.failed_assertions, 1);
572        Ok(())
573    }
574}
rbt/testing/mod.rs

rbt/testing/
mod.rs
