datafusion_functions_window/
planner.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//! SQL planning extensions like [`WindowFunctionPlanner`]
19
20use datafusion_common::Result;
21use datafusion_expr::{
22    expr::{WindowFunction, WindowFunctionParams},
23    expr_rewriter::NamePreserver,
24    planner::{ExprPlanner, PlannerResult, RawWindowExpr},
25    utils::COUNT_STAR_EXPANSION,
26    Expr,
27};
28
29#[derive(Debug)]
30pub struct WindowFunctionPlanner;
31
32impl ExprPlanner for WindowFunctionPlanner {
33    fn plan_window(
34        &self,
35        raw_expr: RawWindowExpr,
36    ) -> Result<PlannerResult<RawWindowExpr>> {
37        let RawWindowExpr {
38            func_def,
39            args,
40            partition_by,
41            order_by,
42            window_frame,
43            filter,
44            null_treatment,
45            distinct,
46        } = raw_expr;
47
48        let origin_expr = Expr::from(WindowFunction {
49            fun: func_def,
50            params: WindowFunctionParams {
51                args,
52                partition_by,
53                order_by,
54                window_frame,
55                filter,
56                null_treatment,
57                distinct,
58            },
59        });
60
61        let saved_name = NamePreserver::new_for_projection().save(&origin_expr);
62
63        let Expr::WindowFunction(window_fun) = origin_expr else {
64            unreachable!("")
65        };
66        let WindowFunction {
67            fun,
68            params:
69                WindowFunctionParams {
70                    args,
71                    partition_by,
72                    order_by,
73                    window_frame,
74                    null_treatment,
75                    distinct,
76                    filter,
77                },
78        } = *window_fun;
79        let raw_expr = RawWindowExpr {
80            func_def: fun,
81            args,
82            partition_by,
83            order_by,
84            window_frame,
85            filter,
86            null_treatment,
87            distinct,
88        };
89
90        // TODO: remove the next line after `Expr::Wildcard` is removed
91        #[expect(deprecated)]
92        if raw_expr.func_def.name() == "count"
93            && (raw_expr.args.len() == 1
94                && matches!(raw_expr.args[0], Expr::Wildcard { .. })
95                || raw_expr.args.is_empty())
96        {
97            let RawWindowExpr {
98                func_def,
99                args: _,
100                partition_by,
101                order_by,
102                window_frame,
103                filter,
104                null_treatment,
105                distinct,
106            } = raw_expr;
107
108            let new_expr = Expr::from(WindowFunction {
109                fun: func_def,
110                params: WindowFunctionParams {
111                    args: vec![Expr::Literal(COUNT_STAR_EXPANSION, None)],
112                    partition_by,
113                    order_by,
114                    window_frame,
115                    filter,
116                    null_treatment,
117                    distinct,
118                },
119            });
120            let new_expr = saved_name.restore(new_expr);
121
122            return Ok(PlannerResult::Planned(new_expr));
123        }
124
125        Ok(PlannerResult::Original(raw_expr))
126    }
127}
datafusion_functions_window/planner.rs

datafusion_functions_window/
planner.rs
