Repository: prql/prql
Pull request #461 (2022-05-07 04:27:06)

Add Operator to AST
Rather than just parsing as `Raw`. This will let us have decent code for
the `sort (-amount)` case.



Patch:
diff --git a/README.md b/README.md
index c6e52880..a2ff0050 100644
--- a/README.md
+++ b/README.md
@@ -160,14 +160,14 @@ Here's the same query with PRQL:
 ```elm
 prql version:0.3 db:snowflake                         # PRQL version & database name.
 
-func excess x -> (x - interest_rate) / 252             # Functions are clean and simple.
+func excess x -> (x - interest_rate) / 252            # Functions are clean and simple.
 func if_valid x -> is_valid_price ? x : null
-func lag_day x -> group sec_id (                       # `group` is used for window partitions too
-    sort date
-    window (                                          # `window` runs a pipeline over each window
-      lag 1 x                                         # `lag 1 x` lags the `x` col by 1
-    )
+func lag_day x -> group sec_id (                      # `group` is used for window partitions too
+  sort date
+  window (                                            # `window` runs a pipeline over each window
+    lag 1 x                                           # `lag 1 x` lags the `x` col by 1
   )
+)
 
 func ret x -> x / (x | lag_day) - 1 + dividend_return
 
@@ -185,7 +185,7 @@ select [                                              # `select` only includes u
     ln
     group sec_id (                                    # Complicated logic remains clear(er)
       sort date
-      window (
+      window ..current (                              # Rolling sum
         sum
       )
     )
diff --git a/prql-compiler/src/ast/ast_fold.rs b/prql-compiler/src/ast/ast_fold.rs
index a7b1aa55..093a0ae8 100644
--- a/prql-compiler/src/ast/ast_fold.rs
+++ b/prql-compiler/src/ast/ast_fold.rs
@@ -123,6 +123,7 @@ pub fn fold_item<T: ?Sized + AstFold>(fold: &mut T, item: Item) -> Result<Item>
         // them.
         Item::String(_)
         | Item::Raw(_)
+        | Item::Operator(_)
         | Item::Interval(_)
         | Item::Date(_)
         | Item::Time(_)
diff --git a/prql-compiler/src/ast/item.rs b/prql-compiler/src/ast/item.rs
index 39069289..3dbf18c1 100644
--- a/prql-compiler/src/ast/item.rs
+++ b/prql-compiler/src/ast/item.rs
@@ -15,6 +15,7 @@ pub enum Item {
     Ident(Ident),
     String(String),
     Raw(String),
+    Operator(String),
     Assign(NamedExpr),
     NamedArg(NamedExpr),
     Query(Query),
@@ -150,6 +151,9 @@ impl Display for Item {
             Item::String(s) => {
                 write!(f, "\"{s}\"")?;
             }
+            Item::Operator(o) => {
+                f.write_str(o)?;
+            }
             Item::Raw(r) => {
                 f.write_str(r)?;
             }
diff --git a/prql-compiler/src/parser.rs b/prql-compiler/src/parser.rs
index 65ea9682..378cbb73 100644
--- a/prql-compiler/src/parser.rs
+++ b/prql-compiler/src/parser.rs
@@ -252,16 +252,16 @@ fn ast_of_parse_tree(pairs: Pairs<Rule>) -> Result<Vec<Node>> {
                     Item::Type(Type { name, param })
                 }
                 Rule::boolean => Item::Boolean(pair.as_str() == "true"),
+                Rule::operator_unary
+                | Rule::operator_mul
+                | Rule::operator_add
+                | Rule::operator_compare
+                | Rule::operator_logical => Item::Operator(pair.as_str().to_owned()),
                 Rule::number
                 | Rule::interval_kind
                 | Rule::date_inner
                 | Rule::time_inner
-                | Rule::timestamp_inner
-                | Rule::operator_unary
-                | Rule::operator_mul
-                | Rule::operator_add
-                | Rule::operator_compare
-                | Rule::operator_logical => Item::Raw(pair.as_str().to_owned()),
+                | Rule::timestamp_inner => Item::Raw(pair.as_str().to_owned()),
 
                 _ => unreachable!("{pair}"),
             };
@@ -434,7 +434,7 @@ Canada
           - Expr:
               Expr:
                 - Raw: "2"
-                - Raw: +
+                - Operator: +
                 - Raw: "2"
           - String: )
         "###);
@@ -449,7 +449,7 @@ Canada
         List:
           - Expr:
               - Raw: "1"
-              - Raw: +
+              - Operator: +
               - Raw: "1"
           - Raw: "2"
         "###);
@@ -458,7 +458,7 @@ Canada
         List:
           - Expr:
               - Raw: "1"
-              - Raw: +
+              - Operator: +
               - FuncCall:
                   name: f
                   args:
@@ -534,7 +534,7 @@ Canada
                       args:
                         - Expr:
                             - Ident: country
-                            - Raw: "=="
+                            - Operator: "=="
                             - String: USA
                       named_args: {}
         "###);
@@ -558,7 +558,7 @@ Canada
                                 args:
                                   - Ident: country
                                 named_args: {}
-                            - Raw: "=="
+                            - Operator: "=="
                             - String: USA
                       named_args: {}
         "###
@@ -636,6 +636,21 @@ Canada
         "###);
     }
 
+    #[test]
+    fn test_parse_derive() -> Result<()> {
+        assert_yaml_snapshot!(
+            ast_of_string(r#"derive [x = 5, y = -x]"#, Rule::func_curry)?
+        , @r###"
+        ---
+        FuncCall:
+          name: derive
+          args: []
+          named_args: {}
+        "###);
+
+        Ok(())
+    }
+
     #[test]
     fn test_parse_select() -> Result<()> {
         assert_yaml_snapshot!(
@@ -673,7 +688,7 @@ Canada
         ---
         Expr:
           - Ident: country
-          - Raw: "=="
+          - Operator: "=="
           - String: USA
         "###);
         assert_yaml_snapshot!(ast_of_string(
@@ -689,14 +704,14 @@ Canada
               expr:
                 Expr:
                   - Ident: salary
-                  - Raw: +
+                  - Operator: +
                   - Ident: payroll_tax
           - Assign:
               name: gross_cost
               expr:
                 Expr:
                   - Ident: gross_salary
-                  - Raw: +
+                  - Operator: +
                   - Ident: benefits_cost
         "###);
         assert_yaml_snapshot!(
@@ -712,12 +727,12 @@ Canada
             Expr:
               - Expr:
                   - Ident: salary
-                  - Raw: +
+                  - Operator: +
                   - Ident: payroll_tax
-              - Raw: "*"
+              - Operator: "*"
               - Expr:
                   - Raw: "1"
-                  - Raw: +
+                  - Operator: +
                   - Ident: tax_rate
         "###);
         Ok(())
@@ -789,7 +804,7 @@ take 20
           body:
             Expr:
               - Ident: x
-              - Raw: +
+              - Operator: +
               - Raw: "1"
           return_type: ~
         "###);
@@ -807,7 +822,7 @@ take 20
           body:
             Expr:
               - Ident: x
-              - Raw: +
+              - Operator: +
               - Raw: "1"
           return_type: ~
         "###);
@@ -830,7 +845,7 @@ take 20
               args:
                 - Expr:
                     - Ident: bar
-                    - Raw: +
+                    - Operator: +
                     - Raw: "1"
               named_args: {}
           return_type: ~
@@ -896,7 +911,7 @@ take 20
           body:
             Expr:
               - Ident: x
-              - Raw: +
+              - Operator: +
               - Ident: to
           return_type: ~
         "###);
@@ -947,18 +962,18 @@ take 20
                       - List:
                           - Expr:
                               - Ident: a
-                              - Raw: and
+                              - Operator: and
                               - Expr:
                                   - Ident: b
-                                  - Raw: +
+                                  - Operator: +
                                   - Ident: c
-                              - Raw: or
+                              - Operator: or
                               - FuncCall:
                                   name: d
                                   args:
                                     - Ident: e
                                   named_args: {}
-                              - Raw: and
+                              - Operator: and
                               - Ident: f
                     named_args: {}
         "###);
@@ -1223,11 +1238,11 @@ take 20
                       - List:
                           - Expr:
                               - Ident: first_name
-                              - Raw: "=="
+                              - Operator: "=="
                               - Ident: $1
                           - Expr:
                               - Ident: last_name
-                              - Raw: "=="
+                              - Operator: "=="
                               - Ident: $2.name
                     named_args: {}
         "###);
@@ -1393,7 +1408,7 @@ select [
                               expr:
                                 Expr:
                                   - Ident: age
-                                  - Raw: +
+                                  - Operator: +
                                   - Interval:
                                       n: 2
                                       unit: years
diff --git a/prql-compiler/src/sql/materializer.rs b/prql-compiler/src/sql/materializer.rs
index eabe2907..3424a70a 100644
--- a/prql-compiler/src/sql/materializer.rs
+++ b/prql-compiler/src/sql/materializer.rs
@@ -633,7 +633,7 @@ take 20
                         - Expr:
                             Ident: foo
                         - String: )
-                    - Raw: +
+                    - Operator: +
                     - Raw: "1"
               by: []
         "###);
@@ -885,7 +885,7 @@ take 20
         - Ident: tax
         - Expr:
             - Ident: gross
-            - Raw: "-"
+            - Operator: "-"
             - Ident: tax
         "###);
 
diff --git a/prql-compiler/src/sql/translator.rs b/prql-compiler/src/sql/translator.rs
index 5679bb85..061d0d8c 100644
--- a/prql-compiler/src/sql/translator.rs
+++ b/prql-compiler/src/sql/translator.rs
@@ -442,7 +442,7 @@ impl TryFrom<Item> for Expr {
     type Error = anyhow::Error;
     fn try_from(item: Item) -> Result<Self> {
         Ok(match item {
-            Item::Ident(_) | Item::Raw(_) => Expr::Identifier(item.try_into()?),
+            Item::Ident(_) | Item::Raw(_) | Item::Operator(_) => Expr::Identifier(item.try_into()?),
             // For expressions like `country = "USA"`, we take each one, convert
             // it, and put spaces between them. It's a bit hacky — we could
             // convert each term to a SQL AST item, but it works for the moment.
@@ -553,7 +553,6 @@ impl TryFrom<Item> for Expr {
                 value: literal,
             },
             Item::Boolean(b) => Expr::Value(Value::Boolean(b)),
-
             _ => bail!("Can't convert to Expr; {item:?}"),
         })
     }
@@ -597,14 +596,13 @@ impl TryFrom<Item> for sql_ast::Ident {
     fn try_from(item: Item) -> Result<Self> {
         Ok(match item {
             Item::Ident(ident) => sql_ast::Ident::new(ident),
-            Item::Raw(raw) => {
-                let raw = match raw.as_str() {
-                    "==" => "=".to_string(),
-                    _ => raw,
-                };
-
-                sql_ast::Ident::new(raw.to_uppercase())
-            }
+            Item::Raw(raw) => sql_ast::Ident::new(raw.to_uppercase()),
+            // This isn't ideal; we could instead parse to the sqlparser AST.
+            // But it does OK for the moment.
+            Item::Operator(op) => sql_ast::Ident::new(match op.as_str() {
+                "==" => "=".to_string(),
+                _ => op.to_uppercase(),
+            }),
             _ => bail!("Can't convert to Ident; {item:?}"),
         })
     }
diff --git a/reference/src/functions.md b/reference/src/functions.md
index 117636f4..d0958643 100644
--- a/reference/src/functions.md
+++ b/reference/src/functions.md
@@ -23,7 +23,7 @@ from cities
 derive temp_c = (celsius_of_fahrenheit temp_f)
 ```
 
-This function named `interp`, and has two positional parameters named
+This function is named `interp`, and has two positional parameters named
 `higher` and `x`, and one named parameter named `lower` which takes a default
 argument of `0`. It calculates the proportion of the distance that `x` is
 between `lower` and `higher`.
@@ -41,18 +41,18 @@ derive [
 ## Piping
 
 Consistent with the principles of PRQL, it's possible to pipe values into
-functions, which makes composing functions more readable. When piping a value
-into a function, the value is passed as an argument to the final positional
-parameter of the function. Here's the same result as the examples above with an
-alternative construction:
+functions, which makes composing many functions more readable. When piping a
+value into a function, the value is passed as an argument to the final
+positional parameter of the function. Here's the same result as the examples
+above with an alternative construction:
 
 ```prql
 func interp lower:0 higher x ->  (x - lower) / (higher - lower)
 
 from students
 derive [
-  sat_proportion_1= (sat_score | interp 1600),
-  sat_proportion_2= (sat_score | interp lower:0 1600),
+  sat_proportion_1 = (sat_score | interp 1600),
+  sat_proportion_2 = (sat_score | interp lower:0 1600),
 ]
 ```
 


Test patch:
diff --git a/prql-compiler/src/snapshots/prql_compiler__parser__test__parse_query.snap b/prql-compiler/src/snapshots/prql_compiler__parser__test__parse_query.snap
index f3e6ca21..66ea9947 100644
--- a/prql-compiler/src/snapshots/prql_compiler__parser__test__parse_query.snap
+++ b/prql-compiler/src/snapshots/prql_compiler__parser__test__parse_query.snap
@@ -1,6 +1,6 @@
 ---
 source: prql-compiler/src/parser.rs
-expression: "ast_of_string(r#\"\nfrom employees\nfilter country == \"USA\"                        # Each line transforms the previous result.\nderive [                                      # This adds columns / variables.\n  gross_salary = salary + payroll_tax,\n  gross_cost =   gross_salary + benefits_cost # Variables can use other variables.\n]\nfilter gross_cost > 0\ngroup [title, country] (\naggregate [               # `by` are the columns to group by.\n                   average salary,            # These are aggregation calcs run on each group.\n                   sum salary,\n                   average gross_salary,\n                   sum gross_salary,\n                   average gross_cost,\n  sum_gross_cost = sum gross_cost,\n  ct             = count,\n] )\nsort sum_gross_cost\nfilter ct > 200\ntake 20\n    \"#.trim(),\n              Rule::query)?"
+expression: "ast_of_string(r#\"\nfrom employees\nfilter country == \"USA\"                        # Each line transforms the previous result.\nderive [                                      # This adds columns / variables.\n  gross_salary = salary + payroll_tax,\n  gross_cost =   gross_salary + benefits_cost # Variables can use other variables.\n]\nfilter gross_cost > 0\ngroup [title, country] (\naggregate [               # `by` are the columns to group by.\n                   average salary,            # These are aggregation calcs run on each group.\n                   sum salary,\n                   average gross_salary,\n                   sum gross_salary,\n                   average gross_cost,\n  sum_gross_cost = sum gross_cost,\n  ct             = count,\n] )\nsort sum_gross_cost\nfilter ct > 200\ntake 20\n    \"#.trim(),\n        Rule::query)?"
 ---
 Query:
   version: ~
@@ -19,7 +19,7 @@ Query:
               args:
                 - Expr:
                     - Ident: country
-                    - Raw: "=="
+                    - Operator: "=="
                     - String: USA
               named_args: {}
           - FuncCall:
@@ -31,14 +31,14 @@ Query:
                         expr:
                           Expr:
                             - Ident: salary
-                            - Raw: +
+                            - Operator: +
                             - Ident: payroll_tax
                     - Assign:
                         name: gross_cost
                         expr:
                           Expr:
                             - Ident: gross_salary
-                            - Raw: +
+                            - Operator: +
                             - Ident: benefits_cost
               named_args: {}
           - FuncCall:
@@ -46,7 +46,7 @@ Query:
               args:
                 - Expr:
                     - Ident: gross_cost
-                    - Raw: ">"
+                    - Operator: ">"
                     - Raw: "0"
               named_args: {}
           - FuncCall:
@@ -111,7 +111,7 @@ Query:
               args:
                 - Expr:
                     - Ident: ct
-                    - Raw: ">"
+                    - Operator: ">"
                     - Raw: "200"
               named_args: {}
           - FuncCall:
diff --git a/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__materialize_2.snap b/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__materialize_2.snap
index ae6f0519..d7d151a8 100644
--- a/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__materialize_2.snap
+++ b/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__materialize_2.snap
@@ -11,7 +11,7 @@ expression: mat
     Filter:
       - Expr:
           - Ident: country
-          - Raw: "=="
+          - Operator: "=="
           - String: USA
 - Transform:
     Filter:
@@ -19,11 +19,11 @@ expression: mat
           - Expr:
               - Expr:
                   - Ident: salary
-                  - Raw: +
+                  - Operator: +
                   - Ident: payroll_tax
-              - Raw: +
+              - Operator: +
               - Ident: benefits_cost
-          - Raw: ">"
+          - Operator: ">"
           - Raw: "0"
 - Transform:
     Group:
@@ -52,7 +52,7 @@ expression: mat
                         - Expr:
                             Expr:
                               - Ident: salary
-                              - Raw: +
+                              - Operator: +
                               - Ident: payroll_tax
                         - String: )
                     - SString:
@@ -60,7 +60,7 @@ expression: mat
                         - Expr:
                             Expr:
                               - Ident: salary
-                              - Raw: +
+                              - Operator: +
                               - Ident: payroll_tax
                         - String: )
                     - SString:
@@ -69,9 +69,9 @@ expression: mat
                             Expr:
                               - Expr:
                                   - Ident: salary
-                                  - Raw: +
+                                  - Operator: +
                                   - Ident: payroll_tax
-                              - Raw: +
+                              - Operator: +
                               - Ident: benefits_cost
                         - String: )
                     - SString:
@@ -80,9 +80,9 @@ expression: mat
                             Expr:
                               - Expr:
                                   - Ident: salary
-                                  - Raw: +
+                                  - Operator: +
                                   - Ident: payroll_tax
-                              - Raw: +
+                              - Operator: +
                               - Ident: benefits_cost
                         - String: )
                     - SString:
@@ -99,12 +99,12 @@ expression: mat
                   Expr:
                     - Expr:
                         - Ident: salary
-                        - Raw: +
+                        - Operator: +
                         - Ident: payroll_tax
-                    - Raw: +
+                    - Operator: +
                     - Ident: benefits_cost
               - String: )
-          - Raw: ">"
+          - Operator: ">"
           - Raw: "200"
 - Transform:
     Take: 20
diff --git a/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__replace_variables_2.snap b/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__replace_variables_2.snap
index a44750cc..2cabb62b 100644
--- a/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__replace_variables_2.snap
+++ b/prql-compiler/src/sql/snapshots/prql_compiler__sql__materializer__test__replace_variables_2.snap
@@ -11,7 +11,7 @@ expression: "&mat"
     Filter:
       - Expr:
           - Ident: country
-          - Raw: "=="
+          - Operator: "=="
           - String: USA
 - Transform:
     Filter:
@@ -19,11 +19,11 @@ expression: "&mat"
           - Expr:
               - Expr:
                   - Ident: salary
-                  - Raw: +
+                  - Operator: +
                   - Ident: payroll_tax
-              - Raw: +
+              - Operator: +
               - Ident: benefits_cost
-          - Raw: ">"
+          - Operator: ">"
           - Raw: "0"
 - Transform:
     Group:
@@ -52,7 +52,7 @@ expression: "&mat"
                         - Expr:
                             Expr:
                               - Ident: salary
-                              - Raw: +
+                              - Operator: +
                               - Ident: payroll_tax
                         - String: )
                     - SString:
@@ -60,7 +60,7 @@ expression: "&mat"
                         - Expr:
                             Expr:
                               - Ident: salary
-                              - Raw: +
+                              - Operator: +
                               - Ident: payroll_tax
                         - String: )
                     - SString:
@@ -69,9 +69,9 @@ expression: "&mat"
                             Expr:
                               - Expr:
                                   - Ident: salary
-                                  - Raw: +
+                                  - Operator: +
                                   - Ident: payroll_tax
-                              - Raw: +
+                              - Operator: +
                               - Ident: benefits_cost
                         - String: )
                     - SString:
@@ -80,9 +80,9 @@ expression: "&mat"
                             Expr:
                               - Expr:
                                   - Ident: salary
-                                  - Raw: +
+                                  - Operator: +
                                   - Ident: payroll_tax
-                              - Raw: +
+                              - Operator: +
                               - Ident: benefits_cost
                         - String: )
                     - SString:
@@ -99,12 +99,12 @@ expression: "&mat"
                   Expr:
                     - Expr:
                         - Ident: salary
-                        - Raw: +
+                        - Operator: +
                         - Ident: payroll_tax
-                    - Raw: +
+                    - Operator: +
                     - Ident: benefits_cost
               - String: )
-          - Raw: ">"
+          - Operator: ">"
           - Raw: "200"
 - Transform:
     Take: 20
diff --git a/reference/tests/examples/functions-2.prql b/reference/tests/examples/functions-2.prql
index cd672f84..b1d30227 100644
--- a/reference/tests/examples/functions-2.prql
+++ b/reference/tests/examples/functions-2.prql
@@ -2,6 +2,6 @@ func interp lower:0 higher x ->  (x - lower) / (higher - lower)
 
 from students
 derive [
-  sat_proportion_1= (sat_score | interp 1600),
-  sat_proportion_2= (sat_score | interp lower:0 1600),
+  sat_proportion_1 = (sat_score | interp 1600),
+  sat_proportion_2 = (sat_score | interp lower:0 1600),
 ]