gluesql_core/plan/statement/query/
aggregation.rs1use {
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}