datafusion_quality/rules/
schema.rs

1use crate::{SchemaRule, ValidationError};
2use datafusion::{arrow::datatypes::DataType, common::DFSchema};
3use std::sync::Arc;
4
5/// Rule that checks if a column exists in the schema
6#[derive(Debug, Clone, Default)]
7pub struct ColumnExistsRule {
8    column_name: String,
9}
10
11impl ColumnExistsRule {
12    /// Creates a new ColumnExistsRule
13    ///
14    /// # Arguments
15    ///
16    /// * `column_name` - The name of the column to check for existence
17    pub fn new(column_name: String) -> Self {
18        Self { column_name }
19    }
20}
21
22impl SchemaRule for ColumnExistsRule {
23    fn validate_schema(&self, schema: &DFSchema) -> Result<bool, ValidationError> {
24        match schema.field_with_name(None, &self.column_name) {
25            Ok(_) => Ok(true),
26            Err(_) => Err(ValidationError::ColumnNotFound {
27                column_name: self.column_name.clone(),
28            }),
29        }
30    }
31
32    fn name(&self) -> &str {
33        "column_exists"
34    }
35
36    fn description(&self) -> &str {
37        "Checks if a column exists in the schema"
38    }
39}
40
41/// Creates a rule that checks if a column exists in the schema
42///
43/// # Examples
44///
45/// ```
46/// use datafusion_quality::rules::schema::dfq_column_exists;
47/// use datafusion_quality::RuleSet;
48///
49/// // Create a rule to check if the 'age' column exists
50/// let rule = dfq_column_exists("age");
51/// let mut ruleset = RuleSet::new();
52/// ruleset.with_schema_rule(rule);
53/// ```
54pub fn dfq_column_exists(column_name: impl AsRef<str>) -> Arc<ColumnExistsRule> {
55    Arc::new(ColumnExistsRule::new(column_name.as_ref().to_string()))
56}
57
58/// Rule that checks if a column has a specific data type
59#[derive(Debug, Clone)]
60pub struct ColumnTypeRule {
61    column_name: String,
62    expected_type: DataType,
63}
64
65impl ColumnTypeRule {
66    /// Creates a new ColumnTypeRule
67    ///
68    /// # Arguments
69    ///
70    /// * `column_name` - The name of the column to check
71    /// * `expected_type` - The expected data type of the column
72    pub fn new(column_name: String, expected_type: DataType) -> Self {
73        Self {
74            column_name,
75            expected_type,
76        }
77    }
78}
79
80impl SchemaRule for ColumnTypeRule {
81    fn validate_schema(&self, schema: &DFSchema) -> Result<bool, ValidationError> {
82        match schema.field_with_name(None, &self.column_name) {
83            Ok(field) => {
84                if field.data_type() == &self.expected_type {
85                    Ok(true)
86                } else {
87                    Err(ValidationError::TypeMismatch {
88                        message: format!(
89                            "Column: {}, Expected type {:?} but got {:?}",
90                            self.column_name,
91                            self.expected_type,
92                            field.data_type()
93                        ),
94                    })
95                }
96            }
97            Err(_) => Err(ValidationError::ColumnNotFound {
98                column_name: self.column_name.clone(),
99            }),
100        }
101    }
102
103    fn name(&self) -> &str {
104        "column_type"
105    }
106
107    fn description(&self) -> &str {
108        "Checks if a column has a specific data type"
109    }
110}
111
112/// Creates a rule that checks if a column has a specific data type
113///
114/// # Examples
115///
116/// ```
117/// use datafusion_quality::rules::schema::dfq_column_type;
118/// use datafusion_quality::RuleSet;
119/// use datafusion::arrow::datatypes::DataType;
120///
121/// // Create a rule to check if the 'age' column is of type Int32
122/// let rule = dfq_column_type("age", DataType::Int32);
123/// let mut ruleset = RuleSet::new();
124/// ruleset.with_schema_rule(rule);
125/// ```
126pub fn dfq_column_type(
127    column_name: impl AsRef<str>,
128    expected_type: DataType,
129) -> Arc<ColumnTypeRule> {
130    Arc::new(ColumnTypeRule::new(
131        column_name.as_ref().to_string(),
132        expected_type,
133    ))
134}
135
136/// Rule that checks if a column is nullable
137#[derive(Debug, Clone)]
138pub struct ColumnNullableRule {
139    column_name: String,
140    expected_nullable: bool,
141}
142
143impl ColumnNullableRule {
144    /// Creates a new ColumnNullableRule
145    ///
146    /// # Arguments
147    ///
148    /// * `column_name` - The name of the column to check
149    /// * `expected_nullable` - Whether the column should be nullable
150    pub fn new(column_name: String, expected_nullable: bool) -> Self {
151        Self {
152            column_name,
153            expected_nullable,
154        }
155    }
156}
157
158impl SchemaRule for ColumnNullableRule {
159    fn validate_schema(&self, schema: &DFSchema) -> Result<bool, ValidationError> {
160        match schema.field_with_name(None, &self.column_name) {
161            Ok(field) => {
162                if field.is_nullable() == self.expected_nullable {
163                    Ok(true)
164                } else {
165                    Err(ValidationError::ColumnNullabilityMismatch {
166                        column_name: self.column_name.clone(),
167                        expected: self.expected_nullable,
168                    })
169                }
170            }
171            Err(_) => Err(ValidationError::ColumnNotFound {
172                column_name: self.column_name.clone(),
173            }),
174        }
175    }
176
177    fn name(&self) -> &str {
178        "column_nullable"
179    }
180
181    fn description(&self) -> &str {
182        "Checks if a column is nullable"
183    }
184}
185
186/// Creates a rule that checks if a column is nullable
187///
188/// # Examples
189///
190/// ```
191/// use datafusion_quality::rules::schema::dfq_column_nullable;
192/// use datafusion_quality::RuleSet;
193///
194/// // Create a rule to check if the 'name' column is nullable
195/// let rule = dfq_column_nullable("name");
196/// let mut ruleset = RuleSet::new();
197/// ruleset.with_schema_rule(rule);
198/// ```
199pub fn dfq_column_nullable(column_name: impl AsRef<str>) -> Arc<ColumnNullableRule> {
200    Arc::new(ColumnNullableRule::new(
201        column_name.as_ref().to_string(),
202        true,
203    ))
204}
205
206/// Creates a rule that checks if a column is not nullable
207///
208/// # Examples
209///
210/// ```
211/// use datafusion_quality::rules::schema::dfq_column_not_nullable;
212/// use datafusion_quality::RuleSet;
213///
214/// // Create a rule to check if the 'id' column is not nullable
215/// let rule = dfq_column_not_nullable("id");
216/// let mut ruleset = RuleSet::new();
217/// ruleset.with_schema_rule(rule);
218/// ```
219pub fn dfq_column_not_nullable(column_name: impl AsRef<str>) -> Arc<ColumnNullableRule> {
220    Arc::new(ColumnNullableRule::new(
221        column_name.as_ref().to_string(),
222        false,
223    ))
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use datafusion::arrow::datatypes::{Field, Schema};
230    use datafusion::common::DFSchema;
231
232    #[test]
233    fn test_column_exists_rule() {
234        let arrow_schema = Schema::new(vec![
235            Field::new("id", DataType::Int32, false),
236            Field::new("name", DataType::Utf8, true),
237        ]);
238        let schema = DFSchema::try_from(arrow_schema).unwrap();
239
240        let rule = dfq_column_exists("id");
241        assert!(rule.validate_schema(&schema).unwrap());
242
243        let rule = dfq_column_exists("nonexistent");
244        assert!(rule.validate_schema(&schema).is_err());
245    }
246
247    #[test]
248    fn test_column_type_rule() {
249        let arrow_schema = Schema::new(vec![
250            Field::new("id", DataType::Int32, false),
251            Field::new("name", DataType::Utf8, true),
252        ]);
253        let schema = DFSchema::try_from(arrow_schema).unwrap();
254
255        let rule = dfq_column_type("id", DataType::Int32);
256        assert!(rule.validate_schema(&schema).unwrap());
257
258        let rule = dfq_column_type("id", DataType::Utf8);
259        assert!(rule.validate_schema(&schema).is_err());
260
261        let rule = dfq_column_type("nonexistent", DataType::Int32);
262        assert!(rule.validate_schema(&schema).is_err());
263    }
264
265    #[test]
266    fn test_column_nullable_rule() {
267        let arrow_schema = Schema::new(vec![
268            Field::new("id", DataType::Int32, false),
269            Field::new("name", DataType::Utf8, true),
270        ]);
271        let schema = DFSchema::try_from(arrow_schema).unwrap();
272
273        let rule = dfq_column_not_nullable("id");
274        assert!(rule.validate_schema(&schema).unwrap());
275
276        let rule = dfq_column_nullable("name");
277        assert!(rule.validate_schema(&schema).unwrap());
278
279        let rule = dfq_column_nullable("id");
280        assert!(rule.validate_schema(&schema).is_err());
281
282        let rule = dfq_column_nullable("nonexistent");
283        assert!(rule.validate_schema(&schema).is_err());
284    }
285}
datafusion_quality/rules/schema.rs

datafusion_quality/rules/
schema.rs
