Repository: prql/prql
Pull request #369 (2022-04-16 03:40:22)

Implement intervals
This doesn't cover the full range of intervals, and not all dialects can do `{date} + {interval}` without a `DATE_ADD` function. But it's a start.


Patch:
diff --git a/prql-compiler/src/ast/ast_fold.rs b/prql-compiler/src/ast/ast_fold.rs
index 6bc381f0..54097c4e 100644
--- a/prql-compiler/src/ast/ast_fold.rs
+++ b/prql-compiler/src/ast/ast_fold.rs
@@ -136,7 +136,7 @@ pub fn fold_item<T: ?Sized + AstFold>(fold: &mut T, item: Item) -> Result<Item>
         }),
         // None of these capture variables, so we don't need to replace
         // them.
-        Item::String(_) | Item::Raw(_) => item,
+        Item::String(_) | Item::Raw(_) | Item::Interval(_) => item,
     })
 }
 
diff --git a/prql-compiler/src/ast/mod.rs b/prql-compiler/src/ast/mod.rs
index 7b78dacc..55686daa 100644
--- a/prql-compiler/src/ast/mod.rs
+++ b/prql-compiler/src/ast/mod.rs
@@ -43,6 +43,7 @@ pub enum Item {
     Table(Table),
     SString(Vec<InterpolateItem>),
     FString(Vec<InterpolateItem>),
+    Interval(Interval),
 }
 
 #[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
@@ -233,6 +234,18 @@ pub struct Range {
     pub end: Option<Box<Node>>,
 }
 
+// I could imagine there being a wrapper of this to represent "2 days 3 hours".
+// Or should that be written as `2days + 3hours`?
+//
+// #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+// pub struct Interval(pub Vec<IntervalPart>);
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct Interval {
+    pub n: i64,       // Do any DBs use floats or decimals for this?
+    pub unit: String, // Could be an enum IntervalType,
+}
+
 impl Node {
     /// For lists that only have one item in each ListItem this returns a Vec of
     /// those terms. (e.g. `[1, a b]` but not `[1 + 2]`, because `+` in an
diff --git a/prql-compiler/src/parser.rs b/prql-compiler/src/parser.rs
index b4b9a4e6..bf57295e 100644
--- a/prql-compiler/src/parser.rs
+++ b/prql-compiler/src/parser.rs
@@ -239,7 +239,6 @@ fn ast_of_parse_tree(pairs: Pairs<Rule>) -> Result<Vec<Node>> {
                         functions,
                     })
                 }
-                Rule::operator | Rule::number => Item::Raw(pair.as_str().to_owned()),
                 Rule::range => {
                     // a bit hacky, but eh
                     let no_start = &pair.as_span().as_str()[0..2] == "..";
@@ -264,6 +263,21 @@ fn ast_of_parse_tree(pairs: Pairs<Rule>) -> Result<Vec<Node>> {
                     };
                     Item::Range(Range { start, end })
                 }
+                Rule::interval => {
+                    let parsed = ast_of_parse_tree(pair.into_inner())?;
+                    // unimplemented!();
+                    let [n, unit]: [Node; 2] = parsed
+                        .try_into()
+                        .map_err(|e| anyhow!("Expected two items; {e:?}"))?;
+
+                    Item::Interval(Interval {
+                        n: n.item.as_raw().unwrap().parse()?,
+                        unit: unit.item.as_raw().unwrap().clone(),
+                    })
+                }
+                Rule::operator | Rule::number | Rule::interval_kind => {
+                    Item::Raw(pair.as_str().to_owned())
+                }
                 _ => unreachable!(),
             };
 
@@ -1466,4 +1480,31 @@ select [
                             Raw: "9"
         "###);
     }
+
+    #[test]
+    fn test_interval() {
+        assert_yaml_snapshot!(parse("
+        from employees
+        derive age_plus_two_years: (age + 2years)
+        ").unwrap(), @r###"
+        ---
+        version: ~
+        dialect: Generic
+        nodes:
+          - Pipeline:
+              - From:
+                  name: employees
+                  alias: ~
+              - Derive:
+                  - NamedExpr:
+                      name: age_plus_two_years
+                      expr:
+                        Expr:
+                          - Ident: age
+                          - Raw: +
+                          - Interval:
+                              n: 2
+                              unit: years
+        "###);
+    }
 }
diff --git a/prql-compiler/src/prql.pest b/prql-compiler/src/prql.pest
index 59d3cdfd..84e0c32c 100644
--- a/prql-compiler/src/prql.pest
+++ b/prql-compiler/src/prql.pest
@@ -71,7 +71,7 @@ func_call = { ident ~ (named_term_simple | term_simple)+ }
 func_curry = { ident ~ (named_term_simple | term_simple)* }
 
 term = _{ ( s_string | f_string | func_call | term_simple ) }
-term_simple = _{ ( s_string | f_string | ident | parenthesized_expr | list | range | number | string_literal | inline_pipeline ) }
+term_simple = _{ ( s_string | f_string | ident | parenthesized_expr | list | range | interval | number | string_literal | inline_pipeline ) }
 list = { "[" ~ NEWLINE* ~ named_expr ~ ("," ~ NEWLINE* ~ named_expr )* ~ ","? ~ NEWLINE* ~ "]" }
 parenthesized_expr = _{ "(" ~ expr ~ ")" }
 
@@ -111,3 +111,6 @@ interpolate_string = { ( !( "\\" | PEEK | "{" ) ~ ANY )+ }
 named_expr = { (ident ~ ":")? ~ expr }
 named_expr_simple = { (ident ~ ":")? ~ expr_simple }
 named_term_simple = { (ident ~ ":")? ~ term_simple }
+
+interval_kind = { "microseconds" | "milliseconds" | "seconds" | "minutes" | "hours" | "days" | "weeks" | "months" | "years" }
+interval = ${ number ~ interval_kind }
diff --git a/prql-compiler/src/translator.rs b/prql-compiler/src/translator.rs
index de926ed3..fdcdfa6c 100644
--- a/prql-compiler/src/translator.rs
+++ b/prql-compiler/src/translator.rs
@@ -11,12 +11,12 @@
 use anyhow::{anyhow, bail, Result};
 use itertools::Itertools;
 use sqlformat::{format, FormatOptions, QueryParams};
-use sqlparser::ast::Value;
 use sqlparser::ast::{
     self as sql_ast, Expr, Function, FunctionArg, FunctionArgExpr, Join, JoinConstraint,
     JoinOperator, ObjectName, OrderByExpr, Select, SelectItem, SetExpr, TableAlias, TableFactor,
     TableWithJoins, Top,
 };
+use sqlparser::ast::{DateTimeField, Value};
 use std::collections::HashMap;
 
 use super::ast::*;
@@ -528,6 +528,24 @@ impl TryFrom<Item> for Expr {
                     over: None,
                 })
             }
+            Item::Interval(interval) => {
+                let sql_parser_datetime = match interval.unit.as_str() {
+                    "years" => DateTimeField::Year,
+                    "months" => DateTimeField::Month,
+                    "days" => DateTimeField::Day,
+                    "hours" => DateTimeField::Hour,
+                    "minutes" => DateTimeField::Minute,
+                    "seconds" => DateTimeField::Second,
+                    _ => bail!("Unsupported interval unit: {}", interval.unit),
+                };
+                Expr::Value(Value::Interval {
+                    value: interval.n.to_string(),
+                    leading_field: Some(sql_parser_datetime),
+                    leading_precision: None,
+                    last_field: None,
+                    fractional_seconds_precision: None,
+                })
+            }
             _ => bail!("Can't convert to Expr at the moment; {item:?}"),
         })
     }
@@ -1246,4 +1264,24 @@ take 20
 
         Ok(())
     }
+
+    #[test]
+    fn test_interval() -> Result<()> {
+        let query: Query = parse(
+            r###"
+        from projects
+        derive first_check_in: start + 10days
+        "###,
+        )?;
+
+        assert_display_snapshot!((translate(&query)?), @r###"
+        SELECT
+          projects.*,
+          start + INTERVAL '10' DAY AS first_check_in
+        FROM
+          projects
+        "###);
+
+        Ok(())
+    }
 }
diff --git a/reference/src/SUMMARY.md b/reference/src/SUMMARY.md
index 5c3262a2..1ac2c393 100644
--- a/reference/src/SUMMARY.md
+++ b/reference/src/SUMMARY.md
@@ -6,11 +6,9 @@
 - [Language features](./language-features/README.md)
 
   - [F-Strings](./language-features/f-strings.md)
-    <!-- - Ranges — `BETWEEN 1 AND 3` -> `in 1..3`? -->
-  - [Ranges]()
+  - [Ranges](./language-features/ranges.md)
     <!-- - - Dates — `"2021-01-01"` -> `@2021-01-01`? And `DATE_TRUNC(foo_date, YEAR)` -> `foo_date.year`? Or -> `foo_date | as year`? Or `foo_date | to year`? -->
-    <!-- - Offsets — `DATE_ADD(DATE "2008-12-25", INTERVAL 5 DAY)` -> `@2008-12-25 - 5day`? -->
-  - [Dates]()
+  - [Dates](./language-features/dates.md)
     <!--   - Regex — `REGEX_MATCH(foo, "\\w{3}")` -> `foo ~ r"\w{3}"`? Or -> `regex foo r"\w{3}"`? -->
   - [Regex]()
   - [S-Strings](./language-features/s-strings.md)
diff --git a/reference/src/language-features/dates.md b/reference/src/language-features/dates.md
new file mode 100644
index 00000000..bd48cd15
--- /dev/null
+++ b/reference/src/language-features/dates.md
@@ -0,0 +1,10 @@
+# Dates
+
+### Intervals
+
+It's possible to represent intervals with a single expression:
+
+```prql
+from projects
+derive first_check_in: start + 10days
+```
diff --git a/reference/src/language-features/ranges.md b/reference/src/language-features/ranges.md
new file mode 100644
index 00000000..7f0b300d
--- /dev/null
+++ b/reference/src/language-features/ranges.md
@@ -0,0 +1,10 @@
+# Ranges
+
+PRQL has a concise range syntax `start..end`, which can be accessed with the `in` function:
+
+```prql
+from employees
+filter (age | in 18..40)
+```
+
+Like in SQL, ranges are inclusive.


Test patch:
diff --git a/reference/tests/examples/language-features/dates-0.prql b/reference/tests/examples/language-features/dates-0.prql
new file mode 100644
index 00000000..7ab19809
--- /dev/null
+++ b/reference/tests/examples/language-features/dates-0.prql
@@ -0,0 +1,2 @@
+from projects
+derive first_check_in: start + 10days
diff --git a/reference/tests/examples/language-features/ranges-0.prql b/reference/tests/examples/language-features/ranges-0.prql
new file mode 100644
index 00000000..dc46261f
--- /dev/null
+++ b/reference/tests/examples/language-features/ranges-0.prql
@@ -0,0 +1,2 @@
+from employees
+filter (age | in 18..40)
diff --git a/reference/tests/snapshots/snapshot__run_reference_examples@dates-0.prql.snap b/reference/tests/snapshots/snapshot__run_reference_examples@dates-0.prql.snap
new file mode 100644
index 00000000..518f2a7d
--- /dev/null
+++ b/reference/tests/snapshots/snapshot__run_reference_examples@dates-0.prql.snap
@@ -0,0 +1,10 @@
+---
+source: reference/tests/snapshot.rs
+expression: sql
+input_file: reference/tests/examples/language-features/dates-0.prql
+---
+SELECT
+  projects.*,
+  start + INTERVAL '10' DAY AS first_check_in
+FROM
+  projects
diff --git a/reference/tests/snapshots/snapshot__run_reference_examples@ranges-0.prql.snap b/reference/tests/snapshots/snapshot__run_reference_examples@ranges-0.prql.snap
new file mode 100644
index 00000000..890de50a
--- /dev/null
+++ b/reference/tests/snapshots/snapshot__run_reference_examples@ranges-0.prql.snap
@@ -0,0 +1,12 @@
+---
+source: reference/tests/snapshot.rs
+expression: sql
+input_file: reference/tests/examples/language-features/ranges-0.prql
+---
+SELECT
+  employees.*
+FROM
+  employees
+WHERE
+  age BETWEEN 18
+  AND 40