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, ExprFunctionExt,
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            null_treatment,
44        } = raw_expr;
45
46        let origin_expr = Expr::WindowFunction(WindowFunction {
47            fun: func_def,
48            params: WindowFunctionParams {
49                args,
50                partition_by,
51                order_by,
52                window_frame,
53                null_treatment,
54            },
55        });
56
57        let saved_name = NamePreserver::new_for_projection().save(&origin_expr);
58
59        let Expr::WindowFunction(WindowFunction {
60            fun,
61            params:
62                WindowFunctionParams {
63                    args,
64                    partition_by,
65                    order_by,
66                    window_frame,
67                    null_treatment,
68                },
69        }) = origin_expr
70        else {
71            unreachable!("")
72        };
73        let raw_expr = RawWindowExpr {
74            func_def: fun,
75            args,
76            partition_by,
77            order_by,
78            window_frame,
79            null_treatment,
80        };
81
82        // TODO: remove the next line after `Expr::Wildcard` is removed
83        #[expect(deprecated)]
84        if raw_expr.func_def.name() == "count"
85            && (raw_expr.args.len() == 1
86                && matches!(raw_expr.args[0], Expr::Wildcard { .. })
87                || raw_expr.args.is_empty())
88        {
89            let RawWindowExpr {
90                func_def,
91                args: _,
92                partition_by,
93                order_by,
94                window_frame,
95                null_treatment,
96            } = raw_expr;
97
98            let new_expr = Expr::WindowFunction(WindowFunction::new(
99                func_def,
100                vec![Expr::Literal(COUNT_STAR_EXPANSION)],
101            ))
102            .partition_by(partition_by)
103            .order_by(order_by)
104            .window_frame(window_frame)
105            .null_treatment(null_treatment)
106            .build()?;
107
108            let new_expr = saved_name.restore(new_expr);
109
110            return Ok(PlannerResult::Planned(new_expr));
111        }
112
113        Ok(PlannerResult::Original(raw_expr))
114    }
115}
datafusion_functions_window/planner.rs

datafusion_functions_window/
planner.rs
