Skip to main content

gluesql_core/query_builder/select/
distinct.rs

1use {
2    super::{BuildProjectPlan, BuildQuery, BuildQueryPlan, BuildSelect},
3    crate::{
4        ast::{OrderByExpr, Query, Select, SetExpr},
5        plan::{DistinctInputPlan, DistinctPlan, QueryPlan},
6        query_builder::{
7            ExprNode, FilterNode, GroupByNode, HavingNode, InnerHashJoinNode,
8            InnerJoinConditionNode, InnerNestedLoopJoinNode, LeftOuterHashJoinNode,
9            LeftOuterJoinConditionNode, LeftOuterNestedLoopJoinNode, LimitNode, OffsetNode,
10            ProjectNode, QueryNode, SelectNode, SelectOrderByNode, SourceNode,
11        },
12        result::Result,
13    },
14};
15
16#[derive(Clone, Debug)]
17pub(super) enum PrevNode<'a> {
18    Select(SelectNode<'a>),
19    Having(HavingNode<'a>),
20    GroupBy(GroupByNode<'a>),
21    Filter(FilterNode<'a>),
22    InnerNestedLoop(Box<InnerNestedLoopJoinNode<'a>>),
23    LeftOuterNestedLoop(Box<LeftOuterNestedLoopJoinNode<'a>>),
24    InnerHash(Box<InnerHashJoinNode<'a>>),
25    LeftOuterHash(Box<LeftOuterHashJoinNode<'a>>),
26    InnerCondition(Box<InnerJoinConditionNode<'a>>),
27    LeftOuterCondition(Box<LeftOuterJoinConditionNode<'a>>),
28    Project(Box<ProjectNode<'a>>),
29    SelectOrderBy(SelectOrderByNode<'a>),
30}
31
32impl PrevNode<'_> {
33    fn build_distinct_input_plan(self) -> Result<DistinctInputPlan> {
34        match self {
35            Self::Select(node) => node.build_project_plan().map(DistinctInputPlan::Project),
36            Self::Having(node) => node.build_project_plan().map(DistinctInputPlan::Project),
37            Self::GroupBy(node) => node.build_project_plan().map(DistinctInputPlan::Project),
38            Self::Filter(node) => node.build_project_plan().map(DistinctInputPlan::Project),
39            Self::InnerNestedLoop(node) => {
40                node.build_project_plan().map(DistinctInputPlan::Project)
41            }
42            Self::LeftOuterNestedLoop(node) => {
43                node.build_project_plan().map(DistinctInputPlan::Project)
44            }
45            Self::InnerHash(node) => node.build_project_plan().map(DistinctInputPlan::Project),
46            Self::LeftOuterHash(node) => node.build_project_plan().map(DistinctInputPlan::Project),
47            Self::InnerCondition(node) => node.build_project_plan().map(DistinctInputPlan::Project),
48            Self::LeftOuterCondition(node) => {
49                node.build_project_plan().map(DistinctInputPlan::Project)
50            }
51            Self::Project(node) => node.build_project_plan().map(DistinctInputPlan::Project),
52            Self::SelectOrderBy(node) => node
53                .build_select_order_by_plan()
54                .map(DistinctInputPlan::SelectOrderBy),
55        }
56    }
57
58    fn build_select(self) -> Result<(Select, Vec<OrderByExpr>)> {
59        let select = match self {
60            Self::Select(node) => node.build_select(),
61            Self::Having(node) => node.build_select(),
62            Self::GroupBy(node) => node.build_select(),
63            Self::Filter(node) => node.build_select(),
64            Self::InnerNestedLoop(node) => node.build_select(),
65            Self::LeftOuterNestedLoop(node) => node.build_select(),
66            Self::InnerHash(node) => node.build_select(),
67            Self::LeftOuterHash(node) => node.build_select(),
68            Self::InnerCondition(node) => node.build_select(),
69            Self::LeftOuterCondition(node) => node.build_select(),
70            Self::Project(node) => node.build_select(),
71            Self::SelectOrderBy(node) => return node.build_select_order_by(),
72        }?;
73
74        Ok((select, Vec::new()))
75    }
76}
77
78macro_rules! impl_from_select_node {
79    ($type: ident, $variant: ident) => {
80        impl<'a> From<$type<'a>> for PrevNode<'a> {
81            fn from(node: $type<'a>) -> Self {
82                Self::$variant(node)
83            }
84        }
85    };
86}
87
88impl_from_select_node!(SelectNode, Select);
89impl_from_select_node!(HavingNode, Having);
90impl_from_select_node!(GroupByNode, GroupBy);
91impl_from_select_node!(FilterNode, Filter);
92impl_from_select_node!(SelectOrderByNode, SelectOrderBy);
93
94impl<'a> From<InnerNestedLoopJoinNode<'a>> for PrevNode<'a> {
95    fn from(node: InnerNestedLoopJoinNode<'a>) -> Self {
96        Self::InnerNestedLoop(Box::new(node))
97    }
98}
99
100impl<'a> From<LeftOuterNestedLoopJoinNode<'a>> for PrevNode<'a> {
101    fn from(node: LeftOuterNestedLoopJoinNode<'a>) -> Self {
102        Self::LeftOuterNestedLoop(Box::new(node))
103    }
104}
105
106impl<'a> From<InnerHashJoinNode<'a>> for PrevNode<'a> {
107    fn from(node: InnerHashJoinNode<'a>) -> Self {
108        Self::InnerHash(Box::new(node))
109    }
110}
111
112impl<'a> From<LeftOuterHashJoinNode<'a>> for PrevNode<'a> {
113    fn from(node: LeftOuterHashJoinNode<'a>) -> Self {
114        Self::LeftOuterHash(Box::new(node))
115    }
116}
117
118impl<'a> From<InnerJoinConditionNode<'a>> for PrevNode<'a> {
119    fn from(node: InnerJoinConditionNode<'a>) -> Self {
120        Self::InnerCondition(Box::new(node))
121    }
122}
123
124impl<'a> From<LeftOuterJoinConditionNode<'a>> for PrevNode<'a> {
125    fn from(node: LeftOuterJoinConditionNode<'a>) -> Self {
126        Self::LeftOuterCondition(Box::new(node))
127    }
128}
129
130impl<'a> From<ProjectNode<'a>> for PrevNode<'a> {
131    fn from(node: ProjectNode<'a>) -> Self {
132        Self::Project(Box::new(node))
133    }
134}
135
136#[derive(Clone, Debug)]
137pub struct DistinctNode<'a> {
138    prev_node: PrevNode<'a>,
139}
140
141impl<'a> DistinctNode<'a> {
142    pub(super) fn new<N: Into<PrevNode<'a>>>(prev_node: N) -> Self {
143        Self {
144            prev_node: prev_node.into(),
145        }
146    }
147
148    pub fn offset<T: Into<ExprNode<'a>>>(self, expr: T) -> OffsetNode<'a> {
149        OffsetNode::new(self, expr)
150    }
151
152    pub fn limit<T: Into<ExprNode<'a>>>(self, expr: T) -> LimitNode<'a> {
153        LimitNode::new(self, expr)
154    }
155
156    pub fn alias_as(self, table_alias: &'a str) -> SourceNode<'a> {
157        QueryNode::DistinctNode(self).alias_as(table_alias)
158    }
159
160    pub(super) fn build_distinct_plan(self) -> Result<DistinctPlan> {
161        self.prev_node
162            .build_distinct_input_plan()
163            .map(|input| DistinctPlan { input })
164    }
165}
166
167impl BuildQueryPlan for DistinctNode<'_> {
168    fn build_query_plan(self) -> Result<QueryPlan> {
169        self.build_distinct_plan().map(QueryPlan::Distinct)
170    }
171}
172
173impl BuildQuery for DistinctNode<'_> {
174    fn build_query(self) -> Result<Query> {
175        let (mut select, order_by) = self.prev_node.build_select()?;
176        select.distinct = true;
177
178        Ok(Query {
179            body: SetExpr::Select(Box::new(select)),
180            order_by,
181            limit: None,
182            offset: None,
183        })
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use crate::{
190        plan::{
191            DistinctInputPlan, DistinctPlan, HashJoinInputPlan, HashJoinPlan, InnerJoinInputPlan,
192            InnerJoinPlan, ProjectInputPlan, ProjectPlan, ProjectionPlan, QueryPlan, SourcePlan,
193            StatementPlan, TableAccessPlan, TableSourcePlan,
194        },
195        query_builder::{
196            Build, QueryBuilderError, SelectItemList, col, select::BuildQuery, table,
197            test_query_builder,
198        },
199        result::Error,
200    };
201    use pretty_assertions::assert_eq;
202
203    #[test]
204    fn distinct_accepts_select_stage_inputs() {
205        let actual = table("Item").select().distinct();
206        let expected = "SELECT DISTINCT * FROM Item";
207        test_query_builder(actual, expected);
208
209        let actual = table("Item").select().filter("id > 0").distinct();
210        let expected = "SELECT DISTINCT * FROM Item WHERE id > 0";
211        test_query_builder(actual, expected);
212
213        let actual = table("Item").select().group_by("name").distinct();
214        let expected = "SELECT DISTINCT * FROM Item GROUP BY name";
215        test_query_builder(actual, expected);
216
217        let actual = table("Item")
218            .select()
219            .group_by("name")
220            .having("COUNT(*) > 0")
221            .distinct();
222        let expected = "SELECT DISTINCT * FROM Item GROUP BY name HAVING COUNT(*) > 0";
223        test_query_builder(actual, expected);
224
225        let actual = table("Item").select().join("Category").distinct();
226        let expected = "SELECT DISTINCT * FROM Item JOIN Category";
227        test_query_builder(actual, expected);
228
229        let actual = table("Item")
230            .select()
231            .join("Category")
232            .on("Item.category_id = Category.id")
233            .distinct();
234        let expected =
235            "SELECT DISTINCT * FROM Item JOIN Category ON Item.category_id = Category.id";
236        test_query_builder(actual, expected);
237
238        let actual = table("Item").select().project("name").distinct();
239        let expected = "SELECT DISTINCT name FROM Item";
240        test_query_builder(actual, expected);
241
242        let actual = table("Item")
243            .select()
244            .project("name")
245            .order_by("name")
246            .distinct();
247        let expected = "SELECT DISTINCT name FROM Item ORDER BY name";
248        test_query_builder(actual, expected);
249    }
250
251    #[test]
252    fn distinct_connects_to_terminal_stages() {
253        let actual = table("Item").select().distinct().offset(1);
254        let expected = "SELECT DISTINCT * FROM Item OFFSET 1";
255        test_query_builder(actual, expected);
256
257        let actual = table("Item").select().distinct().limit(1);
258        let expected = "SELECT DISTINCT * FROM Item LIMIT 1";
259        test_query_builder(actual, expected);
260
261        let actual = table("Item")
262            .select()
263            .distinct()
264            .alias_as("DistinctItem")
265            .select();
266        let expected = "SELECT * FROM (SELECT DISTINCT * FROM Item) DistinctItem";
267        test_query_builder(actual, expected);
268    }
269
270    #[test]
271    fn distinct_preserves_plan_only_hash_join() {
272        let actual = table("Item")
273            .select()
274            .join("Category")
275            .hash_executor("Item.category_id", "Category.id")
276            .distinct()
277            .build_query();
278        let expected = Err(Error::QueryBuilder(
279            QueryBuilderError::HashJoinExecutorRequiresPlan,
280        ));
281        assert_eq!(actual, expected);
282
283        let actual = table("Item")
284            .select()
285            .join("Category")
286            .hash_executor("Item.category_id", "Category.id")
287            .distinct()
288            .build();
289        let expected = {
290            let join = InnerJoinPlan {
291                input: InnerJoinInputPlan::Hash(HashJoinPlan {
292                    input: HashJoinInputPlan::Source(SourcePlan::Table(TableSourcePlan {
293                        name: "Item".to_owned(),
294                        alias: None,
295                        access: TableAccessPlan::FullScan,
296                    })),
297                    right: SourcePlan::Table(TableSourcePlan {
298                        name: "Category".to_owned(),
299                        alias: None,
300                        access: TableAccessPlan::FullScan,
301                    }),
302                    input_key: col("Category.id").build_expr_plan().unwrap(),
303                    right_key: col("Item.category_id").build_expr_plan().unwrap(),
304                    right_filter: None,
305                }),
306            };
307            let project = ProjectPlan {
308                input: ProjectInputPlan::InnerJoin(Box::new(join)),
309                projection: ProjectionPlan::SelectItems(
310                    SelectItemList::from("*").build_select_items_plan().unwrap(),
311                ),
312            };
313
314            Ok(StatementPlan::Query(QueryPlan::Distinct(DistinctPlan {
315                input: DistinctInputPlan::Project(project),
316            })))
317        };
318        assert_eq!(actual, expected);
319    }
320}