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