Repository: prql/prql (rust)

Issue:
`group` not working as intended with `take`
Possibly we need to implement #420 given our plan of how to implement this query? Though in the meantime, should this raise an error rather than ignoring `group`?

```elm
from employees
group role (
  sort join_date  # taken from above
  take 1
)
```

```sql
SELECT
  employees.*
FROM
  employees
ORDER BY
  join_date
LIMIT
  1
```

Pull request:
Group and take a range
Closes #421

Patch:
diff --git a/prql-compiler/src/ast/ast_fold.rs b/prql-compiler/src/ast/ast_fold.rs
index a59a529c..aaeee04d 100644
--- a/prql-compiler/src/ast/ast_fold.rs
+++ b/prql-compiler/src/ast/ast_fold.rs
@@ -208,9 +208,10 @@ pub fn fold_transform<T: ?Sized + AstFold>(
             kind,
             pipeline: Box::new(fold.fold_node(*pipeline)?),
         },
-        Transform::Take { by, range } => Transform::Take {
+        Transform::Take { by, range, sort } => Transform::Take {
             range: fold_range(fold, range)?,
             by: fold.fold_nodes(by)?,
+            sort: fold.fold_column_sorts(sort)?,
         },
         Transform::Unique => Transform::Unique,
     })
diff --git a/prql-compiler/src/ast/query.rs b/prql-compiler/src/ast/query.rs
index 4ea95e7f..025207ed 100644
--- a/prql-compiler/src/ast/query.rs
+++ b/prql-compiler/src/ast/query.rs
@@ -43,6 +43,7 @@ pub enum Transform {
     Take {
         range: Range,
         by: Vec<Node>,
+        sort: Vec<ColumnSort<Node>>,
     },
     Join {
         side: JoinSide,
diff --git a/prql-compiler/src/semantic/resolver.rs b/prql-compiler/src/semantic/resolver.rs
index f9f13cda..4269cdd7 100644
--- a/prql-compiler/src/semantic/resolver.rs
+++ b/prql-compiler/src/semantic/resolver.rs
@@ -396,12 +396,26 @@ impl Resolver {
             | Transform::Window { .. } => {
                 // ok
             }
-            Transform::Aggregate { by, .. } | Transform::Take { by, .. } => {
+            Transform::Aggregate { by, .. } => {
                 *by = (self.within_group)
                     .iter()
                     .map(|id| Node::new_ident("<ref>", *id))
                     .collect();
             }
+            Transform::Take { by, sort, .. } => {
+                *by = (self.within_group)
+                    .iter()
+                    .map(|id| Node::new_ident("<ref>", *id))
+                    .collect();
+
+                *sort = (self.sorted)
+                    .iter()
+                    .map(|s| ColumnSort {
+                        column: Node::new_ident("<ref>", s.column),
+                        direction: s.direction.clone(),
+                    })
+                    .collect();
+            }
             _ => {
                 // TODO: attach span to this error
                 bail!(Error::new(Reason::Simple(format!(
@@ -675,10 +689,23 @@ mod tests {
             take 1
         )
         "#).and_then(resolve_and_translate).unwrap(), @r###"
+        WITH table_0 AS (
+          SELECT
+            employees.*,
+            ROW NUMBER() OVER (
+              PARTITION BY last_name
+              ORDER BY
+                first_name
+            ) AS _rn
+          FROM
+            employees
+        )
         SELECT
-          DISTINCT employees.*
+          table_0.*
         FROM
-          employees
+          table_0
+        WHERE
+          _rn <= 1
         "###);
 
         let res = parse(
diff --git a/prql-compiler/src/semantic/transforms.rs b/prql-compiler/src/semantic/transforms.rs
index 4d768d46..a9325c29 100644
--- a/prql-compiler/src/semantic/transforms.rs
+++ b/prql-compiler/src/semantic/transforms.rs
@@ -123,7 +123,11 @@ pub fn cast_transform(func_call: FuncCall, span: Option<Span>) -> Result<Transfo
                 Item::Range(range) => range,
                 _ => unimplemented!(),
             };
-            Transform::Take { range, by: vec![] }
+            Transform::Take {
+                range,
+                by: vec![],
+                sort: vec![],
+            }
         }
         "join" => {
             let ([with, filter], [side]) = unpack(func_call, ["side"])?;
@@ -364,6 +368,7 @@ mod tests {
                                       Integer: 1
                                 by:
                                   - Ident: "<ref>"
+                                sort: []
         "###);
 
         // oops, two arguments #339
diff --git a/prql-compiler/src/sql/distinct.rs b/prql-compiler/src/sql/distinct.rs
index aca72925..0c5cb9ba 100644
--- a/prql-compiler/src/sql/distinct.rs
+++ b/prql-compiler/src/sql/distinct.rs
@@ -1,18 +1,23 @@
-use anyhow::{bail, Result};
+use anyhow::Result;
 
-use crate::{
-    ast::{ast_fold::AstFold, *},
-    error::{Error, Reason},
-};
+use crate::ast::{ast_fold::AstFold, *};
+use crate::Declaration;
 
-pub fn take_to_distinct(nodes: Vec<Node>) -> Result<Vec<Node>> {
-    let mut d = DistinctMaker {};
+use super::materializer::MaterializationContext;
+
+pub fn take_to_distinct(
+    nodes: Vec<Node>,
+    context: &mut MaterializationContext,
+) -> Result<Vec<Node>> {
+    let mut d = DistinctMaker { context };
     d.fold_nodes(nodes)
 }
 /// Creates [Transform::Unique] from [Transform::Take]
-struct DistinctMaker {}
+struct DistinctMaker<'a> {
+    context: &'a mut MaterializationContext,
+}
 
-impl AstFold for DistinctMaker {
+impl<'a> AstFold for DistinctMaker<'a> {
     fn fold_nodes(&mut self, nodes: Vec<Node>) -> Result<Vec<Node>> {
         let mut res = Vec::new();
 
@@ -22,27 +27,21 @@ impl AstFold for DistinctMaker {
                     res.push(node);
                 }
 
-                Item::Transform(Transform::Take { range, .. }) => {
-                    let range = range.into_int()?;
-                    let take_only_first =
-                        range.start.unwrap_or(1) == 1 && matches!(range.end, Some(1));
+                Item::Transform(Transform::Take { range, by, sort }) => {
+                    let range_int = range.clone().into_int()?;
 
-                    if take_only_first {
+                    let take_only_first =
+                        range_int.start.unwrap_or(1) == 1 && matches!(range_int.end, Some(1));
+                    if take_only_first && sort.is_empty() {
                         // TODO: use distinct only if `by == all columns in frame`
                         res.push(Item::Transform(Transform::Unique).into());
-                    } else {
-                        bail!(Error::new(Reason::Simple(
-                            "`take` within `group` currently only supports argument `1`".to_string(),
-                        ))
-                        .with_span(node.span)
-                        .with_help("For now, you can derive a row number within a group and do a filter on that."));
+                        continue;
                     }
 
-                    // TODO: else
-                    //  return `
-                    //    derive _rn = row number over (order by)
-                    //    filter (_rn | in range)
-                    //  `
+                    // convert `take range` into:
+                    //   derive _rn = s"ROW NUMBER"
+                    //   filter (_rn | in range)
+                    res.extend(self.filter_row_number(range, sort, by));
                 }
                 _ => {
                     res.push(node);
@@ -56,3 +55,75 @@ impl AstFold for DistinctMaker {
         Ok(function)
     }
 }
+
+impl<'a> DistinctMaker<'a> {
+    fn filter_row_number(
+        &mut self,
+        range: Range,
+        sort: Vec<ColumnSort>,
+        by: Vec<Node>,
+    ) -> Vec<Node> {
+        let range_int = range.clone().into_int().unwrap();
+
+        // declare new column
+        let decl = Node::from(Item::SString(vec![InterpolateItem::String(
+            "ROW NUMBER()".to_string(),
+        )]));
+        let is_unsorted = sort.is_empty();
+        let windowed = Windowed {
+            expr: Box::new(decl),
+            group: by,
+            sort,
+            window: if is_unsorted {
+                (WindowKind::Rows, Range::unbounded())
+            } else {
+                (WindowKind::Range, Range::from_ints(None, Some(0)))
+            },
+        };
+        let decl = Declaration::Expression(Box::new(Item::Windowed(windowed).into()));
+        let row_number_id = self.context.declare(decl);
+
+        // name it _rn
+        let mut ident = Node::from(Item::Ident("_rn".to_string()));
+        ident.declared_at = Some(row_number_id);
+
+        // add the two transforms
+        let transforms = vec![
+            Transform::Derive(vec![ident.clone()]),
+            Transform::Filter(Box::new(match (range_int.start, range_int.end) {
+                (Some(s), Some(e)) if s == e => Item::Expr(vec![
+                    ident,
+                    Item::Operator("==".to_string()).into(),
+                    Item::Literal(Literal::Integer(s)).into(),
+                ])
+                .into(),
+                (Some(s), None) => Item::Expr(vec![
+                    ident,
+                    Item::Operator(">=".to_string()).into(),
+                    Item::Literal(Literal::Integer(s)).into(),
+                ])
+                .into(),
+                (None, Some(e)) => Item::Expr(vec![
+                    ident,
+                    Item::Operator("<=".to_string()).into(),
+                    Item::Literal(Literal::Integer(e)).into(),
+                ])
+                .into(),
+                (Some(_), Some(_)) => Item::SString(vec![
+                    InterpolateItem::Expr(Box::new(ident)),
+                    InterpolateItem::String(" BETWEEN ".to_string()),
+                    InterpolateItem::Expr(Box::new(Item::Range(range).into())),
+                ])
+                .into(),
+                (None, None) => Item::Literal(Literal::Boolean(true)).into(),
+            })),
+        ];
+        transforms
+            .into_iter()
+            .map(|t| Node {
+                is_complex: true, // this transform DOES contain windowed functions
+                ..Node::from(Item::Transform(t))
+            })
+            .collect()
+    }
+}
diff --git a/prql-compiler/src/sql/materializer.rs b/prql-compiler/src/sql/materializer.rs
index 3880fffb..f3c579ab 100644
--- a/prql-compiler/src/sql/materializer.rs
+++ b/prql-compiler/src/sql/materializer.rs
@@ -62,6 +62,7 @@ pub struct MaterializedFrame {
     pub sort: Vec<ColumnSort>,
 }
 
+#[derive(Default)]
 pub struct MaterializationContext {
     /// Map of all accessible names (for each namespace)
     pub(super) frame: Frame,
@@ -71,11 +72,13 @@ pub struct MaterializationContext {
 }
 
 impl MaterializationContext {
-    pub fn declare_table(&mut self, table: &str) -> usize {
-        let decl = Declaration::Table(table.to_string());
-        self.declarations.push(decl);
+    pub fn declare(&mut self, dec: Declaration) -> usize {
+        self.declarations.push(dec);
+        self.declarations.len() - 1
+    }
 
-        let id = self.declarations.len() - 1;
+    pub fn declare_table(&mut self, table: &str) -> usize {
+        let id = self.declare(Declaration::Table(table.to_string()));
 
         self.frame.tables.push(id);
         id
@@ -884,6 +887,7 @@ take 20
                   Literal:
                     Integer: 20
               by: []
+              sort: []
         "###);
         assert_yaml_snapshot!(frame.columns, @r###"
         ---
diff --git a/prql-compiler/src/sql/translator.rs b/prql-compiler/src/sql/translator.rs
index b9971649..55b269b7 100644
--- a/prql-compiler/src/sql/translator.rs
+++ b/prql-compiler/src/sql/translator.rs
@@ -253,7 +253,10 @@ 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>> {
+fn atomic_pipelines_of_pipeline(
+    pipeline: Pipeline,
+    context: &mut MaterializationContext,
+) -> Result<Vec<AtomicTable>> {
     // Insert a cut, when we find transformation that out of order:
     // - joins (no limit),
     // - filters (for WHERE)
@@ -268,7 +271,7 @@ fn atomic_pipelines_of_pipeline(pipeline: Pipeline) -> Result<Vec<AtomicTable>>
     // which we'll then compose together.
     let pipeline = Ok(pipeline.functions)
         .and_then(un_group::un_group)
-        .and_then(distinct::take_to_distinct)?;
+        .and_then(|x| distinct::take_to_distinct(x, context))?;
 
     let mut counts: HashMap<&str, u32> = HashMap::new();
     let mut splits = vec![0];
@@ -327,7 +330,7 @@ fn atomic_tables_of_tables(
     for table in tables {
         // split table into atomics
         let pipeline = table.pipeline.item.into_pipeline()?;
-        let mut t_atomics: Vec<_> = atomic_pipelines_of_pipeline(pipeline)?;
+        let mut t_atomics: Vec<_> = atomic_pipelines_of_pipeline(pipeline, context)?;
 
         let (last, ctes) = t_atomics
             .split_last_mut()
@@ -558,7 +561,7 @@ impl TryFrom<Item> for Expr {
                 let default_frame = if window.sort.is_empty() {
                     (WindowKind::Rows, Range::unbounded())
                 } else {
-                    (WindowKind::Range, Range::from_ints(None, Some(1)))
+                    (WindowKind::Range, Range::from_ints(None, Some(0)))
                 };
 
                 let window = WindowSpec {
@@ -931,6 +934,8 @@ SString:
 
     #[test]
     fn test_ctes_of_pipeline() -> Result<()> {
+        let mut context = MaterializationContext::default();
+
         // One aggregate, take at the end
         let prql: &str = r###"
         from employees
@@ -941,7 +946,7 @@ SString:
         "###;
 
         let pipeline = parse_and_resolve(prql)?;
-        let queries = atomic_pipelines_of_pipeline(pipeline)?;
+        let queries = atomic_pipelines_of_pipeline(pipeline, &mut context)?;
         assert_eq!(queries.len(), 1);
 
         // One aggregate, but take at the top
@@ -954,7 +959,7 @@ SString:
         "###;
 
         let pipeline = parse_and_resolve(prql)?;
-        let queries = atomic_pipelines_of_pipeline(pipeline)?;
+        let queries = atomic_pipelines_of_pipeline(pipeline, &mut context)?;
         assert_eq!(queries.len(), 2);
 
         // A take, then two aggregates
@@ -968,7 +973,7 @@ SString:
         "###;
 
         let pipeline = parse_and_resolve(prql)?;
-        let queries = atomic_pipelines_of_pipeline(pipeline)?;
+        let queries = atomic_pipelines_of_pipeline(pipeline, &mut context)?;
         assert_eq!(queries.len(), 3);
 
         // A take, then a select
@@ -979,7 +984,7 @@ SString:
         "###;
 
         let pipeline = parse_and_resolve(prql)?;
-        let queries = atomic_pipelines_of_pipeline(pipeline)?;
+        let queries = atomic_pipelines_of_pipeline(pipeline, &mut context)?;
         assert_eq!(queries.len(), 1);
         Ok(())
     }
@@ -2000,5 +2005,51 @@ take 20
         FROM
           employees
         "###);
+
+        // head
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        group department ( | take 3)
+        "###,
+        ).unwrap()).unwrap()), @r###"
+        WITH table_0 AS (
+          SELECT
+            employees.*,
+            ROW NUMBER() OVER (PARTITION BY department) AS _rn
+          FROM
+            employees
+        )
+        SELECT
+          table_0.*
+        FROM
+          table_0
+        WHERE
+          _rn <= 3
+        "###);
+
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from employees
+        group department ( | sort salary | take 2..3)
+        "###,
+        ).unwrap()).unwrap()), @r###"
+        WITH table_0 AS (
+          SELECT
+            employees.*,
+            ROW NUMBER() OVER (
+              PARTITION BY department
+              ORDER BY
+                salary
+            ) AS _rn
+          FROM
+            employees
+        )
+        SELECT
+          table_0.*
+        FROM
+          table_0
+        WHERE
+          _rn BETWEEN 2
+          AND 3
+        "###);
     }
 }
diff --git a/reference/src/language-features/distinct.md b/reference/src/language-features/distinct.md
index adb6197a..53a520cf 100644
--- a/reference/src/language-features/distinct.md
+++ b/reference/src/language-features/distinct.md
@@ -20,13 +20,14 @@ group department ( | take 1) # Note below
 
 > Note: `|` is here temporarily, until we finish work on function & pipeline currying.
 
-## Roadmap
 
-Soon we'll be able to [select range of rows in each
+-----
+
+We are be able to [select range of rows in each
 group](https://stackoverflow.com/questions/3800551/select-first-row-in-each-group-by-group)
 by combining `group` and `sort`:
 
-```prql_no_test
+```prql
 # youngest employee from each department
 from employees
 group department (
@@ -34,34 +35,23 @@ group department (
   take 1
 )
 ```
+## Roadmap
 
-... which will produces ...
+When using Postgres dialect, we are planning to compile:
 
-```sql
-WITH table_0 = (
-  SELECT ROW_NUMBER() OVER(PARTITION BY department ORDER_BY age) as _row_number
-  FROM employees
+```prql_no_test
+# youngest employee from each department
+from employees
+group department (
+  sort age
+  take 1
 )
-SELECT *
-FROM table_0
-WHERE _row_number = 1
 ```
 
-... or in Postgres dialect ...
+... to ...
 
 ```sql
 SELECT DISTINCT ON (department) *
 FROM employees
 ORDER BY department, age
 ```
-
-Until this is implemented, we can use `row_number`:
-
-```prql
-from employees
-group department (
-  sort age
-  derive rank = row_number
-)
-filter rank == 1
-```


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 6a46edce..5caf300e 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
@@ -117,4 +117,5 @@ expression: mat
           Literal:
             Integer: 20
       by: []
+      sort: []
 
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 ba9304f7..6da76710 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
@@ -117,4 +117,5 @@ expression: "&mat"
           Literal:
             Integer: 20
       by: []
+      sort: []
 
diff --git a/reference/tests/prql/language-features/distinct-2.prql b/reference/tests/prql/language-features/distinct-2.prql
index 7b3428d3..f224527e 100644
--- a/reference/tests/prql/language-features/distinct-2.prql
+++ b/reference/tests/prql/language-features/distinct-2.prql
@@ -1,6 +1,6 @@
+# youngest employee from each department
 from employees
 group department (
   sort age
-  derive rank = row_number
+  take 1
 )
-filter rank == 1
diff --git a/reference/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap b/reference/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap
index 5b025aaa..4ed761b8 100644
--- a/reference/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap
+++ b/reference/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap
@@ -6,12 +6,11 @@ input_file: reference/tests/prql/language-features/distinct-2.prql
 WITH table_0 AS (
   SELECT
     employees.*,
-    ROW_NUMBER() OVER (
+    ROW NUMBER() OVER (
       PARTITION BY department
       ORDER BY
-        age ROWS BETWEEN UNBOUNDED PRECEDING
-        AND UNBOUNDED FOLLOWING
-    ) AS rank
+        age
+    ) AS _rn
   FROM
     employees
 )
@@ -20,4 +19,4 @@ SELECT
 FROM
   table_0
 WHERE
-  rank = 1
+  _rn <= 1