Skip to main content

drizzle_core/expr/
window.rs

1//! Window functions and OVER clause support.
2//!
3//! Provides:
4//! - `WindowSpec` builder for PARTITION BY, ORDER BY, and frame clauses
5//! - `.over()` method on aggregate `SQLExpr` to convert Agg → Scalar
6//! - Pure window functions: `row_number`, `rank`, `dense_rank`, `ntile`,
7//!   `percent_rank`, `cume_dist`, `lag`, `lead`, `first_value`, `last_value`,
8//!   `nth_value`
9//!
10//! # Example
11//!
12//! ```rust
13//! # let _ = r####"
14//! use drizzle_core::expr::*;
15//!
16//! // Aggregate as window function
17//! count_all().over(window().partition_by([users.dept]))
18//! // → SQLExpr<CountType, NonNull, Scalar>
19//!
20//! // Pure window function
21//! row_number().over(window().order_by([asc(users.id)]))
22//! // → SQLExpr<CountType, NonNull, Scalar>
23//! # "####;
24//! ```
25
26use core::marker::PhantomData;
27
28use crate::sql::{SQL, Token};
29use crate::traits::{SQLParam, ToSQL};
30use crate::types::{BooleanLike, Compatible, DataType};
31
32use super::agg::{CountPolicy, FloatPolicy};
33use super::null::NullOr;
34use super::{Agg, Expr, NonNull, Null, Nullability, SQLExpr, Scalar};
35
36// =============================================================================
37// Frame Bounds
38// =============================================================================
39
40/// Specifies a bound for a window frame (ROWS/RANGE BETWEEN).
41#[derive(Debug, Clone, Copy)]
42pub enum FrameBound {
43    /// UNBOUNDED PRECEDING
44    UnboundedPreceding,
45    /// N PRECEDING
46    Preceding(u64),
47    /// CURRENT ROW
48    CurrentRow,
49    /// N FOLLOWING
50    Following(u64),
51    /// UNBOUNDED FOLLOWING
52    UnboundedFollowing,
53}
54
55impl FrameBound {
56    fn write_sql<'a, V: SQLParam>(&self) -> SQL<'a, V> {
57        match self {
58            Self::UnboundedPreceding => SQL::from(Token::UNBOUNDED).push(Token::PRECEDING),
59            Self::Preceding(n) => {
60                SQL::number(usize::try_from(*n).unwrap_or(usize::MAX)).push(Token::PRECEDING)
61            }
62            Self::CurrentRow => SQL::from(Token::CURRENT).push(Token::ROW),
63            Self::Following(n) => {
64                SQL::number(usize::try_from(*n).unwrap_or(usize::MAX)).push(Token::FOLLOWING)
65            }
66            Self::UnboundedFollowing => SQL::from(Token::UNBOUNDED).push(Token::FOLLOWING),
67        }
68    }
69}
70
71// =============================================================================
72// WindowSpec
73// =============================================================================
74
75/// Builder for a window specification (the content inside `OVER (...)`).
76///
77/// # Example
78///
79/// ```rust
80/// # let _ = r####"
81/// window()
82///     .partition_by([users.dept])
83///     .order_by([asc(users.salary)])
84///     .rows_between(FrameBound::UnboundedPreceding, FrameBound::CurrentRow)
85/// # "####;
86/// ```
87#[derive(Debug, Clone)]
88pub struct WindowSpec<'a, V: SQLParam> {
89    partition: Option<SQL<'a, V>>,
90    order: Option<SQL<'a, V>>,
91    frame: Option<SQL<'a, V>>,
92}
93
94/// Create an empty window specification.
95#[must_use]
96pub const fn window<'a, V: SQLParam>() -> WindowSpec<'a, V> {
97    WindowSpec {
98        partition: None,
99        order: None,
100        frame: None,
101    }
102}
103
104impl<'a, V: SQLParam + 'a> WindowSpec<'a, V> {
105    /// Set the PARTITION BY clause.
106    #[must_use]
107    pub fn partition_by<I>(mut self, exprs: I) -> Self
108    where
109        I: IntoIterator,
110        I::Item: ToSQL<'a, V>,
111    {
112        self.partition = Some(
113            SQL::from(Token::PARTITION)
114                .push(Token::BY)
115                .append(SQL::join(exprs, Token::COMMA)),
116        );
117        self
118    }
119
120    /// Set the ORDER BY clause.
121    #[must_use]
122    pub fn order_by<T: ToSQL<'a, V>>(mut self, exprs: T) -> Self {
123        self.order = Some(
124            SQL::from(Token::ORDER)
125                .push(Token::BY)
126                .append(exprs.into_sql()),
127        );
128        self
129    }
130
131    /// Set a ROWS frame specification.
132    #[must_use]
133    pub fn rows_between(mut self, start: FrameBound, end: FrameBound) -> Self {
134        self.frame = Some(
135            SQL::from(Token::ROWS)
136                .push(Token::BETWEEN)
137                .append(start.write_sql())
138                .push(Token::AND)
139                .append(end.write_sql()),
140        );
141        self
142    }
143
144    /// Set a RANGE frame specification.
145    #[must_use]
146    pub fn range_between(mut self, start: FrameBound, end: FrameBound) -> Self {
147        self.frame = Some(
148            SQL::from(Token::RANGE)
149                .push(Token::BETWEEN)
150                .append(start.write_sql())
151                .push(Token::AND)
152                .append(end.write_sql()),
153        );
154        self
155    }
156
157    /// Build the window spec into SQL (contents inside the OVER parentheses).
158    fn into_sql(self) -> SQL<'a, V> {
159        let mut sql = SQL::empty();
160        if let Some(p) = self.partition {
161            sql.append_mut(p);
162        }
163        if let Some(o) = self.order {
164            sql.append_mut(o);
165        }
166        if let Some(f) = self.frame {
167            sql.append_mut(f);
168        }
169        sql
170    }
171}
172
173// =============================================================================
174// .over() on aggregate expressions — Agg → Scalar
175// =============================================================================
176
177impl<'a, V, T, N> SQLExpr<'a, V, T, N, Agg>
178where
179    V: SQLParam + 'a,
180    T: DataType,
181    N: Nullability,
182{
183    /// Apply a window specification to this aggregate expression.
184    ///
185    /// Converts the expression from `Agg` to `Scalar`, generating
186    /// `<expr> OVER (...)`.
187    ///
188    /// # Example
189    ///
190    /// ```rust
191    /// # let _ = r####"
192    /// sum(orders.amount).over(
193    ///     window()
194    ///         .partition_by([orders.customer_id])
195    ///         .order_by([asc(orders.date)])
196    /// )
197    /// # "####;
198    /// ```
199    pub fn over(self, spec: WindowSpec<'a, V>) -> SQLExpr<'a, V, T, N, Scalar> {
200        let sql = self
201            .into_sql()
202            .push(Token::OVER)
203            .push(Token::LPAREN)
204            .append(spec.into_sql())
205            .push(Token::RPAREN);
206        SQLExpr::new(sql)
207    }
208
209    /// Apply a FILTER clause to this aggregate (`PostgreSQL` extension).
210    ///
211    /// Generates `<agg> FILTER (WHERE <condition>)`.
212    #[must_use]
213    pub fn filter<C>(self, condition: C) -> Self
214    where
215        C: Expr<'a, V>,
216        C::SQLType: BooleanLike,
217    {
218        let sql = self
219            .into_sql()
220            .push(Token::FILTER)
221            .push(Token::LPAREN)
222            .push(Token::WHERE)
223            .append(condition.into_sql())
224            .push(Token::RPAREN);
225        SQLExpr::new(sql)
226    }
227}
228
229// =============================================================================
230// WindowFnExpr — pure window functions that require .over()
231// =============================================================================
232
233/// A window function expression that is not yet valid SQL.
234///
235/// Pure window functions like `ROW_NUMBER`, RANK, LAG, etc. MUST have an
236/// `.over()` call before they can be used in a query. This type enforces
237/// that at compile time by not implementing `Expr` or `ToSQL`.
238#[derive(Debug, Clone)]
239pub struct WindowFnExpr<'a, V: SQLParam, T: DataType, N: Nullability> {
240    sql: SQL<'a, V>,
241    _marker: PhantomData<(T, N)>,
242}
243
244impl<'a, V, T, N> WindowFnExpr<'a, V, T, N>
245where
246    V: SQLParam + 'a,
247    T: DataType,
248    N: Nullability,
249{
250    const fn new(sql: SQL<'a, V>) -> Self {
251        Self {
252            sql,
253            _marker: PhantomData,
254        }
255    }
256
257    /// Apply a window specification, producing a usable scalar expression.
258    ///
259    /// Generates `<fn> OVER (...)`.
260    pub fn over(self, spec: WindowSpec<'a, V>) -> SQLExpr<'a, V, T, N, Scalar> {
261        let sql = self
262            .sql
263            .push(Token::OVER)
264            .push(Token::LPAREN)
265            .append(spec.into_sql())
266            .push(Token::RPAREN);
267        SQLExpr::new(sql)
268    }
269}
270
271// =============================================================================
272// Pure Window Functions
273// =============================================================================
274
275/// `ROW_NUMBER()` — sequential row number within the partition.
276///
277/// Returns an integer, never NULL.
278#[must_use]
279pub fn row_number<'a, V>() -> WindowFnExpr<'a, V, <V::DialectMarker as CountPolicy>::Count, NonNull>
280where
281    V: SQLParam + 'a,
282    V::DialectMarker: CountPolicy,
283{
284    WindowFnExpr::new(SQL::raw("ROW_NUMBER()"))
285}
286
287/// `RANK()` — rank with gaps for ties.
288///
289/// Returns an integer, never NULL.
290#[must_use]
291pub fn rank<'a, V>() -> WindowFnExpr<'a, V, <V::DialectMarker as CountPolicy>::Count, NonNull>
292where
293    V: SQLParam + 'a,
294    V::DialectMarker: CountPolicy,
295{
296    WindowFnExpr::new(SQL::raw("RANK()"))
297}
298
299/// `DENSE_RANK()` — rank without gaps.
300///
301/// Returns an integer, never NULL.
302#[must_use]
303pub fn dense_rank<'a, V>() -> WindowFnExpr<'a, V, <V::DialectMarker as CountPolicy>::Count, NonNull>
304where
305    V: SQLParam + 'a,
306    V::DialectMarker: CountPolicy,
307{
308    WindowFnExpr::new(SQL::raw("DENSE_RANK()"))
309}
310
311/// NTILE(n) — divide rows into n roughly equal groups.
312///
313/// Returns an integer, never NULL.
314#[must_use]
315pub fn ntile<'a, V>(
316    n: usize,
317) -> WindowFnExpr<'a, V, <V::DialectMarker as CountPolicy>::Count, NonNull>
318where
319    V: SQLParam + 'a,
320    V::DialectMarker: CountPolicy,
321{
322    WindowFnExpr::new(SQL::func("NTILE", SQL::number(n)))
323}
324
325/// `PERCENT_RANK()` — relative rank of the current row: (rank - 1) / (total rows - 1).
326///
327/// Returns a float between 0.0 and 1.0, never NULL.
328#[must_use]
329pub fn percent_rank<'a, V>()
330-> WindowFnExpr<'a, V, <V::DialectMarker as FloatPolicy>::Float, NonNull>
331where
332    V: SQLParam + 'a,
333    V::DialectMarker: FloatPolicy,
334{
335    WindowFnExpr::new(SQL::raw("PERCENT_RANK()"))
336}
337
338/// `CUME_DIST()` — cumulative distribution: fraction of rows <= current row.
339///
340/// Returns a float between 0.0 and 1.0 (exclusive of 0), never NULL.
341#[must_use]
342pub fn cume_dist<'a, V>() -> WindowFnExpr<'a, V, <V::DialectMarker as FloatPolicy>::Float, NonNull>
343where
344    V: SQLParam + 'a,
345    V::DialectMarker: FloatPolicy,
346{
347    WindowFnExpr::new(SQL::raw("CUME_DIST()"))
348}
349
350/// LAG(expr) — value of expr from the previous row.
351///
352/// Returns the same type as expr, always nullable (no previous row → NULL).
353pub fn lag<'a, V, E>(expr: E) -> WindowFnExpr<'a, V, E::SQLType, Null>
354where
355    V: SQLParam + 'a,
356    E: Expr<'a, V>,
357{
358    WindowFnExpr::new(SQL::func("LAG", expr.into_sql()))
359}
360
361/// LAG(expr, offset, default) — value of expr from N rows back with a default.
362///
363/// Nullability is the combination of the expression's and default's nullability.
364pub fn lag_with_default<'a, V, E, D>(
365    expr: E,
366    offset: usize,
367    default: D,
368) -> WindowFnExpr<'a, V, E::SQLType, <E::Nullable as NullOr<D::Nullable>>::Output>
369where
370    V: SQLParam + 'a,
371    E: Expr<'a, V>,
372    D: Expr<'a, V>,
373    E::SQLType: Compatible<D::SQLType>,
374    E::Nullable: NullOr<D::Nullable>,
375    D::Nullable: Nullability,
376{
377    let args = expr
378        .into_sql()
379        .push(Token::COMMA)
380        .append(SQL::number(offset))
381        .push(Token::COMMA)
382        .append(default.into_sql());
383    WindowFnExpr::new(SQL::func("LAG", args))
384}
385
386/// LEAD(expr) — value of expr from the next row.
387///
388/// Returns the same type as expr, always nullable (no next row → NULL).
389pub fn lead<'a, V, E>(expr: E) -> WindowFnExpr<'a, V, E::SQLType, Null>
390where
391    V: SQLParam + 'a,
392    E: Expr<'a, V>,
393{
394    WindowFnExpr::new(SQL::func("LEAD", expr.into_sql()))
395}
396
397/// LEAD(expr, offset, default) — value of expr from N rows ahead with a default.
398///
399/// Nullability is the combination of the expression's and default's nullability.
400pub fn lead_with_default<'a, V, E, D>(
401    expr: E,
402    offset: usize,
403    default: D,
404) -> WindowFnExpr<'a, V, E::SQLType, <E::Nullable as NullOr<D::Nullable>>::Output>
405where
406    V: SQLParam + 'a,
407    E: Expr<'a, V>,
408    D: Expr<'a, V>,
409    E::SQLType: Compatible<D::SQLType>,
410    E::Nullable: NullOr<D::Nullable>,
411    D::Nullable: Nullability,
412{
413    let args = expr
414        .into_sql()
415        .push(Token::COMMA)
416        .append(SQL::number(offset))
417        .push(Token::COMMA)
418        .append(default.into_sql());
419    WindowFnExpr::new(SQL::func("LEAD", args))
420}
421
422/// `FIRST_VALUE(expr)` — value of expr from the first row of the frame.
423///
424/// Always nullable (frame may be empty for some edge cases).
425pub fn first_value<'a, V, E>(expr: E) -> WindowFnExpr<'a, V, E::SQLType, Null>
426where
427    V: SQLParam + 'a,
428    E: Expr<'a, V>,
429{
430    WindowFnExpr::new(SQL::func("FIRST_VALUE", expr.into_sql()))
431}
432
433/// `LAST_VALUE(expr)` — value of expr from the last row of the frame.
434///
435/// Always nullable (frame boundaries affect result).
436pub fn last_value<'a, V, E>(expr: E) -> WindowFnExpr<'a, V, E::SQLType, Null>
437where
438    V: SQLParam + 'a,
439    E: Expr<'a, V>,
440{
441    WindowFnExpr::new(SQL::func("LAST_VALUE", expr.into_sql()))
442}
443
444/// `NTH_VALUE(expr`, n) — value of expr from the nth row of the frame.
445///
446/// Always nullable (n may exceed frame size).
447pub fn nth_value<'a, V, E>(expr: E, n: usize) -> WindowFnExpr<'a, V, E::SQLType, Null>
448where
449    V: SQLParam + 'a,
450    E: Expr<'a, V>,
451{
452    let args = expr.into_sql().push(Token::COMMA).append(SQL::number(n));
453    WindowFnExpr::new(SQL::func("NTH_VALUE", args))
454}