datafusion_python/expr/
window.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
18use datafusion::common::{DataFusionError, ScalarValue};
19use datafusion::logical_expr::expr::WindowFunction;
20use datafusion::logical_expr::{Expr, Window, WindowFrame, WindowFrameBound, WindowFrameUnits};
21use pyo3::prelude::*;
22use std::fmt::{self, Display, Formatter};
23
24use crate::common::data_type::PyScalarValue;
25use crate::common::df_schema::PyDFSchema;
26use crate::errors::{py_type_err, PyDataFusionResult};
27use crate::expr::logical_node::LogicalNode;
28use crate::expr::sort_expr::{py_sort_expr_list, PySortExpr};
29use crate::expr::PyExpr;
30use crate::sql::logical::PyLogicalPlan;
31
32use super::py_expr_list;
33
34use crate::errors::py_datafusion_err;
35
36#[pyclass(name = "WindowExpr", module = "datafusion.expr", subclass)]
37#[derive(Clone)]
38pub struct PyWindowExpr {
39    window: Window,
40}
41
42#[pyclass(name = "WindowFrame", module = "datafusion.expr", subclass)]
43#[derive(Clone)]
44pub struct PyWindowFrame {
45    window_frame: WindowFrame,
46}
47
48impl From<PyWindowFrame> for WindowFrame {
49    fn from(window_frame: PyWindowFrame) -> Self {
50        window_frame.window_frame
51    }
52}
53
54impl From<WindowFrame> for PyWindowFrame {
55    fn from(window_frame: WindowFrame) -> PyWindowFrame {
56        PyWindowFrame { window_frame }
57    }
58}
59
60#[pyclass(name = "WindowFrameBound", module = "datafusion.expr", subclass)]
61#[derive(Clone)]
62pub struct PyWindowFrameBound {
63    frame_bound: WindowFrameBound,
64}
65
66impl From<PyWindowExpr> for Window {
67    fn from(window: PyWindowExpr) -> Window {
68        window.window
69    }
70}
71
72impl From<Window> for PyWindowExpr {
73    fn from(window: Window) -> PyWindowExpr {
74        PyWindowExpr { window }
75    }
76}
77
78impl From<WindowFrameBound> for PyWindowFrameBound {
79    fn from(frame_bound: WindowFrameBound) -> Self {
80        PyWindowFrameBound { frame_bound }
81    }
82}
83
84impl Display for PyWindowExpr {
85    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
86        write!(
87            f,
88            "Over\n
89            Window Expr: {:?}
90            Schema: {:?}",
91            &self.window.window_expr, &self.window.schema
92        )
93    }
94}
95
96impl Display for PyWindowFrame {
97    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
98        write!(
99            f,
100            "OVER ({} BETWEEN {} AND {})",
101            self.window_frame.units, self.window_frame.start_bound, self.window_frame.end_bound
102        )
103    }
104}
105
106#[pymethods]
107impl PyWindowExpr {
108    /// Returns the schema of the Window
109    pub fn schema(&self) -> PyResult<PyDFSchema> {
110        Ok(self.window.schema.as_ref().clone().into())
111    }
112
113    /// Returns window expressions
114    pub fn get_window_expr(&self) -> PyResult<Vec<PyExpr>> {
115        py_expr_list(&self.window.window_expr)
116    }
117
118    /// Returns order by columns in a window function expression
119    pub fn get_sort_exprs(&self, expr: PyExpr) -> PyResult<Vec<PySortExpr>> {
120        match expr.expr.unalias() {
121            Expr::WindowFunction(WindowFunction { order_by, .. }) => py_sort_expr_list(&order_by),
122            other => Err(not_window_function_err(other)),
123        }
124    }
125
126    /// Return partition by columns in a window function expression
127    pub fn get_partition_exprs(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>> {
128        match expr.expr.unalias() {
129            Expr::WindowFunction(WindowFunction { partition_by, .. }) => {
130                py_expr_list(&partition_by)
131            }
132            other => Err(not_window_function_err(other)),
133        }
134    }
135
136    /// Return input args for window function
137    pub fn get_args(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>> {
138        match expr.expr.unalias() {
139            Expr::WindowFunction(WindowFunction { args, .. }) => py_expr_list(&args),
140            other => Err(not_window_function_err(other)),
141        }
142    }
143
144    /// Return window function name
145    pub fn window_func_name(&self, expr: PyExpr) -> PyResult<String> {
146        match expr.expr.unalias() {
147            Expr::WindowFunction(WindowFunction { fun, .. }) => Ok(fun.to_string()),
148            other => Err(not_window_function_err(other)),
149        }
150    }
151
152    /// Returns a Pywindow frame for a given window function expression
153    pub fn get_frame(&self, expr: PyExpr) -> Option<PyWindowFrame> {
154        match expr.expr.unalias() {
155            Expr::WindowFunction(WindowFunction { window_frame, .. }) => Some(window_frame.into()),
156            _ => None,
157        }
158    }
159}
160
161fn not_window_function_err(expr: Expr) -> PyErr {
162    py_type_err(format!(
163        "Provided {} Expr {:?} is not a WindowFunction type",
164        expr.variant_name(),
165        expr
166    ))
167}
168
169#[pymethods]
170impl PyWindowFrame {
171    #[new]
172    #[pyo3(signature=(unit, start_bound, end_bound))]
173    pub fn new(
174        unit: &str,
175        start_bound: Option<PyScalarValue>,
176        end_bound: Option<PyScalarValue>,
177    ) -> PyResult<Self> {
178        let units = unit.to_ascii_lowercase();
179        let units = match units.as_str() {
180            "rows" => WindowFrameUnits::Rows,
181            "range" => WindowFrameUnits::Range,
182            "groups" => WindowFrameUnits::Groups,
183            _ => {
184                return Err(py_datafusion_err(DataFusionError::NotImplemented(format!(
185                    "{:?}",
186                    units,
187                ))));
188            }
189        };
190        let start_bound = match start_bound {
191            Some(start_bound) => WindowFrameBound::Preceding(start_bound.0),
192            None => match units {
193                WindowFrameUnits::Range => WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
194                WindowFrameUnits::Rows => WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
195                WindowFrameUnits::Groups => {
196                    return Err(py_datafusion_err(DataFusionError::NotImplemented(format!(
197                        "{:?}",
198                        units,
199                    ))));
200                }
201            },
202        };
203        let end_bound = match end_bound {
204            Some(end_bound) => WindowFrameBound::Following(end_bound.0),
205            None => match units {
206                WindowFrameUnits::Rows => WindowFrameBound::Following(ScalarValue::UInt64(None)),
207                WindowFrameUnits::Range => WindowFrameBound::Following(ScalarValue::UInt64(None)),
208                WindowFrameUnits::Groups => {
209                    return Err(py_datafusion_err(DataFusionError::NotImplemented(format!(
210                        "{:?}",
211                        units,
212                    ))));
213                }
214            },
215        };
216        Ok(PyWindowFrame {
217            window_frame: WindowFrame::new_bounds(units, start_bound, end_bound),
218        })
219    }
220
221    /// Returns the window frame units for the bounds
222    pub fn get_frame_units(&self) -> PyResult<String> {
223        Ok(self.window_frame.units.to_string())
224    }
225    /// Returns starting bound
226    pub fn get_lower_bound(&self) -> PyResult<PyWindowFrameBound> {
227        Ok(self.window_frame.start_bound.clone().into())
228    }
229    /// Returns end bound
230    pub fn get_upper_bound(&self) -> PyResult<PyWindowFrameBound> {
231        Ok(self.window_frame.end_bound.clone().into())
232    }
233
234    /// Get a String representation of this window frame
235    fn __repr__(&self) -> String {
236        format!("{}", self)
237    }
238}
239
240#[pymethods]
241impl PyWindowFrameBound {
242    /// Returns if the frame bound is current row
243    pub fn is_current_row(&self) -> bool {
244        matches!(self.frame_bound, WindowFrameBound::CurrentRow)
245    }
246
247    /// Returns if the frame bound is preceding
248    pub fn is_preceding(&self) -> bool {
249        matches!(self.frame_bound, WindowFrameBound::Preceding(_))
250    }
251
252    /// Returns if the frame bound is following
253    pub fn is_following(&self) -> bool {
254        matches!(self.frame_bound, WindowFrameBound::Following(_))
255    }
256    /// Returns the offset of the window frame
257    pub fn get_offset(&self) -> PyDataFusionResult<Option<u64>> {
258        match &self.frame_bound {
259            WindowFrameBound::Preceding(val) | WindowFrameBound::Following(val) => match val {
260                x if x.is_null() => Ok(None),
261                ScalarValue::UInt64(v) => Ok(*v),
262                // The cast below is only safe because window bounds cannot be negative
263                ScalarValue::Int64(v) => Ok(v.map(|n| n as u64)),
264                ScalarValue::Utf8(Some(s)) => match s.parse::<u64>() {
265                    Ok(s) => Ok(Some(s)),
266                    Err(_e) => Err(DataFusionError::Plan(format!(
267                        "Unable to parse u64 from Utf8 value '{s}'"
268                    ))
269                    .into()),
270                },
271                ref x => {
272                    Err(DataFusionError::Plan(format!("Unexpected window frame bound: {x}")).into())
273                }
274            },
275            WindowFrameBound::CurrentRow => Ok(None),
276        }
277    }
278    /// Returns if the frame bound is unbounded
279    pub fn is_unbounded(&self) -> PyResult<bool> {
280        match &self.frame_bound {
281            WindowFrameBound::Preceding(v) | WindowFrameBound::Following(v) => Ok(v.is_null()),
282            WindowFrameBound::CurrentRow => Ok(false),
283        }
284    }
285}
286
287impl LogicalNode for PyWindowExpr {
288    fn inputs(&self) -> Vec<PyLogicalPlan> {
289        vec![self.window.input.as_ref().clone().into()]
290    }
291
292    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
293        Ok(self.clone().into_py(py))
294    }
295}
datafusion_python/expr/window.rs

datafusion_python/expr/
window.rs
