datafusion_physical_expr/equivalence/
mod.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::borrow::Borrow;
19use std::sync::Arc;
20
21use crate::PhysicalExpr;
22
23use arrow::compute::SortOptions;
24use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
25
26mod class;
27mod ordering;
28mod properties;
29
30pub use class::{AcrossPartitions, ConstExpr, EquivalenceClass, EquivalenceGroup};
31pub use ordering::OrderingEquivalenceClass;
32// Re-export for backwards compatibility, we recommend importing from
33// datafusion_physical_expr::projection instead
34pub use crate::projection::{ProjectionMapping, project_ordering, project_orderings};
35pub use properties::{
36    EquivalenceProperties, calculate_union, join_equivalence_properties,
37};
38
39// Convert each tuple to a `PhysicalSortExpr` and construct a vector.
40pub fn convert_to_sort_exprs<T: Borrow<Arc<dyn PhysicalExpr>>>(
41    args: &[(T, SortOptions)],
42) -> Vec<PhysicalSortExpr> {
43    args.iter()
44        .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr.borrow()), *options))
45        .collect()
46}
47
48// Convert each vector of tuples to a `LexOrdering`.
49pub fn convert_to_orderings<T: Borrow<Arc<dyn PhysicalExpr>>>(
50    args: &[Vec<(T, SortOptions)>],
51) -> Vec<LexOrdering> {
52    args.iter()
53        .filter_map(|sort_exprs| LexOrdering::new(convert_to_sort_exprs(sort_exprs)))
54        .collect()
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::expressions::{Column, col};
61    use crate::{LexRequirement, PhysicalSortExpr};
62
63    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
64    use datafusion_common::Result;
65    use datafusion_physical_expr_common::sort_expr::PhysicalSortRequirement;
66
67    /// Converts a string to a physical sort expression
68    ///
69    /// # Example
70    /// * `"a"` -> (`"a"`, `SortOptions::default()`)
71    /// * `"a ASC"` -> (`"a"`, `SortOptions { descending: false, nulls_first: false }`)
72    pub fn parse_sort_expr(name: &str, schema: &SchemaRef) -> PhysicalSortExpr {
73        let mut parts = name.split_whitespace();
74        let name = parts.next().expect("empty sort expression");
75        let mut sort_expr = PhysicalSortExpr::new(
76            col(name, schema).expect("invalid column name"),
77            SortOptions::default(),
78        );
79
80        if let Some(options) = parts.next() {
81            sort_expr = match options {
82                "ASC" => sort_expr.asc(),
83                "DESC" => sort_expr.desc(),
84                _ => panic!(
85                    "unknown sort options. Expected 'ASC' or 'DESC', got {options}"
86                ),
87            }
88        }
89
90        assert!(
91            parts.next().is_none(),
92            "unexpected tokens in column name. Expected 'name' / 'name ASC' / 'name DESC' but got  '{name}'"
93        );
94
95        sort_expr
96    }
97
98    // Generate a schema which consists of 8 columns (a, b, c, d, e, f, g, h)
99    pub fn create_test_schema() -> Result<SchemaRef> {
100        let a = Field::new("a", DataType::Int32, true);
101        let b = Field::new("b", DataType::Int32, true);
102        let c = Field::new("c", DataType::Int32, true);
103        let d = Field::new("d", DataType::Int32, true);
104        let e = Field::new("e", DataType::Int32, true);
105        let f = Field::new("f", DataType::Int32, true);
106        let g = Field::new("g", DataType::Int32, true);
107        let h = Field::new("h", DataType::Int32, true);
108        let schema = Arc::new(Schema::new(vec![a, b, c, d, e, f, g, h]));
109
110        Ok(schema)
111    }
112
113    /// Construct a schema with following properties
114    /// Schema satisfies following orderings:
115    /// [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC]
116    /// and
117    /// Column [a=c] (e.g they are aliases).
118    pub fn create_test_params() -> Result<(SchemaRef, EquivalenceProperties)> {
119        let test_schema = create_test_schema()?;
120        let col_a = col("a", &test_schema)?;
121        let col_b = col("b", &test_schema)?;
122        let col_c = col("c", &test_schema)?;
123        let col_d = col("d", &test_schema)?;
124        let col_e = col("e", &test_schema)?;
125        let col_f = col("f", &test_schema)?;
126        let col_g = col("g", &test_schema)?;
127        let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema));
128        eq_properties.add_equal_conditions(Arc::clone(&col_a), Arc::clone(&col_c))?;
129
130        let option_asc = SortOptions {
131            descending: false,
132            nulls_first: false,
133        };
134        let option_desc = SortOptions {
135            descending: true,
136            nulls_first: true,
137        };
138        let orderings = vec![
139            // [a ASC]
140            vec![(col_a, option_asc)],
141            // [d ASC, b ASC]
142            vec![(col_d, option_asc), (col_b, option_asc)],
143            // [e DESC, f ASC, g ASC]
144            vec![
145                (col_e, option_desc),
146                (col_f, option_asc),
147                (col_g, option_asc),
148            ],
149        ];
150        let orderings = convert_to_orderings(&orderings);
151        eq_properties.add_orderings(orderings);
152        Ok((test_schema, eq_properties))
153    }
154
155    // Convert each tuple to a `PhysicalSortRequirement` and construct a
156    // a `LexRequirement` from them.
157    pub fn convert_to_sort_reqs(
158        args: &[(&Arc<dyn PhysicalExpr>, Option<SortOptions>)],
159    ) -> LexRequirement {
160        let exprs = args.iter().map(|(expr, options)| {
161            PhysicalSortRequirement::new(Arc::clone(*expr), *options)
162        });
163        LexRequirement::new(exprs).unwrap()
164    }
165
166    #[test]
167    fn add_equal_conditions_test() -> Result<()> {
168        let schema = Arc::new(Schema::new(vec![
169            Field::new("a", DataType::Int64, true),
170            Field::new("b", DataType::Int64, true),
171            Field::new("c", DataType::Int64, true),
172            Field::new("x", DataType::Int64, true),
173            Field::new("y", DataType::Int64, true),
174        ]));
175
176        let mut eq_properties = EquivalenceProperties::new(schema);
177        let col_a = Arc::new(Column::new("a", 0)) as _;
178        let col_b = Arc::new(Column::new("b", 1)) as _;
179        let col_c = Arc::new(Column::new("c", 2)) as _;
180        let col_x = Arc::new(Column::new("x", 3)) as _;
181        let col_y = Arc::new(Column::new("y", 4)) as _;
182
183        // a and b are aliases
184        eq_properties.add_equal_conditions(Arc::clone(&col_a), Arc::clone(&col_b))?;
185        assert_eq!(eq_properties.eq_group().len(), 1);
186
187        // This new entry is redundant, size shouldn't increase
188        eq_properties.add_equal_conditions(Arc::clone(&col_b), Arc::clone(&col_a))?;
189        assert_eq!(eq_properties.eq_group().len(), 1);
190        let eq_groups = eq_properties.eq_group().iter().next().unwrap();
191        assert_eq!(eq_groups.len(), 2);
192        assert!(eq_groups.contains(&col_a));
193        assert!(eq_groups.contains(&col_b));
194
195        // b and c are aliases. Existing equivalence class should expand,
196        // however there shouldn't be any new equivalence class
197        eq_properties.add_equal_conditions(Arc::clone(&col_b), Arc::clone(&col_c))?;
198        assert_eq!(eq_properties.eq_group().len(), 1);
199        let eq_groups = eq_properties.eq_group().iter().next().unwrap();
200        assert_eq!(eq_groups.len(), 3);
201        assert!(eq_groups.contains(&col_a));
202        assert!(eq_groups.contains(&col_b));
203        assert!(eq_groups.contains(&col_c));
204
205        // This is a new set of equality. Hence equivalent class count should be 2.
206        eq_properties.add_equal_conditions(Arc::clone(&col_x), Arc::clone(&col_y))?;
207        assert_eq!(eq_properties.eq_group().len(), 2);
208
209        // This equality bridges distinct equality sets.
210        // Hence equivalent class count should decrease from 2 to 1.
211        eq_properties.add_equal_conditions(Arc::clone(&col_x), Arc::clone(&col_a))?;
212        assert_eq!(eq_properties.eq_group().len(), 1);
213        let eq_groups = eq_properties.eq_group().iter().next().unwrap();
214        assert_eq!(eq_groups.len(), 5);
215        assert!(eq_groups.contains(&col_a));
216        assert!(eq_groups.contains(&col_b));
217        assert!(eq_groups.contains(&col_c));
218        assert!(eq_groups.contains(&col_x));
219        assert!(eq_groups.contains(&col_y));
220
221        Ok(())
222    }
223}
datafusion_physical_expr/equivalence/mod.rs

datafusion_physical_expr/equivalence/
mod.rs
