datafusion_optimizer/
eliminate_one_union.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//! [`EliminateOneUnion`]  eliminates single element `Union`
19
20use crate::{OptimizerConfig, OptimizerRule};
21use datafusion_common::{tree_node::Transformed, Result};
22use datafusion_expr::logical_plan::{LogicalPlan, Union};
23use std::sync::Arc;
24
25use crate::optimizer::ApplyOrder;
26
27#[derive(Default, Debug)]
28/// An optimization rule that eliminates union with one element.
29pub struct EliminateOneUnion;
30
31impl EliminateOneUnion {
32    #[allow(missing_docs)]
33    pub fn new() -> Self {
34        Self {}
35    }
36}
37
38impl OptimizerRule for EliminateOneUnion {
39    fn name(&self) -> &str {
40        "eliminate_one_union"
41    }
42
43    fn supports_rewrite(&self) -> bool {
44        true
45    }
46
47    fn rewrite(
48        &self,
49        plan: LogicalPlan,
50        _config: &dyn OptimizerConfig,
51    ) -> Result<Transformed<LogicalPlan>> {
52        match plan {
53            LogicalPlan::Union(Union { mut inputs, .. }) if inputs.len() == 1 => Ok(
54                Transformed::yes(Arc::unwrap_or_clone(inputs.pop().unwrap())),
55            ),
56            _ => Ok(Transformed::no(plan)),
57        }
58    }
59
60    fn apply_order(&self) -> Option<ApplyOrder> {
61        Some(ApplyOrder::TopDown)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use crate::test::*;
69    use arrow::datatypes::{DataType, Field, Schema};
70    use datafusion_common::ToDFSchema;
71    use datafusion_expr::{
72        expr_rewriter::coerce_plan_expr_for_schema, logical_plan::table_scan,
73    };
74    use std::sync::Arc;
75
76    fn schema() -> Schema {
77        Schema::new(vec![
78            Field::new("id", DataType::Int32, false),
79            Field::new("key", DataType::Utf8, false),
80            Field::new("value", DataType::Int32, false),
81        ])
82    }
83
84    fn assert_optimized_plan_equal(plan: LogicalPlan, expected: &str) -> Result<()> {
85        assert_optimized_plan_with_rules(
86            vec![Arc::new(EliminateOneUnion::new())],
87            plan,
88            expected,
89            true,
90        )
91    }
92
93    #[test]
94    fn eliminate_nothing() -> Result<()> {
95        let plan_builder = table_scan(Some("table"), &schema(), None)?;
96
97        let plan = plan_builder.clone().union(plan_builder.build()?)?.build()?;
98
99        let expected = "\
100        Union\
101        \n  TableScan: table\
102        \n  TableScan: table";
103        assert_optimized_plan_equal(plan, expected)
104    }
105
106    #[test]
107    fn eliminate_one_union() -> Result<()> {
108        let table_plan = coerce_plan_expr_for_schema(
109            table_scan(Some("table"), &schema(), None)?.build()?,
110            &schema().to_dfschema()?,
111        )?;
112        let schema = Arc::clone(table_plan.schema());
113        let single_union_plan = LogicalPlan::Union(Union {
114            inputs: vec![Arc::new(table_plan)],
115            schema,
116        });
117
118        let expected = "TableScan: table";
119        assert_optimized_plan_equal(single_union_plan, expected)
120    }
121}
datafusion_optimizer/eliminate_one_union.rs

datafusion_optimizer/
eliminate_one_union.rs
