polars_core/
testing.rs

1//! Testing utilities.
2
3use crate::prelude::*;
4
5impl Series {
6    /// Check if series are equal. Note that `None == None` evaluates to `false`
7    pub fn equals(&self, other: &Series) -> bool {
8        if self.null_count() > 0 || other.null_count() > 0 {
9            false
10        } else {
11            self.equals_missing(other)
12        }
13    }
14
15    /// Check if all values in series are equal where `None == None` evaluates to `true`.
16    pub fn equals_missing(&self, other: &Series) -> bool {
17        match (self.dtype(), other.dtype()) {
18            // Two [`Datetime`](DataType::Datetime) series are *not* equal if their timezones
19            // are different, regardless if they represent the same UTC time or not.
20            #[cfg(feature = "timezones")]
21            (DataType::Datetime(_, tz_lhs), DataType::Datetime(_, tz_rhs)) => {
22                if tz_lhs != tz_rhs {
23                    return false;
24                }
25            },
26            _ => {},
27        }
28
29        // Differs from Partial::eq in that numerical dtype may be different
30        self.len() == other.len() && self.null_count() == other.null_count() && {
31            let eq = self.equal_missing(other);
32            match eq {
33                Ok(b) => b.all(),
34                Err(_) => false,
35            }
36        }
37    }
38}
39
40impl PartialEq for Series {
41    fn eq(&self, other: &Self) -> bool {
42        self.equals_missing(other)
43    }
44}
45
46impl DataFrame {
47    /// Check if [`DataFrame`]' schemas are equal.
48    pub fn schema_equal(&self, other: &DataFrame) -> PolarsResult<()> {
49        for (lhs, rhs) in self.iter().zip(other.iter()) {
50            polars_ensure!(
51                lhs.name() == rhs.name(),
52                SchemaMismatch: "column name mismatch: left-hand = '{}', right-hand = '{}'",
53                lhs.name(), rhs.name()
54            );
55            polars_ensure!(
56                lhs.dtype() == rhs.dtype(),
57                SchemaMismatch: "column datatype mismatch: left-hand = '{}', right-hand = '{}'",
58                lhs.dtype(), rhs.dtype()
59            );
60        }
61        Ok(())
62    }
63
64    /// Check if [`DataFrame`]s are equal. Note that `None == None` evaluates to `false`
65    ///
66    /// # Example
67    ///
68    /// ```rust
69    /// # use polars_core::prelude::*;
70    /// let df1: DataFrame = df!("Atomic number" => &[1, 51, 300],
71    ///                         "Element" => &[Some("Hydrogen"), Some("Antimony"), None])?;
72    /// let df2: DataFrame = df!("Atomic number" => &[1, 51, 300],
73    ///                         "Element" => &[Some("Hydrogen"), Some("Antimony"), None])?;
74    ///
75    /// assert!(!df1.equals(&df2));
76    /// # Ok::<(), PolarsError>(())
77    /// ```
78    pub fn equals(&self, other: &DataFrame) -> bool {
79        if self.shape() != other.shape() {
80            return false;
81        }
82        for (left, right) in self.get_columns().iter().zip(other.get_columns()) {
83            if left.name() != right.name() || !left.equals(right) {
84                return false;
85            }
86        }
87        true
88    }
89
90    /// Check if all values in [`DataFrame`]s are equal where `None == None` evaluates to `true`.
91    ///
92    /// # Example
93    ///
94    /// ```rust
95    /// # use polars_core::prelude::*;
96    /// let df1: DataFrame = df!("Atomic number" => &[1, 51, 300],
97    ///                         "Element" => &[Some("Hydrogen"), Some("Antimony"), None])?;
98    /// let df2: DataFrame = df!("Atomic number" => &[1, 51, 300],
99    ///                         "Element" => &[Some("Hydrogen"), Some("Antimony"), None])?;
100    ///
101    /// assert!(df1.equals_missing(&df2));
102    /// # Ok::<(), PolarsError>(())
103    /// ```
104    pub fn equals_missing(&self, other: &DataFrame) -> bool {
105        if self.shape() != other.shape() {
106            return false;
107        }
108        for (left, right) in self.get_columns().iter().zip(other.get_columns()) {
109            if left.name() != right.name() || !left.equals_missing(right) {
110                return false;
111            }
112        }
113        true
114    }
115}
116
117impl PartialEq for DataFrame {
118    fn eq(&self, other: &Self) -> bool {
119        self.shape() == other.shape()
120            && self
121                .columns
122                .iter()
123                .zip(other.columns.iter())
124                .all(|(s1, s2)| s1.equals_missing(s2))
125    }
126}
127
128/// Asserts that two expressions of type [`DataFrame`] are equal according to [`DataFrame::equals`]
129/// at runtime.
130///
131/// If the expression are not equal, the program will panic with a message that displays
132/// both dataframes.
133#[macro_export]
134macro_rules! assert_df_eq {
135    ($a:expr, $b:expr $(,)?) => {
136        let a: &$crate::frame::DataFrame = &$a;
137        let b: &$crate::frame::DataFrame = &$b;
138        assert!(a.equals(b), "expected {:?}\nto equal {:?}", a, b);
139    };
140}
141
142#[cfg(test)]
143mod test {
144    use crate::prelude::*;
145
146    #[test]
147    fn test_series_equals() {
148        let a = Series::new("a".into(), &[1_u32, 2, 3]);
149        let b = Series::new("a".into(), &[1_u32, 2, 3]);
150        assert!(a.equals(&b));
151
152        let s = Series::new("foo".into(), &[None, Some(1i64)]);
153        assert!(s.equals_missing(&s));
154    }
155
156    #[test]
157    fn test_series_dtype_not_equal() {
158        let s_i32 = Series::new("a".into(), &[1_i32, 2_i32]);
159        let s_i64 = Series::new("a".into(), &[1_i64, 2_i64]);
160        assert!(s_i32.dtype() != s_i64.dtype());
161        assert!(s_i32.equals(&s_i64));
162    }
163
164    #[test]
165    fn test_df_equal() {
166        let a = Column::new("a".into(), [1, 2, 3].as_ref());
167        let b = Column::new("b".into(), [1, 2, 3].as_ref());
168
169        let df1 = DataFrame::new(vec![a, b]).unwrap();
170        assert!(df1.equals(&df1))
171    }
172
173    #[test]
174    fn assert_df_eq_passes() {
175        let df = df!("a" => [1], "b" => [2]).unwrap();
176        assert_df_eq!(df, df);
177        drop(df); // Ensure `assert_df_eq!` does not consume its arguments.
178    }
179
180    #[test]
181    #[should_panic(expected = "to equal")]
182    fn assert_df_eq_panics() {
183        assert_df_eq!(df!("a" => [1]).unwrap(), df!("a" => [2]).unwrap(),);
184    }
185
186    #[test]
187    fn test_df_partialeq() {
188        let df1 = df!("a" => &[1, 2, 3],
189                      "b" => &[4, 5, 6])
190        .unwrap();
191        let df2 = df!("b" => &[4, 5, 6],
192                      "a" => &[1, 2, 3])
193        .unwrap();
194        let df3 = df!("" => &[Some(1), None]).unwrap();
195        let df4 = df!("" => &[f32::NAN]).unwrap();
196
197        assert_eq!(df1, df1);
198        assert_ne!(df1, df2);
199        assert_eq!(df2, df2);
200        assert_ne!(df2, df3);
201        assert_eq!(df3, df3);
202        assert_eq!(df4, df4);
203    }
204}