laminar_sql/parser/window_rewriter/
mod.rs

1//! Production implementation for window function extraction and rewriting
2//!
3//! This module handles:
4//! - TUMBLE(time_col, interval) - tumbling windows
5//! - HOP(time_col, slide, size) / SLIDE(...) - sliding/hopping windows
6//! - SESSION(time_col, gap) - session windows
7
8use sqlparser::ast::{
9    Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, Query, Select, SelectItem,
10    SetExpr, Statement,
11};
12
13use super::{ParseError, WindowFunction};
14
15/// Rewrites window functions in SQL queries
16pub struct WindowRewriter;
17
18impl WindowRewriter {
19    /// Rewrite a SQL statement to expand window functions.
20    ///
21    /// Transforms window functions like TUMBLE into appropriate table functions
22    /// and adds `window_start`/`window_end` columns.
23    ///
24    /// # Errors
25    ///
26    /// Returns `ParseError::WindowError` if a window function cannot be rewritten.
27    ///
28    /// # Example
29    ///
30    /// ```sql
31    /// -- Input:
32    /// SELECT COUNT(*) FROM events
33    /// GROUP BY TUMBLE(event_time, INTERVAL '5' MINUTE)
34    ///
35    /// -- Output:
36    /// SELECT window_start, window_end, COUNT(*)
37    /// FROM events
38    /// GROUP BY window_start, window_end
39    /// ```
40    pub fn rewrite_statement(stmt: &mut Statement) -> Result<(), ParseError> {
41        if let Statement::Query(query) = stmt {
42            Self::rewrite_query(query)?;
43        }
44        Ok(())
45    }
46
47    /// Rewrite a query
48    fn rewrite_query(query: &mut Query) -> Result<(), ParseError> {
49        if let SetExpr::Select(select) = &mut *query.body {
50            Self::rewrite_select(select)?;
51        }
52        Ok(())
53    }
54
55    /// Rewrite a SELECT statement to expand window functions.
56    ///
57    /// This processes GROUP BY to find window functions and adds
58    /// window_start/window_end to the projection.
59    fn rewrite_select(select: &mut Select) -> Result<(), ParseError> {
60        // Find window function in GROUP BY
61        let window_func = Self::find_window_in_group_by(select)?;
62
63        if let Some(_window) = window_func {
64            // Add window_start and window_end to projection if not already present
65            Self::ensure_window_columns_in_projection(select);
66        }
67
68        Ok(())
69    }
70
71    /// Find window function in GROUP BY clause.
72    fn find_window_in_group_by(select: &Select) -> Result<Option<WindowFunction>, ParseError> {
73        // Check the GROUP BY expressions
74        match &select.group_by {
75            sqlparser::ast::GroupByExpr::Expressions(exprs, _modifiers) => {
76                for expr in exprs {
77                    if let Some(window) = Self::extract_window_function(expr)? {
78                        return Ok(Some(window));
79                    }
80                }
81            }
82            sqlparser::ast::GroupByExpr::All(_) => {}
83        }
84        Ok(None)
85    }
86
87    /// Ensure window_start and window_end columns are in projection.
88    fn ensure_window_columns_in_projection(select: &mut Select) {
89        let has_window_start = Self::has_projection_column(select, "window_start");
90        let has_window_end = Self::has_projection_column(select, "window_end");
91
92        // Add window_start at the beginning if not present
93        if !has_window_start {
94            select.projection.insert(
95                0,
96                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("window_start"))),
97            );
98        }
99
100        // Add window_end after window_start if not present
101        if !has_window_end {
102            select.projection.insert(
103                1,
104                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("window_end"))),
105            );
106        }
107    }
108
109    /// Check if a named column exists in the SELECT projection.
110    fn has_projection_column(select: &Select, name: &str) -> bool {
111        select.projection.iter().any(|item| {
112            if let SelectItem::UnnamedExpr(Expr::Identifier(ident)) = item {
113                ident.value.eq_ignore_ascii_case(name)
114            } else if let SelectItem::ExprWithAlias { alias, .. } = item {
115                alias.value.eq_ignore_ascii_case(name)
116            } else {
117                false
118            }
119        })
120    }
121
122    /// Check if expression contains a window function.
123    #[must_use]
124    pub fn contains_window_function(expr: &Expr) -> bool {
125        match expr {
126            Expr::Function(func) => {
127                if let Some(name) = func.name.0.last() {
128                    let func_name = name.to_string().to_uppercase();
129                    matches!(
130                        func_name.as_str(),
131                        "TUMBLE" | "HOP" | "SLIDE" | "SESSION" | "CUMULATE"
132                    )
133                } else {
134                    false
135                }
136            }
137            _ => false,
138        }
139    }
140
141    /// Extract window function details from expression.
142    ///
143    /// Parses the actual arguments from TUMBLE/HOP/SESSION function calls.
144    ///
145    /// # Supported syntax
146    ///
147    /// - `TUMBLE(time_column, interval)` - 2 arguments
148    /// - `HOP(time_column, slide_interval, window_size)` - 3 arguments
149    /// - `SLIDE(time_column, slide_interval, window_size)` - alias for HOP
150    /// - `SESSION(time_column, gap_interval)` - 2 arguments
151    ///
152    /// # Errors
153    ///
154    /// Returns `ParseError::WindowError` if:
155    /// - Function has empty name
156    /// - Wrong number of arguments for window type
157    /// - Arguments cannot be extracted
158    pub fn extract_window_function(expr: &Expr) -> Result<Option<WindowFunction>, ParseError> {
159        match expr {
160            Expr::Function(func) => {
161                let name =
162                    func.name.0.last().ok_or_else(|| {
163                        ParseError::WindowError("Empty function name".to_string())
164                    })?;
165
166                let func_name = name.to_string().to_uppercase();
167
168                // Extract arguments from the function
169                let args = Self::extract_function_args(&func.args)?;
170
171                match func_name.as_str() {
172                    "TUMBLE" => Self::parse_tumble_args(&args),
173                    "HOP" | "SLIDE" => Self::parse_hop_args(&args),
174                    "SESSION" => Self::parse_session_args(&args),
175                    "CUMULATE" => Self::parse_cumulate_args(&args),
176                    _ => Ok(None),
177                }
178            }
179            _ => Ok(None),
180        }
181    }
182
183    /// Extract function arguments as a vector of expressions.
184    fn extract_function_args(args: &FunctionArguments) -> Result<Vec<Expr>, ParseError> {
185        match args {
186            FunctionArguments::List(arg_list) => {
187                let mut result = Vec::new();
188                for arg in &arg_list.args {
189                    if let Some(expr) = Self::extract_arg_expr(arg) {
190                        result.push(expr);
191                    }
192                }
193                Ok(result)
194            }
195            FunctionArguments::None => Ok(vec![]),
196            FunctionArguments::Subquery(_) => Err(ParseError::WindowError(
197                "Subquery arguments not supported for window functions".to_string(),
198            )),
199        }
200    }
201
202    /// Extract expression from a function argument.
203    fn extract_arg_expr(arg: &FunctionArg) -> Option<Expr> {
204        match arg {
205            FunctionArg::Unnamed(arg_expr) => match arg_expr {
206                FunctionArgExpr::Expr(expr) => Some(expr.clone()),
207                FunctionArgExpr::Wildcard | FunctionArgExpr::QualifiedWildcard(_) => None,
208            },
209            FunctionArg::Named { arg, .. } | FunctionArg::ExprNamed { arg, .. } => match arg {
210                FunctionArgExpr::Expr(expr) => Some(expr.clone()),
211                FunctionArgExpr::Wildcard | FunctionArgExpr::QualifiedWildcard(_) => None,
212            },
213        }
214    }
215
216    /// Parse TUMBLE(time_column, interval [, offset]) arguments.
217    fn parse_tumble_args(args: &[Expr]) -> Result<Option<WindowFunction>, ParseError> {
218        if args.len() < 2 || args.len() > 3 {
219            return Err(ParseError::WindowError(format!(
220                "TUMBLE requires 2-3 arguments (time_column, interval [, offset]), got {}",
221                args.len()
222            )));
223        }
224
225        Ok(Some(WindowFunction::Tumble {
226            time_column: Box::new(args[0].clone()),
227            interval: Box::new(args[1].clone()),
228            offset: args.get(2).map(|e| Box::new(e.clone())),
229        }))
230    }
231
232    /// Parse HOP/SLIDE(time_column, slide_interval, window_size [, offset]) arguments.
233    fn parse_hop_args(args: &[Expr]) -> Result<Option<WindowFunction>, ParseError> {
234        if args.len() < 3 || args.len() > 4 {
235            return Err(ParseError::WindowError(format!(
236                "HOP/SLIDE requires 3-4 arguments (time_column, slide_interval, window_size [, offset]), got {}",
237                args.len()
238            )));
239        }
240
241        Ok(Some(WindowFunction::Hop {
242            time_column: Box::new(args[0].clone()),
243            slide_interval: Box::new(args[1].clone()),
244            window_interval: Box::new(args[2].clone()),
245            offset: args.get(3).map(|e| Box::new(e.clone())),
246        }))
247    }
248
249    /// Parse SESSION(time_column, gap_interval) arguments.
250    fn parse_session_args(args: &[Expr]) -> Result<Option<WindowFunction>, ParseError> {
251        if args.len() != 2 {
252            return Err(ParseError::WindowError(format!(
253                "SESSION requires 2 arguments (time_column, gap_interval), got {}",
254                args.len()
255            )));
256        }
257
258        Ok(Some(WindowFunction::Session {
259            time_column: Box::new(args[0].clone()),
260            gap_interval: Box::new(args[1].clone()),
261        }))
262    }
263
264    /// Parse CUMULATE(time_column, step_interval, max_size_interval) arguments.
265    fn parse_cumulate_args(args: &[Expr]) -> Result<Option<WindowFunction>, ParseError> {
266        if args.len() != 3 {
267            return Err(ParseError::WindowError(format!(
268                "CUMULATE requires 3 arguments (time_column, step_interval, max_size_interval), got {}",
269                args.len()
270            )));
271        }
272
273        Ok(Some(WindowFunction::Cumulate {
274            time_column: Box::new(args[0].clone()),
275            step_interval: Box::new(args[1].clone()),
276            max_size_interval: Box::new(args[2].clone()),
277        }))
278    }
279
280    /// Extract the time column name from a window function.
281    ///
282    /// Returns the column name as a string if extractable.
283    #[must_use]
284    pub fn get_time_column_name(window: &WindowFunction) -> Option<String> {
285        let expr = match window {
286            WindowFunction::Tumble { time_column, .. }
287            | WindowFunction::Hop { time_column, .. }
288            | WindowFunction::Session { time_column, .. }
289            | WindowFunction::Cumulate { time_column, .. } => time_column.as_ref(),
290        };
291
292        match expr {
293            Expr::Identifier(ident) => Some(ident.value.clone()),
294            Expr::CompoundIdentifier(parts) => parts.last().map(|p| p.value.clone()),
295            _ => None,
296        }
297    }
298
299    /// Parse an INTERVAL expression to Duration.
300    ///
301    /// Supports: SECOND, MINUTE, HOUR, DAY
302    ///
303    /// # Errors
304    ///
305    /// Returns `ParseError::WindowError` if the expression is not a valid interval.
306    pub fn parse_interval_to_duration(expr: &Expr) -> Result<std::time::Duration, ParseError> {
307        match expr {
308            Expr::Interval(interval) => {
309                // WHY: `INTERVAL '5 minutes'` carries its unit in the string, not the leading field.
310                if let (
311                    None,
312                    Expr::Value(sqlparser::ast::ValueWithSpan {
313                        value: sqlparser::ast::Value::SingleQuotedString(s),
314                        ..
315                    }),
316                ) = (&interval.leading_field, interval.value.as_ref())
317                {
318                    return Self::parse_interval_string(s);
319                }
320
321                // Extract the value
322                let value = Self::extract_interval_value(&interval.value)?;
323
324                // Get the unit (defaults to SECOND)
325                let unit = interval
326                    .leading_field
327                    .clone()
328                    .unwrap_or(sqlparser::ast::DateTimeField::Second);
329
330                match unit {
331                    sqlparser::ast::DateTimeField::Millisecond
332                    | sqlparser::ast::DateTimeField::Milliseconds => {
333                        return Ok(std::time::Duration::from_millis(value));
334                    }
335                    _ => {}
336                }
337
338                let seconds =
339                    match unit {
340                        sqlparser::ast::DateTimeField::Second
341                        | sqlparser::ast::DateTimeField::Seconds => value,
342                        sqlparser::ast::DateTimeField::Minute
343                        | sqlparser::ast::DateTimeField::Minutes => value * 60,
344                        sqlparser::ast::DateTimeField::Hour
345                        | sqlparser::ast::DateTimeField::Hours => value * 3600,
346                        sqlparser::ast::DateTimeField::Day
347                        | sqlparser::ast::DateTimeField::Days => value * 86400,
348                        _ => {
349                            return Err(ParseError::WindowError(format!(
350                                "Unsupported interval unit: {unit:?}"
351                            )))
352                        }
353                    };
354
355                Ok(std::time::Duration::from_secs(seconds))
356            }
357            // Handle string literal intervals like '5 MINUTES'
358            Expr::Value(value_with_span) => {
359                use sqlparser::ast::Value;
360                if let Value::SingleQuotedString(s) = &value_with_span.value {
361                    Self::parse_interval_string(s)
362                } else {
363                    Err(ParseError::WindowError(format!(
364                        "Expected string value, got: {value_with_span:?}"
365                    )))
366                }
367            }
368            // Handle identifier that might be an interval string
369            Expr::Identifier(ident) => Self::parse_interval_string(&ident.value),
370            _ => Err(ParseError::WindowError(format!(
371                "Expected INTERVAL expression, got: {expr:?}"
372            ))),
373        }
374    }
375
376    /// Extract numeric value from interval expression.
377    fn extract_interval_value(expr: &Expr) -> Result<u64, ParseError> {
378        match expr {
379            Expr::Value(value_with_span) => {
380                use sqlparser::ast::Value;
381                match &value_with_span.value {
382                    Value::Number(n, _) => n.parse::<u64>().map_err(|_| {
383                        ParseError::WindowError(format!("Invalid interval value: {n}"))
384                    }),
385                    Value::SingleQuotedString(s) => {
386                        // Handle '5' or '5 MINUTE'
387                        let num_str = s.split_whitespace().next().unwrap_or(s);
388                        num_str.parse::<u64>().map_err(|_| {
389                            ParseError::WindowError(format!("Invalid interval value: {s}"))
390                        })
391                    }
392                    _ => Err(ParseError::WindowError(format!(
393                        "Unsupported value type in interval: {value_with_span:?}"
394                    ))),
395                }
396            }
397            _ => Err(ParseError::WindowError(format!(
398                "Cannot extract interval value from: {expr:?}"
399            ))),
400        }
401    }
402
403    /// Parse an interval string like "5 MINUTES" or "1 HOUR".
404    pub(crate) fn parse_interval_string(s: &str) -> Result<std::time::Duration, ParseError> {
405        let parts: Vec<&str> = s.split_whitespace().collect();
406        if parts.is_empty() {
407            return Err(ParseError::WindowError("Empty interval string".to_string()));
408        }
409
410        let value: u64 = parts[0].parse().map_err(|_| {
411            ParseError::WindowError(format!("Invalid interval value: {}", parts[0]))
412        })?;
413
414        let unit = if parts.len() > 1 {
415            parts[1].to_uppercase()
416        } else {
417            "SECOND".to_string()
418        };
419
420        if matches!(unit.as_str(), "MILLISECOND" | "MILLISECONDS" | "MS") {
421            return Ok(std::time::Duration::from_millis(value));
422        }
423
424        let seconds = match unit.as_str() {
425            "SECOND" | "SECONDS" | "S" => value,
426            "MINUTE" | "MINUTES" | "M" => value * 60,
427            "HOUR" | "HOURS" | "H" => value * 3600,
428            "DAY" | "DAYS" | "D" => value * 86400,
429            _ => {
430                return Err(ParseError::WindowError(format!(
431                    "Unsupported interval unit: {unit}"
432                )))
433            }
434        };
435
436        Ok(std::time::Duration::from_secs(seconds))
437    }
438}
439
440#[cfg(test)]
441mod tests;
laminar_sql/parser/window_rewriter/mod.rs

laminar_sql/parser/window_rewriter/
mod.rs
