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 Utf8 column belong to the `accepted_values` set.
140    pub fn assert_accepted_values(
141        batch: &RecordBatch,
142        column: &str,
143        accepted_values: &[&str],
144    ) -> Result<()> {
145        let schema = batch.schema();
146        let column_idx = schema
147            .index_of(column)
148            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
149
150        let array = batch.column(column_idx);
151        let string_array = array
152            .as_any()
153            .downcast_ref::<StringArray>()
154            .ok_or_else(|| {
155                anyhow!(
156                    "Column '{}' must be Utf8 type for accepted_values validation",
157                    column
158                )
159            })?;
160
161        let valid_set: HashSet<&str> = accepted_values.iter().copied().collect();
162        for i in 0..string_array.len() {
163            if string_array.is_valid(i) {
164                let val = string_array.value(i);
165                if !valid_set.contains(val) {
166                    return Err(anyhow!(
167                        "Assertion failed: Value '{}' in column '{}' is not in accepted list {:?}",
168                        val,
169                        column,
170                        accepted_values
171                    ));
172                }
173            }
174        }
175
176        Ok(())
177    }
178
179    /// Runs a suite of assertions over an incoming array of RecordBatches.
180    ///
181    /// `Unique` / per-batch not_null still walk batches; `UniqueKey` is global.
182    pub fn validate_batches(batches: &[RecordBatch], assertions: &[Assertion]) -> ValidationResult {
183        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
184        let mut passed = 0;
185        let mut failed = 0;
186        let mut errors = Vec::new();
187
188        for assertion in assertions {
189            let res = match assertion {
190                Assertion::NotNull { column } => {
191                    let mut ok = Ok(());
192                    for batch in batches {
193                        if let Err(e) = Self::assert_not_null(batch, column) {
194                            ok = Err(e);
195                            break;
196                        }
197                    }
198                    ok
199                }
200                Assertion::Unique { column } => {
201                    // Global unique for single column via UniqueKey path
202                    Self::assert_unique_key(batches, std::slice::from_ref(column))
203                }
204                Assertion::UniqueKey { columns } => Self::assert_unique_key(batches, columns),
205                Assertion::AcceptedValues { column, values } => {
206                    let str_refs: Vec<&str> = values.iter().map(|s| s.as_str()).collect();
207                    let mut ok = Ok(());
208                    for batch in batches {
209                        if let Err(e) = Self::assert_accepted_values(batch, column, &str_refs) {
210                            ok = Err(e);
211                            break;
212                        }
213                    }
214                    ok
215                }
216            };
217
218            match res {
219                Ok(()) => passed += 1,
220                Err(e) => {
221                    errors.push(e.to_string());
222                    failed += 1;
223                }
224            }
225        }
226
227        ValidationResult {
228            total_rows,
229            passed_assertions: passed,
230            failed_assertions: failed,
231            errors,
232        }
233    }
234}
235
236/// Streaming frontmatter assertion runner (per-batch + global unique state).
237///
238/// Drop each batch after `observe_batch` so peak RAM stays O(unique cardinality)
239/// rather than O(full result).
240#[derive(Debug, Default)]
241pub struct StreamingAssertionRunner {
242    total_rows: usize,
243    /// (column) not-null checks applied every batch
244    not_null: Vec<String>,
245    /// (column, accepted values) checks every batch
246    accepted: Vec<(String, Vec<String>)>,
247    /// Global unique key trackers
248    unique_trackers: Vec<UniqueKeyTracker>,
249    errors: Vec<String>,
250    fail_fast: bool,
251}
252
253impl StreamingAssertionRunner {
254    pub fn new(assertions: &[Assertion], fail_fast: bool) -> Self {
255        let mut not_null = Vec::new();
256        let mut accepted = Vec::new();
257        let mut unique_trackers = Vec::new();
258        for a in assertions {
259            match a {
260                Assertion::NotNull { column } => not_null.push(column.clone()),
261                Assertion::AcceptedValues { column, values } => {
262                    accepted.push((column.clone(), values.clone()))
263                }
264                Assertion::Unique { column } => {
265                    unique_trackers.push(UniqueKeyTracker::new(vec![column.clone()]))
266                }
267                Assertion::UniqueKey { columns } => {
268                    unique_trackers.push(UniqueKeyTracker::new(columns.clone()))
269                }
270            }
271        }
272        Self {
273            total_rows: 0,
274            not_null,
275            accepted,
276            unique_trackers,
277            errors: Vec::new(),
278            fail_fast,
279        }
280    }
281
282    pub fn is_empty(&self) -> bool {
283        self.not_null.is_empty() && self.accepted.is_empty() && self.unique_trackers.is_empty()
284    }
285
286    /// Observe one batch. On fail-fast, returns Err on first violation.
287    pub fn observe_batch(&mut self, batch: &RecordBatch) -> Result<()> {
288        self.total_rows += batch.num_rows();
289        for col in &self.not_null {
290            if let Err(e) = RecordBatchValidator::assert_not_null(batch, col) {
291                self.errors.push(e.to_string());
292                if self.fail_fast {
293                    return Err(anyhow!("{}", self.errors.last().unwrap()));
294                }
295            }
296        }
297        for (col, vals) in &self.accepted {
298            let refs: Vec<&str> = vals.iter().map(|s| s.as_str()).collect();
299            if let Err(e) = RecordBatchValidator::assert_accepted_values(batch, col, &refs) {
300                self.errors.push(e.to_string());
301                if self.fail_fast {
302                    return Err(anyhow!("{}", self.errors.last().unwrap()));
303                }
304            }
305        }
306        for tracker in &mut self.unique_trackers {
307            if let Err(e) = tracker.observe_batch(batch) {
308                self.errors.push(e.to_string());
309                if self.fail_fast {
310                    return Err(anyhow!("{}", self.errors.last().unwrap()));
311                }
312            }
313        }
314        Ok(())
315    }
316
317    pub fn finish(self) -> ValidationResult {
318        let assertion_count =
319            self.not_null.len() + self.accepted.len() + self.unique_trackers.len();
320        if self.errors.is_empty() {
321            ValidationResult {
322                total_rows: self.total_rows,
323                passed_assertions: assertion_count,
324                failed_assertions: 0,
325                errors: Vec::new(),
326            }
327        } else {
328            // Count distinct error messages as failed assertion signals (fail-fast often = 1).
329            let failed = self.errors.len().min(assertion_count.max(1));
330            ValidationResult {
331                total_rows: self.total_rows,
332                passed_assertions: assertion_count.saturating_sub(failed),
333                failed_assertions: failed,
334                errors: self.errors,
335            }
336        }
337    }
338}
339
340/// Global composite uniqueness tracker for streaming materialize.
341#[derive(Debug)]
342pub struct UniqueKeyTracker {
343    columns: Vec<String>,
344    seen: HashSet<String>,
345    column_idxs: Option<Vec<usize>>,
346}
347
348impl UniqueKeyTracker {
349    pub fn new(columns: Vec<String>) -> Self {
350        Self {
351            columns,
352            seen: HashSet::new(),
353            column_idxs: None,
354        }
355    }
356
357    pub fn observe_batch(&mut self, batch: &RecordBatch) -> Result<()> {
358        if self.columns.is_empty() {
359            return Err(anyhow!("unique_key assertion requires at least one column"));
360        }
361        if self.column_idxs.is_none() {
362            let mut idxs = Vec::with_capacity(self.columns.len());
363            for col in &self.columns {
364                let idx = batch.schema().index_of(col).map_err(|_| {
365                    anyhow!("Column '{col}' not found in RecordBatch schema for unique_key")
366                })?;
367                idxs.push(idx);
368            }
369            self.column_idxs = Some(idxs);
370        }
371        let idxs = self.column_idxs.as_ref().unwrap();
372        for row in 0..batch.num_rows() {
373            let mut key = String::new();
374            for (i, &col_idx) in idxs.iter().enumerate() {
375                if i > 0 {
376                    key.push('\u{1f}');
377                }
378                key.push_str(&array_value_to_key(batch.column(col_idx), row));
379            }
380            if !self.seen.insert(key.clone()) {
381                return Err(anyhow!(
382                    "Assertion failed: Duplicate composite key {:?} on columns {:?}",
383                    key.split('\u{1f}').collect::<Vec<_>>(),
384                    self.columns
385                ));
386            }
387        }
388        Ok(())
389    }
390
391    pub fn cardinality(&self) -> usize {
392        self.seen.len()
393    }
394}
395
396pub(crate) fn array_value_to_key(array: &ArrayRef, row: usize) -> String {
397    if array.is_null(row) {
398        return "\u{0}".to_string();
399    }
400    // Common primitives (fast paths)
401    if let Some(s) = array.as_any().downcast_ref::<StringArray>() {
402        return s.value(row).to_string();
403    }
404    if let Some(s) = array
405        .as_any()
406        .downcast_ref::<arrow::array::StringViewArray>()
407    {
408        return s.value(row).to_string();
409    }
410    if let Some(s) = array
411        .as_any()
412        .downcast_ref::<arrow::array::LargeStringArray>()
413    {
414        return s.value(row).to_string();
415    }
416    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int64Array>() {
417        return a.value(row).to_string();
418    }
419    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Float64Array>() {
420        return a.value(row).to_string();
421    }
422    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int32Array>() {
423        return a.value(row).to_string();
424    }
425    if let Some(a) = array.as_any().downcast_ref::<arrow::array::UInt64Array>() {
426        return a.value(row).to_string();
427    }
428    if let Some(a) = array.as_any().downcast_ref::<arrow::array::BooleanArray>() {
429        return a.value(row).to_string();
430    }
431    // Dictionary / timestamps / views / nested: Arrow display formatter
432    // (Parquet re-read often yields Utf8View or dictionary-encoded strings.)
433    use arrow::util::display::{ArrayFormatter, FormatOptions};
434    let opts = FormatOptions::default().with_display_error(true);
435    if let Ok(fmt) = ArrayFormatter::try_new(array.as_ref(), &opts) {
436        return fmt.value(row).to_string();
437    }
438    format!("row{row}")
439}
440
441/// Build assertions from frontmatter-style test declarations.
442pub fn assertions_from_model_tests(
443    not_null: Option<&[String]>,
444    unique: Option<&[String]>,
445    accepted_values: Option<&std::collections::HashMap<String, Vec<String>>>,
446) -> Vec<Assertion> {
447    let mut out = Vec::new();
448    if let Some(cols) = not_null {
449        for c in cols {
450            out.push(Assertion::NotNull { column: c.clone() });
451        }
452    }
453    if let Some(cols) = unique {
454        if !cols.is_empty() {
455            out.push(Assertion::UniqueKey {
456                columns: cols.to_vec(),
457            });
458        }
459    }
460    if let Some(map) = accepted_values {
461        for (col, vals) in map {
462            out.push(Assertion::AcceptedValues {
463                column: col.clone(),
464                values: vals.clone(),
465            });
466        }
467    }
468    out
469}
470
471#[cfg(test)]
472mod tests {
473    use super::*;
474    use arrow::array::{Int64Array, StringArray};
475    use arrow::datatypes::{DataType, Field, Schema};
476    use std::sync::Arc;
477
478    #[test]
479    fn streaming_unique_tracker_cross_batch() -> Result<()> {
480        let schema = Arc::new(Schema::new(vec![
481            Field::new("id", DataType::Int64, false),
482            Field::new("k", DataType::Utf8, false),
483        ]));
484        let b1 = RecordBatch::try_new(
485            schema.clone(),
486            vec![
487                Arc::new(Int64Array::from(vec![1, 2])),
488                Arc::new(StringArray::from(vec!["a", "b"])),
489            ],
490        )?;
491        let b2 = RecordBatch::try_new(
492            schema,
493            vec![
494                Arc::new(Int64Array::from(vec![3, 2])),
495                Arc::new(StringArray::from(vec!["c", "b"])),
496            ],
497        )?;
498        let mut runner = StreamingAssertionRunner::new(
499            &[Assertion::UniqueKey {
500                columns: vec!["id".into()],
501            }],
502            true,
503        );
504        runner.observe_batch(&b1)?;
505        let err = runner.observe_batch(&b2).unwrap_err().to_string();
506        assert!(err.contains("Duplicate"), "got {err}");
507        Ok(())
508    }
509
510    #[test]
511    fn test_record_batch_assertions() -> Result<()> {
512        let schema = Arc::new(Schema::new(vec![
513            Field::new("id", DataType::Int64, false),
514            Field::new("status", DataType::Utf8, false),
515        ]));
516
517        let batch = RecordBatch::try_new(
518            schema,
519            vec![
520                Arc::new(Int64Array::from(vec![1, 2, 3])),
521                Arc::new(StringArray::from(vec!["active", "pending", "active"])),
522            ],
523        )?;
524
525        RecordBatchValidator::assert_not_null(&batch, "id")?;
526        RecordBatchValidator::assert_accepted_values(
527            &batch,
528            "status",
529            &["active", "pending", "completed"],
530        )?;
531
532        let assertions = vec![
533            Assertion::NotNull {
534                column: "id".to_string(),
535            },
536            Assertion::AcceptedValues {
537                column: "status".to_string(),
538                values: vec!["active".to_string(), "pending".to_string()],
539            },
540            Assertion::UniqueKey {
541                columns: vec!["id".to_string()],
542            },
543        ];
544
545        let result = RecordBatchValidator::validate_batches(&[batch], &assertions);
546        assert_eq!(result.passed_assertions, 3);
547        assert_eq!(result.failed_assertions, 0);
548
549        Ok(())
550    }
551
552    #[test]
553    fn test_composite_unique_global() -> Result<()> {
554        let schema = Arc::new(Schema::new(vec![
555            Field::new("symbol", DataType::Utf8, false),
556            Field::new("ts", DataType::Int64, false),
557        ]));
558        let b1 = RecordBatch::try_new(
559            schema.clone(),
560            vec![
561                Arc::new(StringArray::from(vec!["A", "B"])),
562                Arc::new(Int64Array::from(vec![1, 1])),
563            ],
564        )?;
565        let b2 = RecordBatch::try_new(
566            schema,
567            vec![
568                Arc::new(StringArray::from(vec!["A"])),
569                Arc::new(Int64Array::from(vec![1])),
570            ],
571        )?;
572        let res = RecordBatchValidator::validate_batches(
573            &[b1, b2],
574            &[Assertion::UniqueKey {
575                columns: vec!["symbol".into(), "ts".into()],
576            }],
577        );
578        assert_eq!(res.failed_assertions, 1);
579        Ok(())
580    }
581}
rbt/testing/mod.rs

rbt/testing/
mod.rs
