gluesql_core/query_builder/select/
limit.rs

1use {
2    super::{BuildProjectPlan, BuildQuery, BuildQueryPlan, DistinctNode, values::ValuesNode},
3    crate::{
4        ast::Query,
5        plan::{LimitInputPlan, LimitPlan, QueryPlan},
6        query_builder::{
7            ExprNode, FilterNode, GroupByNode, HavingNode, InnerHashJoinNode,
8            InnerJoinConditionNode, InnerNestedLoopJoinNode, LeftOuterHashJoinNode,
9            LeftOuterJoinConditionNode, LeftOuterNestedLoopJoinNode, ProjectNode, QueryNode,
10            SelectNode, SelectOrderByNode, SourceNode, ValuesOrderByNode,
11        },
12        result::Result,
13    },
14};
15
16#[derive(Clone, Debug)]
17pub(super) enum PrevNode<'a> {
18    Select(SelectNode<'a>),
19    Values(ValuesNode<'a>),
20    GroupBy(GroupByNode<'a>),
21    Having(HavingNode<'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    Filter(FilterNode<'a>),
29    SelectOrderBy(SelectOrderByNode<'a>),
30    ValuesOrderBy(ValuesOrderByNode<'a>),
31    Distinct(DistinctNode<'a>),
32    ProjectNode(Box<ProjectNode<'a>>),
33}
34
35impl PrevNode<'_> {
36    fn build_limit_input_plan(self) -> Result<LimitInputPlan> {
37        match self {
38            Self::Select(node) => node.build_project_plan().map(LimitInputPlan::Project),
39            Self::Values(node) => node.build_values_plan().map(LimitInputPlan::Values),
40            Self::GroupBy(node) => node.build_project_plan().map(LimitInputPlan::Project),
41            Self::Having(node) => node.build_project_plan().map(LimitInputPlan::Project),
42            Self::InnerNestedLoop(node) => node.build_project_plan().map(LimitInputPlan::Project),
43            Self::LeftOuterNestedLoop(node) => {
44                node.build_project_plan().map(LimitInputPlan::Project)
45            }
46            Self::InnerHash(node) => node.build_project_plan().map(LimitInputPlan::Project),
47            Self::LeftOuterHash(node) => node.build_project_plan().map(LimitInputPlan::Project),
48            Self::InnerCondition(node) => node.build_project_plan().map(LimitInputPlan::Project),
49            Self::LeftOuterCondition(node) => {
50                node.build_project_plan().map(LimitInputPlan::Project)
51            }
52            Self::Filter(node) => node.build_project_plan().map(LimitInputPlan::Project),
53            Self::SelectOrderBy(node) => node
54                .build_select_order_by_plan()
55                .map(LimitInputPlan::SelectOrderBy),
56            Self::ValuesOrderBy(node) => node
57                .build_values_order_by_plan()
58                .map(LimitInputPlan::ValuesOrderBy),
59            Self::Distinct(node) => node.build_distinct_plan().map(LimitInputPlan::Distinct),
60            Self::ProjectNode(node) => node.build_project_plan().map(LimitInputPlan::Project),
61        }
62    }
63}
64
65impl BuildQuery for PrevNode<'_> {
66    fn build_query(self) -> Result<Query> {
67        match self {
68            Self::Select(node) => node.build_query(),
69            Self::Values(node) => node.build_query(),
70            Self::GroupBy(node) => node.build_query(),
71            Self::Having(node) => node.build_query(),
72            Self::InnerNestedLoop(node) => node.build_query(),
73            Self::LeftOuterNestedLoop(node) => node.build_query(),
74            Self::InnerHash(node) => node.build_query(),
75            Self::LeftOuterHash(node) => node.build_query(),
76            Self::InnerCondition(node) => node.build_query(),
77            Self::LeftOuterCondition(node) => node.build_query(),
78            Self::Filter(node) => node.build_query(),
79            Self::SelectOrderBy(node) => node.build_query(),
80            Self::ValuesOrderBy(node) => node.build_query(),
81            Self::Distinct(node) => node.build_query(),
82            Self::ProjectNode(node) => node.build_query(),
83        }
84    }
85}
86
87impl<'a> From<SelectNode<'a>> for PrevNode<'a> {
88    fn from(node: SelectNode<'a>) -> Self {
89        PrevNode::Select(node)
90    }
91}
92
93impl<'a> From<ValuesNode<'a>> for PrevNode<'a> {
94    fn from(node: ValuesNode<'a>) -> Self {
95        PrevNode::Values(node)
96    }
97}
98
99impl<'a> From<GroupByNode<'a>> for PrevNode<'a> {
100    fn from(node: GroupByNode<'a>) -> Self {
101        PrevNode::GroupBy(node)
102    }
103}
104
105impl<'a> From<HavingNode<'a>> for PrevNode<'a> {
106    fn from(node: HavingNode<'a>) -> Self {
107        PrevNode::Having(node)
108    }
109}
110
111impl<'a> From<InnerNestedLoopJoinNode<'a>> for PrevNode<'a> {
112    fn from(node: InnerNestedLoopJoinNode<'a>) -> Self {
113        Self::InnerNestedLoop(Box::new(node))
114    }
115}
116
117impl<'a> From<LeftOuterNestedLoopJoinNode<'a>> for PrevNode<'a> {
118    fn from(node: LeftOuterNestedLoopJoinNode<'a>) -> Self {
119        Self::LeftOuterNestedLoop(Box::new(node))
120    }
121}
122
123impl<'a> From<InnerHashJoinNode<'a>> for PrevNode<'a> {
124    fn from(node: InnerHashJoinNode<'a>) -> Self {
125        Self::InnerHash(Box::new(node))
126    }
127}
128
129impl<'a> From<LeftOuterHashJoinNode<'a>> for PrevNode<'a> {
130    fn from(node: LeftOuterHashJoinNode<'a>) -> Self {
131        Self::LeftOuterHash(Box::new(node))
132    }
133}
134
135impl<'a> From<InnerJoinConditionNode<'a>> for PrevNode<'a> {
136    fn from(node: InnerJoinConditionNode<'a>) -> Self {
137        Self::InnerCondition(Box::new(node))
138    }
139}
140
141impl<'a> From<LeftOuterJoinConditionNode<'a>> for PrevNode<'a> {
142    fn from(node: LeftOuterJoinConditionNode<'a>) -> Self {
143        Self::LeftOuterCondition(Box::new(node))
144    }
145}
146
147impl<'a> From<FilterNode<'a>> for PrevNode<'a> {
148    fn from(node: FilterNode<'a>) -> Self {
149        PrevNode::Filter(node)
150    }
151}
152
153impl<'a> From<SelectOrderByNode<'a>> for PrevNode<'a> {
154    fn from(node: SelectOrderByNode<'a>) -> Self {
155        Self::SelectOrderBy(node)
156    }
157}
158
159impl<'a> From<ValuesOrderByNode<'a>> for PrevNode<'a> {
160    fn from(node: ValuesOrderByNode<'a>) -> Self {
161        Self::ValuesOrderBy(node)
162    }
163}
164
165impl<'a> From<DistinctNode<'a>> for PrevNode<'a> {
166    fn from(node: DistinctNode<'a>) -> Self {
167        Self::Distinct(node)
168    }
169}
170
171impl<'a> From<ProjectNode<'a>> for PrevNode<'a> {
172    fn from(node: ProjectNode<'a>) -> Self {
173        PrevNode::ProjectNode(Box::new(node))
174    }
175}
176
177#[derive(Clone, Debug)]
178pub struct LimitNode<'a> {
179    prev_node: PrevNode<'a>,
180    expr: ExprNode<'a>,
181}
182
183impl<'a> LimitNode<'a> {
184    pub(super) fn new<N: Into<PrevNode<'a>>, T: Into<ExprNode<'a>>>(prev_node: N, expr: T) -> Self {
185        Self {
186            prev_node: prev_node.into(),
187            expr: expr.into(),
188        }
189    }
190
191    pub fn alias_as(self, table_alias: &'a str) -> SourceNode<'a> {
192        QueryNode::LimitNode(self).alias_as(table_alias)
193    }
194}
195
196impl BuildQueryPlan for LimitNode<'_> {
197    fn build_query_plan(self) -> Result<QueryPlan> {
198        let count = self.expr.build_expr_plan()?;
199        self.prev_node
200            .build_limit_input_plan()
201            .map(|input| QueryPlan::Limit(LimitPlan { input, count }))
202    }
203}
204
205impl BuildQuery for LimitNode<'_> {
206    fn build_query(self) -> Result<Query> {
207        let mut node_data = self.prev_node.build_query()?;
208        node_data.limit = Some(self.expr.build_expr()?);
209
210        Ok(node_data)
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use {
217        crate::{
218            plan::{
219                HashJoinInputPlan, HashJoinPlan, InnerJoinInputPlan, InnerJoinPlan, LimitInputPlan,
220                LimitPlan, ProjectInputPlan, ProjectPlan, ProjectionPlan, QueryPlan, SourcePlan,
221                StatementPlan, TableAccessPlan, TableSourcePlan,
222            },
223            query_builder::{Build, SelectItemList, col, num, table, test_query_builder},
224        },
225        pretty_assertions::assert_eq,
226    };
227
228    #[test]
229    fn limit() {
230        // select node -> limit node -> build
231        let actual = table("Foo").select().limit(10);
232        let expected = "SELECT * FROM Foo LIMIT 10";
233        test_query_builder(actual, expected);
234
235        // group by node -> limit node -> build
236        let actual = table("Foo").select().group_by("bar").limit(10);
237        let expected = "SELECT * FROM Foo GROUP BY bar LIMIT 10";
238        test_query_builder(actual, expected);
239
240        // having node -> limit node -> build
241        let actual = table("Foo")
242            .select()
243            .group_by("bar")
244            .having("bar = 10")
245            .limit(10);
246        let expected = "SELECT * FROM Foo GROUP BY bar HAVING bar = 10 LIMIT 10";
247        test_query_builder(actual, expected);
248
249        // inner nested loop join node -> limit node -> build
250        let actual = table("Foo").select().join("Bar").limit(10);
251        let expected = "SELECT * FROM Foo JOIN Bar LIMIT 10";
252        test_query_builder(actual, expected);
253
254        // inner nested loop join node -> limit node -> build
255        let actual = table("Foo").select().join_as("Bar", "B").limit(10);
256        let expected = "SELECT * FROM Foo JOIN Bar AS B LIMIT 10";
257        test_query_builder(actual, expected);
258
259        // left outer nested loop join node -> limit node -> build
260        let actual = table("Foo").select().left_join("Bar").limit(10);
261        let expected = "SELECT * FROM Foo LEFT JOIN Bar LIMIT 10";
262        test_query_builder(actual, expected);
263
264        // left outer nested loop join node -> limit node -> build
265        let actual = table("Foo").select().left_join_as("Bar", "B").limit(10);
266        let expected = "SELECT * FROM Foo LEFT JOIN Bar AS B LIMIT 10";
267        test_query_builder(actual, expected);
268
269        // group by node -> limit node -> build
270        let actual = table("Foo").select().group_by("id").limit(10);
271        let expected = "SELECT * FROM Foo GROUP BY id LIMIT 10";
272        test_query_builder(actual, expected);
273
274        // having node -> limit node -> build
275        let actual = table("Foo")
276            .select()
277            .group_by("id")
278            .having(col("id").gt(10))
279            .limit(10);
280        let expected = "SELECT * FROM Foo GROUP BY id HAVING id > 10 LIMIT 10";
281        test_query_builder(actual, expected);
282
283        // inner join condition node -> limit node -> build
284        let actual = table("Foo")
285            .select()
286            .join("Bar")
287            .on("Foo.id = Bar.id")
288            .limit(10);
289        let expected = "SELECT * FROM Foo JOIN Bar ON Foo.id = Bar.id LIMIT 10";
290        test_query_builder(actual, expected);
291
292        // filter node -> limit node -> build
293        let actual = table("World").select().filter(col("id").gt(2)).limit(100);
294        let expected = "SELECT * FROM World WHERE id > 2 LIMIT 100";
295        test_query_builder(actual, expected);
296
297        // order by node -> limit node -> build
298        let actual = table("Hello").select().order_by("score").limit(3);
299        let expected = "SELECT * FROM Hello ORDER BY score LIMIT 3";
300        test_query_builder(actual, expected);
301
302        // project node -> limit node -> build
303        let actual = table("Item").select().project("*").limit(10);
304        let expected = "SELECT * FROM Item LIMIT 10";
305        test_query_builder(actual, expected);
306
307        // inner hash join node -> limit node -> build
308        let actual = table("Player")
309            .select()
310            .join("PlayerItem")
311            .hash_executor("PlayerItem.user_id", "Player.id")
312            .limit(100)
313            .build();
314        let expected = {
315            let join = InnerJoinPlan {
316                input: InnerJoinInputPlan::Hash(HashJoinPlan {
317                    input: HashJoinInputPlan::Source(SourcePlan::Table(TableSourcePlan {
318                        name: "Player".to_owned(),
319                        alias: None,
320                        access: TableAccessPlan::FullScan,
321                    })),
322                    right: SourcePlan::Table(TableSourcePlan {
323                        name: "PlayerItem".to_owned(),
324                        alias: None,
325                        access: TableAccessPlan::FullScan,
326                    }),
327                    input_key: col("Player.id").build_expr_plan().unwrap(),
328                    right_key: col("PlayerItem.user_id").build_expr_plan().unwrap(),
329                    right_filter: None,
330                }),
331            };
332            let project = ProjectPlan {
333                input: ProjectInputPlan::InnerJoin(Box::new(join)),
334                projection: ProjectionPlan::SelectItems(
335                    SelectItemList::from("*").build_select_items_plan().unwrap(),
336                ),
337            };
338
339            let limit = LimitPlan {
340                input: LimitInputPlan::Project(project),
341                count: num(100).build_expr_plan().unwrap(),
342            };
343
344            Ok(StatementPlan::Query(QueryPlan::Limit(limit)))
345        };
346        assert_eq!(actual, expected);
347
348        // select node -> limit node -> derived subquery
349        let actual = table("Foo").select().limit(10).alias_as("Sub").select();
350        let expected = "SELECT * FROM (SELECT * FROM Foo LIMIT 10) Sub";
351        test_query_builder(actual, expected);
352    }
353}
gluesql_core/query_builder/select/limit.rs

gluesql_core/query_builder/select/
limit.rs
