Skip to main content

datatypes/vectors/
eq.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use common_time::interval::IntervalUnit;
18
19use crate::data_type::DataType;
20use crate::types::{DurationType, TimeType, TimestampType};
21use crate::vectors::struct_vector::StructVector;
22use crate::vectors::{
23    BinaryVector, BooleanVector, DateVector, Decimal128Vector, DurationMicrosecondVector,
24    DurationMillisecondVector, DurationNanosecondVector, DurationSecondVector,
25    IntervalDayTimeVector, IntervalMonthDayNanoVector, IntervalYearMonthVector, ListVector,
26    PrimitiveVector, StringVector, TimeMicrosecondVector, TimeMillisecondVector,
27    TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, TimestampMillisecondVector,
28    TimestampNanosecondVector, TimestampSecondVector, Vector,
29};
30use crate::with_match_primitive_type_id;
31
32impl Eq for dyn Vector + '_ {}
33
34impl PartialEq for dyn Vector + '_ {
35    fn eq(&self, other: &dyn Vector) -> bool {
36        equal(self, other)
37    }
38}
39
40impl PartialEq<dyn Vector> for Arc<dyn Vector + '_> {
41    fn eq(&self, other: &dyn Vector) -> bool {
42        equal(&**self, other)
43    }
44}
45
46macro_rules! is_vector_eq {
47    ($VectorType: ident, $lhs: ident, $rhs: ident) => {{
48        let lhs = $lhs.as_any().downcast_ref::<$VectorType>().unwrap();
49        let rhs = $rhs.as_any().downcast_ref::<$VectorType>().unwrap();
50
51        lhs == rhs
52    }};
53}
54
55fn equal(lhs: &dyn Vector, rhs: &dyn Vector) -> bool {
56    if lhs.data_type() != rhs.data_type() || lhs.len() != rhs.len() {
57        return false;
58    }
59
60    use crate::data_type::ConcreteDataType::*;
61
62    let lhs_type = lhs.data_type();
63    match lhs.data_type() {
64        Null(_) => true,
65        Boolean(_) => is_vector_eq!(BooleanVector, lhs, rhs),
66        Binary(_) | Json(_) | Vector(_) => is_vector_eq!(BinaryVector, lhs, rhs),
67        String(_) => is_vector_eq!(StringVector, lhs, rhs),
68        Date(_) => is_vector_eq!(DateVector, lhs, rhs),
69        Timestamp(t) => match t {
70            TimestampType::Second(_) => {
71                is_vector_eq!(TimestampSecondVector, lhs, rhs)
72            }
73            TimestampType::Millisecond(_) => {
74                is_vector_eq!(TimestampMillisecondVector, lhs, rhs)
75            }
76            TimestampType::Microsecond(_) => {
77                is_vector_eq!(TimestampMicrosecondVector, lhs, rhs)
78            }
79            TimestampType::Nanosecond(_) => {
80                is_vector_eq!(TimestampNanosecondVector, lhs, rhs)
81            }
82        },
83        Interval(v) => match v.unit() {
84            IntervalUnit::YearMonth => {
85                is_vector_eq!(IntervalYearMonthVector, lhs, rhs)
86            }
87            IntervalUnit::DayTime => {
88                is_vector_eq!(IntervalDayTimeVector, lhs, rhs)
89            }
90            IntervalUnit::MonthDayNano => {
91                is_vector_eq!(IntervalMonthDayNanoVector, lhs, rhs)
92            }
93        },
94        List(_) => is_vector_eq!(ListVector, lhs, rhs),
95        Struct(_) => is_vector_eq!(StructVector, lhs, rhs),
96        UInt8(_) | UInt16(_) | UInt32(_) | UInt64(_) | Int8(_) | Int16(_) | Int32(_) | Int64(_)
97        | Float32(_) | Float64(_) | Dictionary(_) => {
98            with_match_primitive_type_id!(lhs_type.logical_type_id(), |$T| {
99                let lhs = lhs.as_any().downcast_ref::<PrimitiveVector<$T>>().unwrap();
100                let rhs = rhs.as_any().downcast_ref::<PrimitiveVector<$T>>().unwrap();
101
102                lhs == rhs
103            },
104            {
105                unreachable!("should not compare {} with {}", lhs.vector_type_name(), rhs.vector_type_name())
106            })
107        }
108
109        Time(t) => match t {
110            TimeType::Second(_) => {
111                is_vector_eq!(TimeSecondVector, lhs, rhs)
112            }
113            TimeType::Millisecond(_) => {
114                is_vector_eq!(TimeMillisecondVector, lhs, rhs)
115            }
116            TimeType::Microsecond(_) => {
117                is_vector_eq!(TimeMicrosecondVector, lhs, rhs)
118            }
119            TimeType::Nanosecond(_) => {
120                is_vector_eq!(TimeNanosecondVector, lhs, rhs)
121            }
122        },
123        Duration(d) => match d {
124            DurationType::Second(_) => {
125                is_vector_eq!(DurationSecondVector, lhs, rhs)
126            }
127            DurationType::Millisecond(_) => {
128                is_vector_eq!(DurationMillisecondVector, lhs, rhs)
129            }
130            DurationType::Microsecond(_) => {
131                is_vector_eq!(DurationMicrosecondVector, lhs, rhs)
132            }
133            DurationType::Nanosecond(_) => {
134                is_vector_eq!(DurationNanosecondVector, lhs, rhs)
135            }
136        },
137        Decimal128(_) => {
138            is_vector_eq!(Decimal128Vector, lhs, rhs)
139        }
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano};
146
147    use super::*;
148    use crate::vectors::{
149        DurationMicrosecondVector, DurationMillisecondVector, DurationNanosecondVector,
150        DurationSecondVector, Float32Vector, Float64Vector, Int8Vector, Int16Vector, Int32Vector,
151        Int64Vector, NullVector, UInt8Vector, UInt16Vector, UInt32Vector, UInt64Vector, VectorRef,
152        list,
153    };
154
155    fn assert_vector_ref_eq(vector: VectorRef) {
156        let rhs = vector.clone();
157        assert_eq!(vector, rhs);
158        assert_dyn_vector_eq(&*vector, &*rhs);
159    }
160
161    fn assert_dyn_vector_eq(lhs: &dyn Vector, rhs: &dyn Vector) {
162        assert_eq!(lhs, rhs);
163    }
164
165    fn assert_vector_ref_ne(lhs: VectorRef, rhs: VectorRef) {
166        assert_ne!(lhs, rhs);
167    }
168
169    #[test]
170    fn test_vector_eq() {
171        assert_vector_ref_eq(Arc::new(BinaryVector::from(vec![
172            Some(b"hello".to_vec()),
173            Some(b"world".to_vec()),
174        ])));
175        assert_vector_ref_eq(Arc::new(BooleanVector::from(vec![true, false])));
176        assert_vector_ref_eq(Arc::new(BooleanVector::from(vec![true, false])));
177        assert_vector_ref_eq(Arc::new(DateVector::from(vec![Some(100), Some(120)])));
178        assert_vector_ref_eq(Arc::new(TimestampSecondVector::from_values([100, 120])));
179        assert_vector_ref_eq(Arc::new(TimestampMillisecondVector::from_values([
180            100, 120,
181        ])));
182        assert_vector_ref_eq(Arc::new(TimestampMicrosecondVector::from_values([
183            100, 120,
184        ])));
185        assert_vector_ref_eq(Arc::new(TimestampNanosecondVector::from_values([100, 120])));
186
187        let list_vector = list::tests::new_list_vector(&[
188            Some(vec![Some(1), Some(2)]),
189            None,
190            Some(vec![Some(3), Some(4)]),
191        ]);
192        assert_vector_ref_eq(Arc::new(list_vector));
193
194        assert_vector_ref_eq(Arc::new(NullVector::new(4)));
195        assert_vector_ref_eq(Arc::new(StringVector::from(vec![
196            Some("hello"),
197            Some("world"),
198        ])));
199
200        assert_vector_ref_eq(Arc::new(Int8Vector::from_slice([1, 2, 3, 4])));
201        assert_vector_ref_eq(Arc::new(UInt8Vector::from_slice([1, 2, 3, 4])));
202        assert_vector_ref_eq(Arc::new(Int16Vector::from_slice([1, 2, 3, 4])));
203        assert_vector_ref_eq(Arc::new(UInt16Vector::from_slice([1, 2, 3, 4])));
204        assert_vector_ref_eq(Arc::new(Int32Vector::from_slice([1, 2, 3, 4])));
205        assert_vector_ref_eq(Arc::new(UInt32Vector::from_slice([1, 2, 3, 4])));
206        assert_vector_ref_eq(Arc::new(Int64Vector::from_slice([1, 2, 3, 4])));
207        assert_vector_ref_eq(Arc::new(UInt64Vector::from_slice([1, 2, 3, 4])));
208        assert_vector_ref_eq(Arc::new(Float32Vector::from_slice([1.0, 2.0, 3.0, 4.0])));
209        assert_vector_ref_eq(Arc::new(Float64Vector::from_slice([1.0, 2.0, 3.0, 4.0])));
210
211        assert_vector_ref_eq(Arc::new(TimeSecondVector::from_values([100, 120])));
212        assert_vector_ref_eq(Arc::new(TimeMillisecondVector::from_values([100, 120])));
213        assert_vector_ref_eq(Arc::new(TimeMicrosecondVector::from_values([100, 120])));
214        assert_vector_ref_eq(Arc::new(TimeNanosecondVector::from_values([100, 120])));
215
216        assert_vector_ref_eq(Arc::new(IntervalYearMonthVector::from_values([
217            1000, 2000, 3000, 4000,
218        ])));
219        assert_vector_ref_eq(Arc::new(IntervalDayTimeVector::from_values([
220            IntervalDayTime::new(1, 1000),
221            IntervalDayTime::new(1, 2000),
222            IntervalDayTime::new(1, 3000),
223            IntervalDayTime::new(1, 4000),
224        ])));
225        assert_vector_ref_eq(Arc::new(IntervalMonthDayNanoVector::from_values([
226            IntervalMonthDayNano::new(1, 1, 1000),
227            IntervalMonthDayNano::new(1, 1, 2000),
228            IntervalMonthDayNano::new(1, 1, 3000),
229            IntervalMonthDayNano::new(1, 1, 4000),
230        ])));
231        assert_vector_ref_eq(Arc::new(DurationSecondVector::from_values([300, 310])));
232        assert_vector_ref_eq(Arc::new(DurationMillisecondVector::from_values([300, 310])));
233        assert_vector_ref_eq(Arc::new(DurationMicrosecondVector::from_values([300, 310])));
234        assert_vector_ref_eq(Arc::new(DurationNanosecondVector::from_values([300, 310])));
235        assert_vector_ref_eq(Arc::new(Decimal128Vector::from_values(vec![
236            1i128, 2i128, 3i128,
237        ])));
238    }
239
240    #[test]
241    fn test_vector_ne() {
242        assert_vector_ref_ne(
243            Arc::new(Int32Vector::from_slice([1, 2, 3, 4])),
244            Arc::new(Int32Vector::from_slice([1, 2])),
245        );
246        assert_vector_ref_ne(
247            Arc::new(Int32Vector::from_slice([1, 2, 3, 4])),
248            Arc::new(Int8Vector::from_slice([1, 2, 3, 4])),
249        );
250        assert_vector_ref_ne(
251            Arc::new(Int32Vector::from_slice([1, 2, 3, 4])),
252            Arc::new(BooleanVector::from(vec![true, true])),
253        );
254        assert_vector_ref_ne(Arc::new(NullVector::new(5)), Arc::new(NullVector::new(8)));
255
256        assert_vector_ref_ne(
257            Arc::new(TimeMicrosecondVector::from_values([100, 120])),
258            Arc::new(TimeMicrosecondVector::from_values([200, 220])),
259        );
260
261        assert_vector_ref_ne(
262            Arc::new(IntervalDayTimeVector::from_values([
263                IntervalDayTime::new(1, 1000),
264                IntervalDayTime::new(1, 2000),
265            ])),
266            Arc::new(IntervalDayTimeVector::from_values([
267                IntervalDayTime::new(1, 2100),
268                IntervalDayTime::new(1, 1200),
269            ])),
270        );
271        assert_vector_ref_ne(
272            Arc::new(IntervalMonthDayNanoVector::from_values([
273                IntervalMonthDayNano::new(1, 1, 1000),
274                IntervalMonthDayNano::new(1, 1, 2000),
275            ])),
276            Arc::new(IntervalMonthDayNanoVector::from_values([
277                IntervalMonthDayNano::new(1, 1, 2100),
278                IntervalMonthDayNano::new(1, 1, 1200),
279            ])),
280        );
281        assert_vector_ref_ne(
282            Arc::new(IntervalYearMonthVector::from_values([1000, 2000])),
283            Arc::new(IntervalYearMonthVector::from_values([2100, 1200])),
284        );
285
286        assert_vector_ref_ne(
287            Arc::new(DurationSecondVector::from_values([300, 310])),
288            Arc::new(DurationSecondVector::from_values([300, 320])),
289        );
290
291        assert_vector_ref_ne(
292            Arc::new(Decimal128Vector::from_values([300i128, 310i128])),
293            Arc::new(Decimal128Vector::from_values([300i128, 320i128])),
294        );
295    }
296}