Skip to main content

gluesql_core/plan/statement/query/
aggregation.rs

1use {
2    super::FilterPlan,
3    crate::plan::{AggregateExprPlan, ExprPlan, InnerJoinPlan, LeftOuterJoinPlan, SourcePlan},
4    serde::{Deserialize, Serialize},
5};
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
8pub enum AggregationInputPlan {
9    Source(SourcePlan),
10    InnerJoin(Box<InnerJoinPlan>),
11    LeftOuterJoin(Box<LeftOuterJoinPlan>),
12    Filter(FilterPlan),
13}
14
15impl AggregationInputPlan {
16    pub(crate) fn base_source(&self) -> &SourcePlan {
17        match self {
18            Self::Source(source) => source,
19            Self::InnerJoin(join) => join.base_source(),
20            Self::LeftOuterJoin(join) => join.base_source(),
21            Self::Filter(filter) => filter.input.base_source(),
22        }
23    }
24
25    pub(crate) fn joined_sources(&self) -> Vec<&SourcePlan> {
26        match self {
27            Self::Source(_) => Vec::new(),
28            Self::InnerJoin(join) => join.joined_sources(),
29            Self::LeftOuterJoin(join) => join.joined_sources(),
30            Self::Filter(filter) => filter.input.joined_sources(),
31        }
32    }
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
36pub struct AggregationPlan {
37    pub input: AggregationInputPlan,
38    pub group_by: Vec<ExprPlan>,
39    pub aggregate_slots: Vec<AggregateExprPlan>,
40}
41
42#[cfg(test)]
43mod tests {
44    use {
45        super::{AggregationInputPlan, AggregationPlan},
46        crate::{
47            data::Value,
48            plan::{
49                ExprPlan, FilterInputPlan, FilterPlan, InnerJoinInputPlan, InnerJoinPlan,
50                LeftOuterJoinInputPlan, LeftOuterJoinPlan, NestedLoopJoinInputPlan,
51                NestedLoopJoinPlan, SourcePlan, TableAccessPlan, TableSourcePlan,
52            },
53        },
54        pretty_assertions::assert_eq,
55    };
56
57    fn table(name: &str) -> SourcePlan {
58        SourcePlan::Table(TableSourcePlan {
59            name: name.to_owned(),
60            alias: None,
61            access: TableAccessPlan::FullScan,
62        })
63    }
64
65    #[test]
66    fn aggregation_accepts_relation_join_and_filter_inputs() {
67        let inner_join = InnerJoinPlan {
68            input: InnerJoinInputPlan::NestedLoop(NestedLoopJoinPlan {
69                input: NestedLoopJoinInputPlan::Source(table("A")),
70                right: table("B"),
71            }),
72        };
73        let left_outer_join = LeftOuterJoinPlan {
74            input: LeftOuterJoinInputPlan::NestedLoop(NestedLoopJoinPlan {
75                input: NestedLoopJoinInputPlan::Source(table("A")),
76                right: table("B"),
77            }),
78        };
79        let filter = FilterPlan {
80            input: FilterInputPlan::InnerJoin(Box::new(inner_join.clone())),
81            expr: ExprPlan::Value(Value::Bool(true)),
82        };
83        let relation = AggregationPlan {
84            input: AggregationInputPlan::Source(table("A")),
85            group_by: Vec::new(),
86            aggregate_slots: Vec::new(),
87        };
88        let inner = AggregationPlan {
89            input: AggregationInputPlan::InnerJoin(Box::new(inner_join.clone())),
90            group_by: Vec::new(),
91            aggregate_slots: Vec::new(),
92        };
93        let left_outer = AggregationPlan {
94            input: AggregationInputPlan::LeftOuterJoin(Box::new(left_outer_join.clone())),
95            group_by: Vec::new(),
96            aggregate_slots: Vec::new(),
97        };
98        let filtered = AggregationPlan {
99            input: AggregationInputPlan::Filter(filter.clone()),
100            group_by: Vec::new(),
101            aggregate_slots: Vec::new(),
102        };
103
104        assert_eq!(relation.input, AggregationInputPlan::Source(table("A")));
105        assert_eq!(
106            inner.input,
107            AggregationInputPlan::InnerJoin(Box::new(inner_join))
108        );
109        assert_eq!(
110            left_outer.input,
111            AggregationInputPlan::LeftOuterJoin(Box::new(left_outer_join))
112        );
113        assert_eq!(filtered.input, AggregationInputPlan::Filter(filter));
114
115        assert_eq!(relation.input.base_source(), &table("A"));
116        assert_eq!(relation.input.joined_sources(), Vec::<&SourcePlan>::new());
117        assert_eq!(inner.input.base_source(), &table("A"));
118        let expected = [table("B")];
119        assert_eq!(
120            inner.input.joined_sources(),
121            expected.iter().collect::<Vec<_>>()
122        );
123        assert_eq!(left_outer.input.base_source(), &table("A"));
124        let expected = [table("B")];
125        assert_eq!(
126            left_outer.input.joined_sources(),
127            expected.iter().collect::<Vec<_>>()
128        );
129        assert_eq!(filtered.input.base_source(), &table("A"));
130        let expected = [table("B")];
131        assert_eq!(
132            filtered.input.joined_sources(),
133            expected.iter().collect::<Vec<_>>()
134        );
135    }
136}