Repository: prql/prql
Pull request #471 (2022-05-09 05:17:14)

Allow passing a range to `take`
This is a WIP. Ideally we can do

```
from employees
take 5..10
```

And it'll translate to a `LIMIT` & `OFFSET`

Passing `take 10` would be the same as passing `take 1..10`.



max-sixty: Docs:
<img width="820" alt="image" src="https://user-images.githubusercontent.com/5635139/168457346-34c437ed-b08c-462b-b7de-9faeb2a71698.png">

<img width="802" alt="image" src="https://user-images.githubusercontent.com/5635139/168457338-04e6778b-acce-4968-8536-660bc5284f61.png">

aljazerzen: Sorry, couldn't help it: I did a refactor of `range_of_ranges`. I just like to rotate code around until it looks nice.

But I found a mistake, best seen in:
```elm
take 101..110
```
```diff
- LIMIT 110 OFFSET 100
+ LIMIT 10 OFFSET 100
```
Limit directly tells how much rows should there be in the result - its applied after offset.

Patch:
diff --git a/prql-compiler/src/ast/item.rs b/prql-compiler/src/ast/item.rs
index b84d6a25..0a8c8a9a 100644
--- a/prql-compiler/src/ast/item.rs
+++ b/prql-compiler/src/ast/item.rs
@@ -8,7 +8,7 @@ use enum_as_inner::EnumAsInner;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
-pub use super::*;
+use super::*;
 
 #[derive(Debug, EnumAsInner, PartialEq, Clone, Serialize, Deserialize)]
 pub enum Item {
@@ -97,7 +97,7 @@ impl Range {
         }
     }
 
-    pub fn new_int(start: Option<i64>, end: Option<i64>) -> Self {
+    pub fn from_ints(start: Option<i64>, end: Option<i64>) -> Self {
         let start = start.map(|x| Box::new(Node::from(Item::Literal(Literal::Integer(x)))));
         let end = end.map(|x| Box::new(Node::from(Item::Literal(Literal::Integer(x)))));
         Range { start, end }
diff --git a/prql-compiler/src/ast/mod.rs b/prql-compiler/src/ast/mod.rs
index ab6cdcdf..52cb3e88 100644
--- a/prql-compiler/src/ast/mod.rs
+++ b/prql-compiler/src/ast/mod.rs
@@ -44,7 +44,7 @@ impl Node {
         node
     }
 
-    /// Return an error if this is named expression.
+    /// Return an error if this is a named expression.
     pub fn discard_name(self) -> Result<Node, Error> {
         // TODO: replace this function with a prior type checking
 
diff --git a/prql-compiler/src/ast/query.rs b/prql-compiler/src/ast/query.rs
index 329242bb..af94d806 100644
--- a/prql-compiler/src/ast/query.rs
+++ b/prql-compiler/src/ast/query.rs
@@ -40,7 +40,7 @@ pub enum Transform {
         by: Vec<Node>,
     },
     Sort(Vec<ColumnSort<Node>>),
-    Take(i64),
+    Take(Range),
     Join {
         side: JoinSide,
         with: TableRef,
diff --git a/prql-compiler/src/parser.rs b/prql-compiler/src/parser.rs
index 3295c57f..0d8d8de3 100644
--- a/prql-compiler/src/parser.rs
+++ b/prql-compiler/src/parser.rs
@@ -362,6 +362,39 @@ mod test {
     use super::*;
     use insta::{assert_debug_snapshot, assert_yaml_snapshot};
 
+    #[test]
+    fn test_parse_take() -> Result<()> {
+        parse_tree_of_str("take 10", Rule::query)?;
+
+        assert_yaml_snapshot!(ast_of_string(r#"take 10"#, Rule::expr_call)?, @r###"
+        ---
+        FuncCall:
+          name: take
+          args:
+            - Literal:
+                Integer: 10
+          named_args: {}
+        "###);
+
+        // Currently this parses but doesn't translate.
+        assert_yaml_snapshot!(ast_of_string(r#"take 1..10"#, Rule::expr_call)?, @r###"
+        ---
+        FuncCall:
+          name: take
+          args:
+            - Range:
+                start:
+                  Literal:
+                    Integer: 1
+                end:
+                  Literal:
+                    Integer: 10
+          named_args: {}
+        "###);
+
+        Ok(())
+    }
+
     #[test]
     fn test_parse_pipeline_parse_tree() {
         assert_debug_snapshot!(parse_tree_of_str(
diff --git a/prql-compiler/src/semantic/transforms.rs b/prql-compiler/src/semantic/transforms.rs
index 5424ec82..8915fbd6 100644
--- a/prql-compiler/src/semantic/transforms.rs
+++ b/prql-compiler/src/semantic/transforms.rs
@@ -118,7 +118,13 @@ pub fn cast_transform(func_call: FuncCall, span: Option<Span>) -> Result<Transfo
         "take" => {
             let ([expr], []) = unpack(func_call, [])?;
 
-            Transform::Take(expr.discard_name()?.item.into_literal()?.into_integer()?)
+            match expr.discard_name()?.item {
+                Item::Literal(Literal::Integer(n)) => {
+                    Transform::Take(Range::from_ints(None, Some(n)))
+                }
+                Item::Range(range) => Transform::Take(range),
+                _ => unimplemented!(),
+            }
         }
         "join" => {
             let ([with, filter], [side]) = unpack(func_call, ["side"])?;
@@ -244,11 +250,11 @@ pub fn cast_transform(func_call: FuncCall, span: Option<Span>) -> Result<Transfo
             };
 
             let (kind, range) = if expanding {
-                (WindowKind::Rows, Range::new_int(None, Some(1)))
+                (WindowKind::Rows, Range::from_ints(None, Some(1)))
             } else if rolling > 0 {
                 (
                     WindowKind::Rows,
-                    Range::new_int(Some(-rolling + 1), Some(1)),
+                    Range::from_ints(Some(-rolling + 1), Some(1)),
                 )
             } else if let Some(range) = rows {
                 (WindowKind::Rows, range)
diff --git a/prql-compiler/src/sql/materializer.rs b/prql-compiler/src/sql/materializer.rs
index 5bda2833..664a6a48 100644
--- a/prql-compiler/src/sql/materializer.rs
+++ b/prql-compiler/src/sql/materializer.rs
@@ -877,7 +877,11 @@ take 20
               alias: ~
               declared_at: 37
         - Transform:
-            Take: 20
+            Take:
+              start: ~
+              end:
+                Literal:
+                  Integer: 20
         "###);
         assert_yaml_snapshot!(frame.columns, @r###"
         ---
diff --git a/prql-compiler/src/sql/translator.rs b/prql-compiler/src/sql/translator.rs
index 356543e3..fb0c4fe0 100644
--- a/prql-compiler/src/sql/translator.rs
+++ b/prql-compiler/src/sql/translator.rs
@@ -22,6 +22,7 @@ use crate::ast::JoinFilter;
 use crate::ast::*;
 use crate::error::{Error, Reason};
 use crate::semantic::Context;
+use crate::utils::OrMap;
 
 use super::materializer::MaterializationContext;
 use super::{un_group, MaterializedFrame};
@@ -225,14 +226,25 @@ fn sql_query_of_atomic_table(table: AtomicTable, dialect: &Dialect) -> Result<sq
     let where_ = filter_of_pipeline(before)?;
     let having = filter_of_pipeline(after)?;
 
-    let take = transforms
+    let takes = transforms
         .iter()
         .filter_map(|t| match t {
-            Transform::Take(take) => Some(*take),
+            Transform::Take(take) => Some(take.clone()),
             _ => None,
         })
-        .min()
-        .map(expr_of_i64);
+        .collect();
+    let take = range_of_ranges(takes)?;
+    let offset = take.start.map(|s| s - 1).unwrap_or(0);
+    let limit = take.end.map(|e| e - offset);
+
+    let offset = if offset == 0 {
+        None
+    } else {
+        Some(sqlparser::ast::Offset {
+            value: Item::Literal(Literal::Integer(offset)).try_into()?,
+            rows: sqlparser::ast::OffsetRows::None,
+        })
+    };
 
     // Use sorting from the frame
     let order_by = (frame.sort)
@@ -253,7 +265,7 @@ fn sql_query_of_atomic_table(table: AtomicTable, dialect: &Dialect) -> Result<sq
         body: SetExpr::Select(Box::new(Select {
             distinct: false,
             top: if dialect.use_top() {
-                take.clone().map(top_of_expr)
+                limit.map(top_of_i64)
             } else {
                 None
             },
@@ -273,8 +285,12 @@ fn sql_query_of_atomic_table(table: AtomicTable, dialect: &Dialect) -> Result<sq
         })),
         order_by,
         with: None,
-        limit: if dialect.use_top() { None } else { take },
-        offset: None,
+        limit: if dialect.use_top() {
+            None
+        } else {
+            limit.map(expr_of_i64)
+        },
+        offset,
         fetch: None,
         lock: None,
     })
@@ -283,12 +299,12 @@ fn sql_query_of_atomic_table(table: AtomicTable, dialect: &Dialect) -> Result<sq
 /// Convert a pipeline into a number of pipelines which can each "fit" into a SELECT.
 fn atomic_pipelines_of_pipeline(pipeline: Pipeline) -> Result<Vec<AtomicTable>> {
     // Insert a cut, when we find transformation that out of order:
-    // - joins,
+    // - joins (no limit),
     // - filters (for WHERE)
     // - aggregate (max 1x)
     // - sort (max 1x)
     // - filters (for HAVING)
-    // - take (max 1x)
+    // - take (no limit)
     //
     // Select and derive should already be extracted during resolving phase.
     //
@@ -314,7 +330,13 @@ fn atomic_pipelines_of_pipeline(pipeline: Pipeline) -> Result<Vec<AtomicTable>>
                     || counts.get("Sort").is_some()
                     || counts.get("Take").is_some()
             }
-            "Filter" | "Sort" | "Take" => counts.get("Take").is_some(),
+            "Filter" | "Sort" => counts.get("Take").is_some(),
+
+            // There can be many takes, but they have to be consecutive
+            // For example `take 100 | sort a | take 10` can't be one CTE.
+            // But this is enforced by transform order anyway.
+            "Take" => false,
+
             _ => false,
         };
 
@@ -390,6 +412,44 @@ fn prepend_with_from(pipeline: &mut Pipeline, table: Option<TableRef>) {
     }
 }
 
+type RangeI64 = core::ops::Range<Option<i64>>;
+
+/// Aggregate several ordered ranges into one, computing the intersection.
+///
+/// Returns a tuple of `(start, end)`, where `end` is optional.
+fn range_of_ranges(ranges: Vec<Range>) -> Result<RangeI64> {
+    #[allow(clippy::boxed_local)]
+    fn cast_bound(bound: Box<Node>) -> Result<i64> {
+        Ok(bound.item.into_literal()?.into_integer()?)
+    }
+
+    let mut current = RangeI64::default();
+    for range in ranges {
+        let mut range = RangeI64 {
+            start: range.start.map(cast_bound).transpose()?,
+            end: range.end.map(cast_bound).transpose()?,
+        };
+
+        // b = b + a.start -1 (take care of 1-based index!)
+        range.start = range.start.or_map(current.start, |a, b| a + b - 1);
+        range.end = range.end.map(|b| current.start.unwrap_or(1) + b - 1);
+
+        // b.end = min(a.end, b.end)
+        range.end = current.end.or_map(range.end, i64::min);
+        current = range;
+    }
+
+    if current
+        .start
+        .zip(current.end)
+        .map(|(s, e)| e <= s)
+        .unwrap_or(false)
+    {
+        bail!("Range end is before its start.");
+    }
+    Ok(current)
+}
+
 fn expr_of_i64(number: i64) -> Expr {
     Expr::Value(Value::Number(
         number.to_string(),
@@ -397,9 +457,9 @@ fn expr_of_i64(number: i64) -> Expr {
     ))
 }
 
-fn top_of_expr(take: Expr) -> Top {
+fn top_of_i64(take: i64) -> Top {
     Top {
-        quantity: Some(take),
+        quantity: Some(Item::Literal(Literal::Integer(take)).try_into().unwrap()),
         with_ties: false,
         percent: false,
     }
@@ -531,7 +591,7 @@ impl TryFrom<Item> for Expr {
                 let default_frame = if window.sort.is_empty() {
                     (WindowKind::Rows, Range::unbounded())
                 } else {
-                    (WindowKind::Range, Range::new_int(None, Some(1)))
+                    (WindowKind::Range, Range::from_ints(None, Some(1)))
                 };
 
                 let window = WindowSpec {
@@ -572,7 +632,6 @@ impl TryFrom<Item> for Expr {
         })
     }
 }
-
 fn try_into_is_null(nodes: &[Node]) -> Result<Option<Expr>> {
     if nodes.len() == 3 {
         if let Item::Operator(op) = &nodes[1].item {
@@ -696,6 +755,74 @@ mod test {
     };
     use serde_yaml::from_str;
 
+    #[test]
+    fn test_range_of_ranges() -> Result<()> {
+        // let range: Node = Item::Range(Range {
+        //     start: Some(Box::new(Item::Raw("1".to_string()).into())),
+        //     end: Some(Box::new(Item::Raw("10".to_string()).into())),
+        // })
+        // .into();
+        let range1 = Range::from_ints(Some(1), Some(10));
+        let range2 = Range::from_ints(Some(5), Some(6));
+        let range3 = Range::from_ints(Some(5), None);
+        let range4 = Range::from_ints(None, Some(8));
+
+        assert!(range_of_ranges(vec![range1.clone()])?.end.is_some());
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range1.clone()])?, @r###"
+        ---
+        start: 1
+        end: 10
+        "###);
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range1.clone(), range1.clone()])?, @r###"
+        ---
+        start: 1
+        end: 10
+        "###);
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range1.clone(), range2.clone()])?, @r###"
+        ---
+        start: 5
+        end: 6
+        "###);
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range2.clone(), range1.clone()])?, @r###"
+        ---
+        start: 5
+        end: 6
+        "###);
+
+        // We can't get 5..6 from 5..6.
+        assert!(range_of_ranges(vec![range2.clone(), range2.clone()]).is_err());
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range3.clone(), range3.clone()])?, @r###"
+        ---
+        start: 9
+        end: ~
+        "###);
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range1, range3])?, @r###"
+        ---
+        start: 5
+        end: 10
+        "###);
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range2, range4.clone()])?, @r###"
+        ---
+        start: 5
+        end: 6
+        "###);
+
+        assert_yaml_snapshot!(range_of_ranges(vec![range4.clone(), range4])?, @r###"
+        ---
+        start: ~
+        end: 8
+        "###);
+
+        Ok(())
+    }
+
     #[test]
     fn test_try_from_s_string_to_expr() -> Result<()> {
         let ast: Node = from_str(
@@ -1712,4 +1839,87 @@ take 20
 
         Ok(())
     }
+
+    #[test]
+    fn test_range() -> Result<()> {
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        take ..10
+        "###,
+        )?)?), @r###"
+        SELECT
+          employees.*
+        FROM
+          employees
+        LIMIT
+          10
+        "###);
+
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        take 5..10
+        "###,
+        )?)?), @r###"
+        SELECT
+          employees.*
+        FROM
+          employees
+        LIMIT
+          6 OFFSET 4
+        "###);
+
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        take 5..
+        "###,
+        )?)?), @r###"
+        SELECT
+          employees.*
+        FROM
+          employees OFFSET 4
+        "###);
+
+        // should be one SELECT
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        take 11..20
+        take 1..5
+        "###,
+        )?)?), @r###"
+        SELECT
+          employees.*
+        FROM
+          employees
+        LIMIT
+          5 OFFSET 10
+        "###);
+
+        // should be two SELECTs
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        take 11..20
+        sort name
+        take 1..5
+        "###,
+        )?)?), @r###"
+        WITH table_0 AS (
+          SELECT
+            employees.*
+          FROM
+            employees
+          LIMIT
+            10 OFFSET 10
+        )
+        SELECT
+          table_0.*
+        FROM
+          table_0
+        ORDER BY
+          name
+        LIMIT
+          5
+        "###);
+
+        Ok(())
+    }
 }
diff --git a/prql-compiler/src/utils.rs b/prql-compiler/src/utils.rs
index 7cf60665..6077b24f 100644
--- a/prql-compiler/src/utils.rs
+++ b/prql-compiler/src/utils.rs
@@ -88,3 +88,24 @@ pub fn diff(a: &str, b: &str) -> String {
     use similar::TextDiff;
     TextDiff::from_lines(a, b).unified_diff().to_string()
 }
+
+pub trait OrMap<T> {
+    /// Merges two options into one using `f`.
+    /// If one of the options is None, results defaults to the other one.
+    fn or_map<F>(self, b: Self, f: F) -> Self
+    where
+        F: FnOnce(T, T) -> T;
+}
+
+impl<T> OrMap<T> for Option<T> {
+    fn or_map<F>(self, b: Self, f: F) -> Self
+    where
+        F: FnOnce(T, T) -> T,
+    {
+        match (self, b) {
+            (Some(a), Some(b)) => Some(f(a, b)),
+            (a, None) => a,
+            (None, b) => b,
+        }
+    }
+}
diff --git a/reference/src/language-features/ranges.md b/reference/src/language-features/ranges.md
index 7f0b300d..7948d202 100644
--- a/reference/src/language-features/ranges.md
+++ b/reference/src/language-features/ranges.md
@@ -8,3 +8,12 @@ filter (age | in 18..40)
 ```
 
 Like in SQL, ranges are inclusive.
+
+As discussed in the [take](../transforms/take.md) docs, ranges can also be used
+in `take`:
+
+```prql
+from orders
+sort [-value, date]
+take 101..110
+```
diff --git a/reference/src/transforms/take.md b/reference/src/transforms/take.md
index 7c77aabb..8dd52702 100644
--- a/reference/src/transforms/take.md
+++ b/reference/src/transforms/take.md
@@ -3,7 +3,7 @@
 Picks rows based on their position.
 
 ```prql_no_test
-take {n}
+take {n|range}
 ```
 
 ## Examples
@@ -13,11 +13,8 @@ from employees
 take 10
 ```
 
-## Todo
-
-We could support a range expression so we can get an offset:
-
-```prql_no_test
-from employees
-take 1..10
+```prql
+from orders
+sort [-value, date]
+take 101..110
 ```


Test patch:
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 7b277760..841fb2aa 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
@@ -110,5 +110,9 @@ expression: mat
         - Literal:
             Integer: 200
 - Transform:
-    Take: 20
+    Take:
+      start: ~
+      end:
+        Literal:
+          Integer: 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 e1a8925e..9ee22e3c 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
@@ -110,5 +110,9 @@ expression: "&mat"
         - Literal:
             Integer: 200
 - Transform:
-    Take: 20
+    Take:
+      start: ~
+      end:
+        Literal:
+          Integer: 20
 
diff --git a/reference/tests/prql/language-features/ranges-1.prql b/reference/tests/prql/language-features/ranges-1.prql
new file mode 100644
index 00000000..f6939b84
--- /dev/null
+++ b/reference/tests/prql/language-features/ranges-1.prql
@@ -0,0 +1,3 @@
+from orders
+sort [-value, date]
+take 101..110
diff --git a/reference/tests/prql/transforms/take-1.prql b/reference/tests/prql/transforms/take-1.prql
new file mode 100644
index 00000000..f6939b84
--- /dev/null
+++ b/reference/tests/prql/transforms/take-1.prql
@@ -0,0 +1,3 @@
+from orders
+sort [-value, date]
+take 101..110
diff --git a/reference/tests/snapshots/snapshot__run_reference_prql@ranges-1.prql.snap b/reference/tests/snapshots/snapshot__run_reference_prql@ranges-1.prql.snap
new file mode 100644
index 00000000..be49701e
--- /dev/null
+++ b/reference/tests/snapshots/snapshot__run_reference_prql@ranges-1.prql.snap
@@ -0,0 +1,14 @@
+---
+source: reference/tests/snapshot.rs
+expression: sql
+input_file: reference/tests/prql/language-features/ranges-1.prql
+---
+SELECT
+  orders.*
+FROM
+  orders
+ORDER BY
+  value DESC,
+  date
+LIMIT
+  10 OFFSET 100
diff --git a/reference/tests/snapshots/snapshot__run_reference_prql@take-1.prql.snap b/reference/tests/snapshots/snapshot__run_reference_prql@take-1.prql.snap
new file mode 100644
index 00000000..3ba0064e
--- /dev/null
+++ b/reference/tests/snapshots/snapshot__run_reference_prql@take-1.prql.snap
@@ -0,0 +1,14 @@
+---
+source: reference/tests/snapshot.rs
+expression: sql
+input_file: reference/tests/prql/transforms/take-1.prql
+---
+SELECT
+  orders.*
+FROM
+  orders
+ORDER BY
+  value DESC,
+  date
+LIMIT
+  10 OFFSET 100