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    // Prefer Utf8; otherwise debug-print via downcast helpers is heavy — use format of scalar
241    if let Some(s) = array.as_any().downcast_ref::<StringArray>() {
242        return s.value(row).to_string();
243    }
244    // Fallback: use Arrow's display via array.to_data() is awkward; stringify via JSON-like
245    use arrow::array::AsArray;
246    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int64Array>() {
247        return a.value(row).to_string();
248    }
249    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Float64Array>() {
250        return a.value(row).to_string();
251    }
252    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int32Array>() {
253        return a.value(row).to_string();
254    }
255    if let Some(a) = array.as_any().downcast_ref::<arrow::array::UInt64Array>() {
256        return a.value(row).to_string();
257    }
258    if let Some(a) = array.as_any().downcast_ref::<arrow::array::BooleanArray>() {
259        return a.value(row).to_string();
260    }
261    // Timestamp / other: use formatted debug of the array slice is overkill — use to_string on array
262    let _ = AsArray::as_string_opt::<i32>(array.as_ref());
263    format!("row{row}")
264}
265
266/// Build assertions from frontmatter-style test declarations.
267pub fn assertions_from_model_tests(
268    not_null: Option<&[String]>,
269    unique: Option<&[String]>,
270    accepted_values: Option<&std::collections::HashMap<String, Vec<String>>>,
271) -> Vec<Assertion> {
272    let mut out = Vec::new();
273    if let Some(cols) = not_null {
274        for c in cols {
275            out.push(Assertion::NotNull { column: c.clone() });
276        }
277    }
278    if let Some(cols) = unique {
279        if !cols.is_empty() {
280            out.push(Assertion::UniqueKey {
281                columns: cols.to_vec(),
282            });
283        }
284    }
285    if let Some(map) = accepted_values {
286        for (col, vals) in map {
287            out.push(Assertion::AcceptedValues {
288                column: col.clone(),
289                values: vals.clone(),
290            });
291        }
292    }
293    out
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299    use arrow::array::{Int64Array, StringArray};
300    use arrow::datatypes::{DataType, Field, Schema};
301    use std::sync::Arc;
302
303    #[test]
304    fn test_record_batch_assertions() -> Result<()> {
305        let schema = Arc::new(Schema::new(vec![
306            Field::new("id", DataType::Int64, false),
307            Field::new("status", DataType::Utf8, false),
308        ]));
309
310        let batch = RecordBatch::try_new(
311            schema,
312            vec![
313                Arc::new(Int64Array::from(vec![1, 2, 3])),
314                Arc::new(StringArray::from(vec!["active", "pending", "active"])),
315            ],
316        )?;
317
318        RecordBatchValidator::assert_not_null(&batch, "id")?;
319        RecordBatchValidator::assert_accepted_values(
320            &batch,
321            "status",
322            &["active", "pending", "completed"],
323        )?;
324
325        let assertions = vec![
326            Assertion::NotNull {
327                column: "id".to_string(),
328            },
329            Assertion::AcceptedValues {
330                column: "status".to_string(),
331                values: vec!["active".to_string(), "pending".to_string()],
332            },
333            Assertion::UniqueKey {
334                columns: vec!["id".to_string()],
335            },
336        ];
337
338        let result = RecordBatchValidator::validate_batches(&[batch], &assertions);
339        assert_eq!(result.passed_assertions, 3);
340        assert_eq!(result.failed_assertions, 0);
341
342        Ok(())
343    }
344
345    #[test]
346    fn test_composite_unique_global() -> Result<()> {
347        let schema = Arc::new(Schema::new(vec![
348            Field::new("symbol", DataType::Utf8, false),
349            Field::new("ts", DataType::Int64, false),
350        ]));
351        let b1 = RecordBatch::try_new(
352            schema.clone(),
353            vec![
354                Arc::new(StringArray::from(vec!["A", "B"])),
355                Arc::new(Int64Array::from(vec![1, 1])),
356            ],
357        )?;
358        let b2 = RecordBatch::try_new(
359            schema,
360            vec![
361                Arc::new(StringArray::from(vec!["A"])),
362                Arc::new(Int64Array::from(vec![1])),
363            ],
364        )?;
365        let res = RecordBatchValidator::validate_batches(
366            &[b1, b2],
367            &[Assertion::UniqueKey {
368                columns: vec!["symbol".into(), "ts".into()],
369            }],
370        );
371        assert_eq!(res.failed_assertions, 1);
372        Ok(())
373    }
374}
rbt/testing/mod.rs

rbt/testing/
mod.rs
