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
236fn array_value_to_key(array: &ArrayRef, row: usize) -> String {
237    if array.is_null(row) {
238        return "\u{0}".to_string();
239    }
240    // Common primitives (fast paths)
241    if let Some(s) = array.as_any().downcast_ref::<StringArray>() {
242        return s.value(row).to_string();
243    }
244    if let Some(s) = array
245        .as_any()
246        .downcast_ref::<arrow::array::StringViewArray>()
247    {
248        return s.value(row).to_string();
249    }
250    if let Some(s) = array
251        .as_any()
252        .downcast_ref::<arrow::array::LargeStringArray>()
253    {
254        return s.value(row).to_string();
255    }
256    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int64Array>() {
257        return a.value(row).to_string();
258    }
259    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Float64Array>() {
260        return a.value(row).to_string();
261    }
262    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int32Array>() {
263        return a.value(row).to_string();
264    }
265    if let Some(a) = array.as_any().downcast_ref::<arrow::array::UInt64Array>() {
266        return a.value(row).to_string();
267    }
268    if let Some(a) = array.as_any().downcast_ref::<arrow::array::BooleanArray>() {
269        return a.value(row).to_string();
270    }
271    // Dictionary / timestamps / views / nested: Arrow display formatter
272    // (Parquet re-read often yields Utf8View or dictionary-encoded strings.)
273    use arrow::util::display::{ArrayFormatter, FormatOptions};
274    let opts = FormatOptions::default().with_display_error(true);
275    if let Ok(fmt) = ArrayFormatter::try_new(array.as_ref(), &opts) {
276        return fmt.value(row).to_string();
277    }
278    format!("row{row}")
279}
280
281/// Build assertions from frontmatter-style test declarations.
282pub fn assertions_from_model_tests(
283    not_null: Option<&[String]>,
284    unique: Option<&[String]>,
285    accepted_values: Option<&std::collections::HashMap<String, Vec<String>>>,
286) -> Vec<Assertion> {
287    let mut out = Vec::new();
288    if let Some(cols) = not_null {
289        for c in cols {
290            out.push(Assertion::NotNull { column: c.clone() });
291        }
292    }
293    if let Some(cols) = unique {
294        if !cols.is_empty() {
295            out.push(Assertion::UniqueKey {
296                columns: cols.to_vec(),
297            });
298        }
299    }
300    if let Some(map) = accepted_values {
301        for (col, vals) in map {
302            out.push(Assertion::AcceptedValues {
303                column: col.clone(),
304                values: vals.clone(),
305            });
306        }
307    }
308    out
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314    use arrow::array::{Int64Array, StringArray};
315    use arrow::datatypes::{DataType, Field, Schema};
316    use std::sync::Arc;
317
318    #[test]
319    fn test_record_batch_assertions() -> Result<()> {
320        let schema = Arc::new(Schema::new(vec![
321            Field::new("id", DataType::Int64, false),
322            Field::new("status", DataType::Utf8, false),
323        ]));
324
325        let batch = RecordBatch::try_new(
326            schema,
327            vec![
328                Arc::new(Int64Array::from(vec![1, 2, 3])),
329                Arc::new(StringArray::from(vec!["active", "pending", "active"])),
330            ],
331        )?;
332
333        RecordBatchValidator::assert_not_null(&batch, "id")?;
334        RecordBatchValidator::assert_accepted_values(
335            &batch,
336            "status",
337            &["active", "pending", "completed"],
338        )?;
339
340        let assertions = vec![
341            Assertion::NotNull {
342                column: "id".to_string(),
343            },
344            Assertion::AcceptedValues {
345                column: "status".to_string(),
346                values: vec!["active".to_string(), "pending".to_string()],
347            },
348            Assertion::UniqueKey {
349                columns: vec!["id".to_string()],
350            },
351        ];
352
353        let result = RecordBatchValidator::validate_batches(&[batch], &assertions);
354        assert_eq!(result.passed_assertions, 3);
355        assert_eq!(result.failed_assertions, 0);
356
357        Ok(())
358    }
359
360    #[test]
361    fn test_composite_unique_global() -> Result<()> {
362        let schema = Arc::new(Schema::new(vec![
363            Field::new("symbol", DataType::Utf8, false),
364            Field::new("ts", DataType::Int64, false),
365        ]));
366        let b1 = RecordBatch::try_new(
367            schema.clone(),
368            vec![
369                Arc::new(StringArray::from(vec!["A", "B"])),
370                Arc::new(Int64Array::from(vec![1, 1])),
371            ],
372        )?;
373        let b2 = RecordBatch::try_new(
374            schema,
375            vec![
376                Arc::new(StringArray::from(vec!["A"])),
377                Arc::new(Int64Array::from(vec![1])),
378            ],
379        )?;
380        let res = RecordBatchValidator::validate_batches(
381            &[b1, b2],
382            &[Assertion::UniqueKey {
383                columns: vec!["symbol".into(), "ts".into()],
384            }],
385        );
386        assert_eq!(res.failed_assertions, 1);
387        Ok(())
388    }
389}
rbt/testing/mod.rs

rbt/testing/
mod.rs
