datafusion_physical_expr/
physical_expr.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use datafusion_common::HashMap;
21pub(crate) use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
22pub use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
23use itertools::izip;
24
25/// This function is similar to the `contains` method of `Vec`. It finds
26/// whether `expr` is among `physical_exprs`.
27pub fn physical_exprs_contains(
28    physical_exprs: &[Arc<dyn PhysicalExpr>],
29    expr: &Arc<dyn PhysicalExpr>,
30) -> bool {
31    physical_exprs
32        .iter()
33        .any(|physical_expr| physical_expr.eq(expr))
34}
35
36/// Checks whether the given physical expression slices are equal.
37pub fn physical_exprs_equal(
38    lhs: &[Arc<dyn PhysicalExpr>],
39    rhs: &[Arc<dyn PhysicalExpr>],
40) -> bool {
41    lhs.len() == rhs.len() && izip!(lhs, rhs).all(|(lhs, rhs)| lhs.eq(rhs))
42}
43
44/// Checks whether the given physical expression slices are equal in the sense
45/// of bags (multi-sets), disregarding their orderings.
46pub fn physical_exprs_bag_equal(
47    lhs: &[Arc<dyn PhysicalExpr>],
48    rhs: &[Arc<dyn PhysicalExpr>],
49) -> bool {
50    let mut multi_set_lhs: HashMap<_, usize> = HashMap::new();
51    let mut multi_set_rhs: HashMap<_, usize> = HashMap::new();
52    for expr in lhs {
53        *multi_set_lhs.entry(expr).or_insert(0) += 1;
54    }
55    for expr in rhs {
56        *multi_set_rhs.entry(expr).or_insert(0) += 1;
57    }
58    multi_set_lhs == multi_set_rhs
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    use crate::expressions::{Column, Literal};
66    use crate::physical_expr::{
67        physical_exprs_bag_equal, physical_exprs_contains, physical_exprs_equal,
68    };
69
70    use datafusion_common::ScalarValue;
71
72    #[test]
73    fn test_physical_exprs_contains() {
74        let lit_true = Arc::new(Literal::new(ScalarValue::Boolean(Some(true))))
75            as Arc<dyn PhysicalExpr>;
76        let lit_false = Arc::new(Literal::new(ScalarValue::Boolean(Some(false))))
77            as Arc<dyn PhysicalExpr>;
78        let lit4 =
79            Arc::new(Literal::new(ScalarValue::Int32(Some(4)))) as Arc<dyn PhysicalExpr>;
80        let lit2 =
81            Arc::new(Literal::new(ScalarValue::Int32(Some(2)))) as Arc<dyn PhysicalExpr>;
82        let lit1 =
83            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
84        let col_a_expr = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
85        let col_b_expr = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
86        let col_c_expr = Arc::new(Column::new("c", 2)) as Arc<dyn PhysicalExpr>;
87
88        // lit(true), lit(false), lit(4), lit(2), Col(a), Col(b)
89        let physical_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
90            Arc::clone(&lit_true),
91            Arc::clone(&lit_false),
92            Arc::clone(&lit4),
93            Arc::clone(&lit2),
94            Arc::clone(&col_a_expr),
95            Arc::clone(&col_b_expr),
96        ];
97        // below expressions are inside physical_exprs
98        assert!(physical_exprs_contains(&physical_exprs, &lit_true));
99        assert!(physical_exprs_contains(&physical_exprs, &lit2));
100        assert!(physical_exprs_contains(&physical_exprs, &col_b_expr));
101
102        // below expressions are not inside physical_exprs
103        assert!(!physical_exprs_contains(&physical_exprs, &col_c_expr));
104        assert!(!physical_exprs_contains(&physical_exprs, &lit1));
105    }
106
107    #[test]
108    fn test_physical_exprs_equal() {
109        let lit_true = Arc::new(Literal::new(ScalarValue::Boolean(Some(true))))
110            as Arc<dyn PhysicalExpr>;
111        let lit_false = Arc::new(Literal::new(ScalarValue::Boolean(Some(false))))
112            as Arc<dyn PhysicalExpr>;
113        let lit1 =
114            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
115        let lit2 =
116            Arc::new(Literal::new(ScalarValue::Int32(Some(2)))) as Arc<dyn PhysicalExpr>;
117        let col_b_expr = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
118
119        let vec1 = vec![Arc::clone(&lit_true), Arc::clone(&lit_false)];
120        let vec2 = vec![Arc::clone(&lit_true), Arc::clone(&col_b_expr)];
121        let vec3 = vec![Arc::clone(&lit2), Arc::clone(&lit1)];
122        let vec4 = vec![Arc::clone(&lit_true), Arc::clone(&lit_false)];
123
124        // these vectors are same
125        assert!(physical_exprs_equal(&vec1, &vec1));
126        assert!(physical_exprs_equal(&vec1, &vec4));
127        assert!(physical_exprs_bag_equal(&vec1, &vec1));
128        assert!(physical_exprs_bag_equal(&vec1, &vec4));
129
130        // these vectors are different
131        assert!(!physical_exprs_equal(&vec1, &vec2));
132        assert!(!physical_exprs_equal(&vec1, &vec3));
133        assert!(!physical_exprs_bag_equal(&vec1, &vec2));
134        assert!(!physical_exprs_bag_equal(&vec1, &vec3));
135    }
136
137    #[test]
138    fn test_physical_exprs_set_equal() {
139        let list1: Vec<Arc<dyn PhysicalExpr>> = vec![
140            Arc::new(Column::new("a", 0)),
141            Arc::new(Column::new("a", 0)),
142            Arc::new(Column::new("b", 1)),
143        ];
144        let list2: Vec<Arc<dyn PhysicalExpr>> = vec![
145            Arc::new(Column::new("b", 1)),
146            Arc::new(Column::new("b", 1)),
147            Arc::new(Column::new("a", 0)),
148        ];
149        assert!(!physical_exprs_bag_equal(
150            list1.as_slice(),
151            list2.as_slice()
152        ));
153        assert!(!physical_exprs_bag_equal(
154            list2.as_slice(),
155            list1.as_slice()
156        ));
157        assert!(!physical_exprs_equal(list1.as_slice(), list2.as_slice()));
158        assert!(!physical_exprs_equal(list2.as_slice(), list1.as_slice()));
159
160        let list3: Vec<Arc<dyn PhysicalExpr>> = vec![
161            Arc::new(Column::new("a", 0)),
162            Arc::new(Column::new("b", 1)),
163            Arc::new(Column::new("c", 2)),
164            Arc::new(Column::new("a", 0)),
165            Arc::new(Column::new("b", 1)),
166        ];
167        let list4: Vec<Arc<dyn PhysicalExpr>> = vec![
168            Arc::new(Column::new("b", 1)),
169            Arc::new(Column::new("b", 1)),
170            Arc::new(Column::new("a", 0)),
171            Arc::new(Column::new("c", 2)),
172            Arc::new(Column::new("a", 0)),
173        ];
174        assert!(physical_exprs_bag_equal(list3.as_slice(), list4.as_slice()));
175        assert!(physical_exprs_bag_equal(list4.as_slice(), list3.as_slice()));
176        assert!(physical_exprs_bag_equal(list3.as_slice(), list3.as_slice()));
177        assert!(physical_exprs_bag_equal(list4.as_slice(), list4.as_slice()));
178        assert!(!physical_exprs_equal(list3.as_slice(), list4.as_slice()));
179        assert!(!physical_exprs_equal(list4.as_slice(), list3.as_slice()));
180        assert!(physical_exprs_bag_equal(list3.as_slice(), list3.as_slice()));
181        assert!(physical_exprs_bag_equal(list4.as_slice(), list4.as_slice()));
182    }
183}