datafusion_expr/
operation.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
18//! This module contains implementations of operations (unary, binary etc.) for DataFusion expressions.
19
20use crate::expr_fn::binary_expr;
21use crate::{Expr, Like};
22use datafusion_expr_common::operator::Operator;
23use std::ops::{self, Not};
24
25/// Support `<expr> + <expr>` fluent style
26impl ops::Add for Expr {
27    type Output = Self;
28
29    fn add(self, rhs: Self) -> Self {
30        binary_expr(self, Operator::Plus, rhs)
31    }
32}
33
34/// Support `<expr> - <expr>` fluent style
35impl ops::Sub for Expr {
36    type Output = Self;
37
38    fn sub(self, rhs: Self) -> Self {
39        binary_expr(self, Operator::Minus, rhs)
40    }
41}
42
43/// Support `<expr> * <expr>` fluent style
44impl ops::Mul for Expr {
45    type Output = Self;
46
47    fn mul(self, rhs: Self) -> Self {
48        binary_expr(self, Operator::Multiply, rhs)
49    }
50}
51
52/// Support `<expr> / <expr>` fluent style
53impl ops::Div for Expr {
54    type Output = Self;
55
56    fn div(self, rhs: Self) -> Self {
57        binary_expr(self, Operator::Divide, rhs)
58    }
59}
60
61/// Support `<expr> % <expr>` fluent style
62impl ops::Rem for Expr {
63    type Output = Self;
64
65    fn rem(self, rhs: Self) -> Self {
66        binary_expr(self, Operator::Modulo, rhs)
67    }
68}
69
70/// Support `<expr> & <expr>` fluent style
71impl ops::BitAnd for Expr {
72    type Output = Self;
73
74    fn bitand(self, rhs: Self) -> Self {
75        binary_expr(self, Operator::BitwiseAnd, rhs)
76    }
77}
78
79/// Support `<expr> | <expr>` fluent style
80impl ops::BitOr for Expr {
81    type Output = Self;
82
83    fn bitor(self, rhs: Self) -> Self {
84        binary_expr(self, Operator::BitwiseOr, rhs)
85    }
86}
87
88/// Support `<expr> ^ <expr>` fluent style
89impl ops::BitXor for Expr {
90    type Output = Self;
91
92    fn bitxor(self, rhs: Self) -> Self {
93        binary_expr(self, Operator::BitwiseXor, rhs)
94    }
95}
96
97/// Support `<expr> << <expr>` fluent style
98impl ops::Shl for Expr {
99    type Output = Self;
100
101    fn shl(self, rhs: Self) -> Self::Output {
102        binary_expr(self, Operator::BitwiseShiftLeft, rhs)
103    }
104}
105
106/// Support `<expr> >> <expr>` fluent style
107impl ops::Shr for Expr {
108    type Output = Self;
109
110    fn shr(self, rhs: Self) -> Self::Output {
111        binary_expr(self, Operator::BitwiseShiftRight, rhs)
112    }
113}
114
115/// Support `- <expr>` fluent style
116impl ops::Neg for Expr {
117    type Output = Self;
118
119    fn neg(self) -> Self::Output {
120        Expr::Negative(Box::new(self))
121    }
122}
123
124/// Support `NOT <expr>` fluent style
125impl Not for Expr {
126    type Output = Self;
127
128    fn not(self) -> Self::Output {
129        match self {
130            Expr::Like(Like {
131                negated,
132                expr,
133                pattern,
134                escape_char,
135                case_insensitive,
136            }) => Expr::Like(Like::new(
137                !negated,
138                expr,
139                pattern,
140                escape_char,
141                case_insensitive,
142            )),
143            Expr::SimilarTo(Like {
144                negated,
145                expr,
146                pattern,
147                escape_char,
148                case_insensitive,
149            }) => Expr::SimilarTo(Like::new(
150                !negated,
151                expr,
152                pattern,
153                escape_char,
154                case_insensitive,
155            )),
156            _ => Expr::Not(Box::new(self)),
157        }
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use crate::lit;
164
165    #[test]
166    fn test_operators() {
167        // Add
168        assert_eq!(
169            format!("{}", lit(1u32) + lit(2u32)),
170            "UInt32(1) + UInt32(2)"
171        );
172        // Sub
173        assert_eq!(
174            format!("{}", lit(1u32) - lit(2u32)),
175            "UInt32(1) - UInt32(2)"
176        );
177        // Mul
178        assert_eq!(
179            format!("{}", lit(1u32) * lit(2u32)),
180            "UInt32(1) * UInt32(2)"
181        );
182        // Div
183        assert_eq!(
184            format!("{}", lit(1u32) / lit(2u32)),
185            "UInt32(1) / UInt32(2)"
186        );
187        // Rem
188        assert_eq!(
189            format!("{}", lit(1u32) % lit(2u32)),
190            "UInt32(1) % UInt32(2)"
191        );
192        // BitAnd
193        assert_eq!(
194            format!("{}", lit(1u32) & lit(2u32)),
195            "UInt32(1) & UInt32(2)"
196        );
197        // BitOr
198        assert_eq!(
199            format!("{}", lit(1u32) | lit(2u32)),
200            "UInt32(1) | UInt32(2)"
201        );
202        // BitXor
203        assert_eq!(
204            format!("{}", lit(1u32) ^ lit(2u32)),
205            "UInt32(1) BIT_XOR UInt32(2)"
206        );
207        // Shl
208        assert_eq!(
209            format!("{}", lit(1u32) << lit(2u32)),
210            "UInt32(1) << UInt32(2)"
211        );
212        // Shr
213        assert_eq!(
214            format!("{}", lit(1u32) >> lit(2u32)),
215            "UInt32(1) >> UInt32(2)"
216        );
217        // Neg
218        assert_eq!(format!("{}", -lit(1u32)), "(- UInt32(1))");
219        // Not
220        assert_eq!(format!("{}", !lit(1u32)), "NOT UInt32(1)");
221    }
222}