rbt_testing/
lib.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 { column: String },
12    /// Single Utf8 column uniqueness (legacy; prefer UniqueKey for multi-type).
13    Unique { column: String },
14    /// Composite uniqueness across one or more columns (global over all batches).
15    UniqueKey { columns: Vec<String> },
16    AcceptedValues { column: String, values: Vec<String> },
17}
18
19/// Validation result summary for a tested model.
20#[derive(Debug, Clone)]
21pub struct ValidationResult {
22    pub total_rows: usize,
23    pub passed_assertions: usize,
24    pub failed_assertions: usize,
25    pub errors: Vec<String>,
26}
27
28/// Streaming data validator for Apache Arrow record batches.
29pub struct RecordBatchValidator;
30
31impl RecordBatchValidator {
32    /// Validates that a specified column contains no NULL entries in the batch.
33    pub fn assert_not_null(batch: &RecordBatch, column: &str) -> Result<()> {
34        let schema = batch.schema();
35        let column_idx = schema
36            .index_of(column)
37            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
38
39        let array = batch.column(column_idx);
40        let null_count = array.null_count();
41        if null_count > 0 {
42            return Err(anyhow!(
43                "Assertion failed: Column '{}' has {} null values",
44                column,
45                null_count
46            ));
47        }
48
49        Ok(())
50    }
51
52    /// Validates that a string/Utf8 column contains strictly unique entries in the batch.
53    pub fn assert_unique(batch: &RecordBatch, column: &str) -> Result<()> {
54        let schema = batch.schema();
55        let column_idx = schema
56            .index_of(column)
57            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
58
59        let array = batch.column(column_idx);
60        let string_array = array
61            .as_any()
62            .downcast_ref::<StringArray>()
63            .ok_or_else(|| anyhow!("Column '{}' must be Utf8 type for unique validation", column))?;
64
65        let mut seen = HashSet::new();
66        for i in 0..string_array.len() {
67            if string_array.is_valid(i) {
68                let val = string_array.value(i);
69                if !seen.insert(val) {
70                    return Err(anyhow!(
71                        "Assertion failed: Duplicate value '{}' found in column '{}'",
72                        val,
73                        column
74                    ));
75                }
76            }
77        }
78
79        Ok(())
80    }
81
82    /// Global composite uniqueness across all batches (stringified cell values).
83    pub fn assert_unique_key(batches: &[RecordBatch], columns: &[String]) -> Result<()> {
84        if columns.is_empty() {
85            return Err(anyhow!("unique_key assertion requires at least one column"));
86        }
87        if batches.is_empty() {
88            return Ok(());
89        }
90        for col in columns {
91            if batches[0].schema().index_of(col).is_err() {
92                return Err(anyhow!(
93                    "Column '{}' not found in RecordBatch schema for unique_key",
94                    col
95                ));
96            }
97        }
98
99        let mut seen: HashSet<String> = HashSet::new();
100        for batch in batches {
101            let idxs: Vec<usize> = columns
102                .iter()
103                .map(|c| batch.schema().index_of(c).unwrap())
104                .collect();
105            for row in 0..batch.num_rows() {
106                let mut key = String::new();
107                for (i, &col_idx) in idxs.iter().enumerate() {
108                    if i > 0 {
109                        key.push('\u{1f}');
110                    }
111                    key.push_str(&array_value_to_key(batch.column(col_idx), row));
112                }
113                if !seen.insert(key.clone()) {
114                    return Err(anyhow!(
115                        "Assertion failed: Duplicate composite key {:?} on columns {:?}",
116                        key.split('\u{1f}').collect::<Vec<_>>(),
117                        columns
118                    ));
119                }
120            }
121        }
122        Ok(())
123    }
124
125    /// Validates that all non-null values in a Utf8 column belong to the `accepted_values` set.
126    pub fn assert_accepted_values(
127        batch: &RecordBatch,
128        column: &str,
129        accepted_values: &[&str],
130    ) -> Result<()> {
131        let schema = batch.schema();
132        let column_idx = schema
133            .index_of(column)
134            .map_err(|_| anyhow!("Column '{}' not found in RecordBatch schema", column))?;
135
136        let array = batch.column(column_idx);
137        let string_array = array.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
138            anyhow!(
139                "Column '{}' must be Utf8 type for accepted_values validation",
140                column
141            )
142        })?;
143
144        let valid_set: HashSet<&str> = accepted_values.iter().copied().collect();
145        for i in 0..string_array.len() {
146            if string_array.is_valid(i) {
147                let val = string_array.value(i);
148                if !valid_set.contains(val) {
149                    return Err(anyhow!(
150                        "Assertion failed: Value '{}' in column '{}' is not in accepted list {:?}",
151                        val,
152                        column,
153                        accepted_values
154                    ));
155                }
156            }
157        }
158
159        Ok(())
160    }
161
162    /// Runs a suite of assertions over an incoming array of RecordBatches.
163    ///
164    /// `Unique` / per-batch not_null still walk batches; `UniqueKey` is global.
165    pub fn validate_batches(batches: &[RecordBatch], assertions: &[Assertion]) -> ValidationResult {
166        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
167        let mut passed = 0;
168        let mut failed = 0;
169        let mut errors = Vec::new();
170
171        for assertion in assertions {
172            let res = match assertion {
173                Assertion::NotNull { column } => {
174                    let mut ok = Ok(());
175                    for batch in batches {
176                        if let Err(e) = Self::assert_not_null(batch, column) {
177                            ok = Err(e);
178                            break;
179                        }
180                    }
181                    ok
182                }
183                Assertion::Unique { column } => {
184                    // Global unique for single column via UniqueKey path
185                    Self::assert_unique_key(batches, &[column.clone()])
186                }
187                Assertion::UniqueKey { columns } => Self::assert_unique_key(batches, columns),
188                Assertion::AcceptedValues { column, values } => {
189                    let str_refs: Vec<&str> = values.iter().map(|s| s.as_str()).collect();
190                    let mut ok = Ok(());
191                    for batch in batches {
192                        if let Err(e) = Self::assert_accepted_values(batch, column, &str_refs) {
193                            ok = Err(e);
194                            break;
195                        }
196                    }
197                    ok
198                }
199            };
200
201            match res {
202                Ok(()) => passed += 1,
203                Err(e) => {
204                    errors.push(e.to_string());
205                    failed += 1;
206                }
207            }
208        }
209
210        ValidationResult {
211            total_rows,
212            passed_assertions: passed,
213            failed_assertions: failed,
214            errors,
215        }
216    }
217}
218
219fn array_value_to_key(array: &ArrayRef, row: usize) -> String {
220    if array.is_null(row) {
221        return "\u{0}".to_string();
222    }
223    // Prefer Utf8; otherwise debug-print via downcast helpers is heavy — use format of scalar
224    if let Some(s) = array.as_any().downcast_ref::<StringArray>() {
225        return s.value(row).to_string();
226    }
227    // Fallback: use Arrow's display via array.to_data() is awkward; stringify via JSON-like
228    use arrow::array::AsArray;
229    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int64Array>() {
230        return a.value(row).to_string();
231    }
232    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Float64Array>() {
233        return a.value(row).to_string();
234    }
235    if let Some(a) = array.as_any().downcast_ref::<arrow::array::Int32Array>() {
236        return a.value(row).to_string();
237    }
238    if let Some(a) = array.as_any().downcast_ref::<arrow::array::UInt64Array>() {
239        return a.value(row).to_string();
240    }
241    if let Some(a) = array.as_any().downcast_ref::<arrow::array::BooleanArray>() {
242        return a.value(row).to_string();
243    }
244    // Timestamp / other: use formatted debug of the array slice is overkill — use to_string on array
245    let _ = AsArray::as_string_opt::<i32>(array.as_ref());
246    format!("row{row}")
247}
248
249/// Build assertions from frontmatter-style test declarations.
250pub fn assertions_from_model_tests(
251    not_null: Option<&[String]>,
252    unique: Option<&[String]>,
253    accepted_values: Option<&std::collections::HashMap<String, Vec<String>>>,
254) -> Vec<Assertion> {
255    let mut out = Vec::new();
256    if let Some(cols) = not_null {
257        for c in cols {
258            out.push(Assertion::NotNull { column: c.clone() });
259        }
260    }
261    if let Some(cols) = unique {
262        if cols.len() == 1 {
263            out.push(Assertion::UniqueKey {
264                columns: cols.to_vec(),
265            });
266        } else if cols.len() > 1 {
267            out.push(Assertion::UniqueKey {
268                columns: cols.to_vec(),
269            });
270        }
271    }
272    if let Some(map) = accepted_values {
273        for (col, vals) in map {
274            out.push(Assertion::AcceptedValues {
275                column: col.clone(),
276                values: vals.clone(),
277            });
278        }
279    }
280    out
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286    use arrow::array::{Int64Array, StringArray};
287    use arrow::datatypes::{DataType, Field, Schema};
288    use std::sync::Arc;
289
290    #[test]
291    fn test_record_batch_assertions() -> Result<()> {
292        let schema = Arc::new(Schema::new(vec![
293            Field::new("id", DataType::Int64, false),
294            Field::new("status", DataType::Utf8, false),
295        ]));
296
297        let batch = RecordBatch::try_new(
298            schema,
299            vec![
300                Arc::new(Int64Array::from(vec![1, 2, 3])),
301                Arc::new(StringArray::from(vec!["active", "pending", "active"])),
302            ],
303        )?;
304
305        RecordBatchValidator::assert_not_null(&batch, "id")?;
306        RecordBatchValidator::assert_accepted_values(
307            &batch,
308            "status",
309            &["active", "pending", "completed"],
310        )?;
311
312        let assertions = vec![
313            Assertion::NotNull {
314                column: "id".to_string(),
315            },
316            Assertion::AcceptedValues {
317                column: "status".to_string(),
318                values: vec!["active".to_string(), "pending".to_string()],
319            },
320            Assertion::UniqueKey {
321                columns: vec!["id".to_string()],
322            },
323        ];
324
325        let result = RecordBatchValidator::validate_batches(&[batch], &assertions);
326        assert_eq!(result.passed_assertions, 3);
327        assert_eq!(result.failed_assertions, 0);
328
329        Ok(())
330    }
331
332    #[test]
333    fn test_composite_unique_global() -> Result<()> {
334        let schema = Arc::new(Schema::new(vec![
335            Field::new("symbol", DataType::Utf8, false),
336            Field::new("ts", DataType::Int64, false),
337        ]));
338        let b1 = RecordBatch::try_new(
339            schema.clone(),
340            vec![
341                Arc::new(StringArray::from(vec!["A", "B"])),
342                Arc::new(Int64Array::from(vec![1, 1])),
343            ],
344        )?;
345        let b2 = RecordBatch::try_new(
346            schema,
347            vec![
348                Arc::new(StringArray::from(vec!["A"])),
349                Arc::new(Int64Array::from(vec![1])),
350            ],
351        )?;
352        let res = RecordBatchValidator::validate_batches(
353            &[b1, b2],
354            &[Assertion::UniqueKey {
355                columns: vec!["symbol".into(), "ts".into()],
356            }],
357        );
358        assert_eq!(res.failed_assertions, 1);
359        Ok(())
360    }
361}
rbt_testing/lib.rs

rbt_testing/
lib.rs
