datafusion_physical_expr/intervals/
test_utils.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//! Test utilities for the interval arithmetic library
19
20use std::sync::Arc;
21
22use crate::expressions::{binary, BinaryExpr, Literal};
23use crate::PhysicalExpr;
24use arrow::datatypes::Schema;
25use datafusion_common::{DataFusionError, ScalarValue};
26use datafusion_expr::Operator;
27
28#[allow(clippy::too_many_arguments)]
29/// This test function generates a conjunctive statement with two numeric
30/// terms with the following form:
31/// left_col (op_1) a  >/>= right_col (op_2) b AND left_col (op_3) c </<= right_col (op_4) d
32pub fn gen_conjunctive_numerical_expr(
33    left_col: Arc<dyn PhysicalExpr>,
34    right_col: Arc<dyn PhysicalExpr>,
35    op: (Operator, Operator, Operator, Operator),
36    a: ScalarValue,
37    b: ScalarValue,
38    c: ScalarValue,
39    d: ScalarValue,
40    bounds: (Operator, Operator),
41) -> Arc<dyn PhysicalExpr> {
42    let (op_1, op_2, op_3, op_4) = op;
43    let left_and_1 = Arc::new(BinaryExpr::new(
44        Arc::clone(&left_col),
45        op_1,
46        Arc::new(Literal::new(a)),
47    ));
48    let left_and_2 = Arc::new(BinaryExpr::new(
49        Arc::clone(&right_col),
50        op_2,
51        Arc::new(Literal::new(b)),
52    ));
53    let right_and_1 =
54        Arc::new(BinaryExpr::new(left_col, op_3, Arc::new(Literal::new(c))));
55    let right_and_2 =
56        Arc::new(BinaryExpr::new(right_col, op_4, Arc::new(Literal::new(d))));
57    let (greater_op, less_op) = bounds;
58
59    let left_expr = Arc::new(BinaryExpr::new(left_and_1, greater_op, left_and_2));
60    let right_expr = Arc::new(BinaryExpr::new(right_and_1, less_op, right_and_2));
61    Arc::new(BinaryExpr::new(left_expr, Operator::And, right_expr))
62}
63
64#[allow(clippy::too_many_arguments)]
65/// This test function generates a conjunctive statement with
66/// two scalar values with the following form:
67/// left_col (op_1) a  > right_col (op_2) b AND left_col (op_3) c < right_col (op_4) d
68pub fn gen_conjunctive_temporal_expr(
69    left_col: Arc<dyn PhysicalExpr>,
70    right_col: Arc<dyn PhysicalExpr>,
71    op_1: Operator,
72    op_2: Operator,
73    op_3: Operator,
74    op_4: Operator,
75    a: ScalarValue,
76    b: ScalarValue,
77    c: ScalarValue,
78    d: ScalarValue,
79    schema: &Schema,
80) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
81    let left_and_1 = binary(
82        Arc::clone(&left_col),
83        op_1,
84        Arc::new(Literal::new(a)),
85        schema,
86    )?;
87    let left_and_2 = binary(
88        Arc::clone(&right_col),
89        op_2,
90        Arc::new(Literal::new(b)),
91        schema,
92    )?;
93    let right_and_1 = binary(left_col, op_3, Arc::new(Literal::new(c)), schema)?;
94    let right_and_2 = binary(right_col, op_4, Arc::new(Literal::new(d)), schema)?;
95    let left_expr = Arc::new(BinaryExpr::new(left_and_1, Operator::Gt, left_and_2));
96    let right_expr = Arc::new(BinaryExpr::new(right_and_1, Operator::Lt, right_and_2));
97    Ok(Arc::new(BinaryExpr::new(
98        left_expr,
99        Operator::And,
100        right_expr,
101    )))
102}