fix: Prevent DISTINCT ON internal sorting from leaking past joins (#5633)

Co-authored-by: Claude <[email protected]>

Author: max-sixty

prqlc/prqlc/src/lib.rs

@@ -614,4 +614,117 @@ mod tests {
             .map(|name| Target::from_str(&name))
             .collect();
     }
+
+    /// Regression test for #4633: sort inside group should not leak to outer query after join.
+    ///
+    /// Per PRQL spec, `group` resets the order. The `sort` inside a group is for
+    /// row selection (which row to keep), not output ordering. After the group,
+    /// there is no defined order, so it should not appear in the outer query.
+    #[test]
+    fn test_sort_not_propagated_after_join() {
+        use insta::assert_snapshot;
+
+        // DISTINCT ON (postgres) uses the sort for row selection within the CTE.
+        // This internal sorting must not leak to the outer query after a join.
+        assert_snapshot!(
+            super::compile(
+                r#"
+                prql target:sql.postgres
+
+                from tracks
+                group media_type_id (
+                    sort name
+                    take 1
+                )
+                join media_types (== media_type_id)
+                select {
+                    tracks.track_id,
+                    media_types.name
+                }
+                "#,
+                &super::Options::default().no_signature()
+            ).unwrap(),
+            @"
+        WITH table_0 AS (
+          SELECT
+            DISTINCT ON (media_type_id) track_id,
+            media_type_id,
+            name
+          FROM
+            tracks
+          ORDER BY
+            media_type_id,
+            name
+        )
+        SELECT
+          table_0.track_id,
+          media_types.name
+        FROM
+          table_0
+          INNER JOIN media_types ON table_0.media_type_id = media_types.media_type_id
+        "
+        );
+    }
+
+    /// Verify that explicit sorts after group are preserved past joins.
+    ///
+    /// Per PRQL spec, `sort` introduces a new order. When the user explicitly
+    /// sorts AFTER a group, that becomes the new output order and should
+    /// propagate through subsequent transforms including joins.
+    #[test]
+    fn test_explicit_sort_after_distinct_on_preserved() {
+        use insta::assert_snapshot;
+
+        // Explicit `sort media_type_id` after the group introduces a new order.
+        // This user-requested ordering should propagate past the join.
+        assert_snapshot!(
+            super::compile(
+                r#"
+                prql target:sql.postgres
+
+                from tracks
+                group media_type_id (
+                    sort name
+                    take 1
+                )
+                sort media_type_id
+                join media_types (== media_type_id)
+                select {
+                    tracks.track_id,
+                    media_types.name
+                }
+                "#,
+                &super::Options::default().no_signature()
+            ).unwrap(),
+            @"
+        WITH table_0 AS (
+          SELECT
+            DISTINCT ON (media_type_id) track_id,
+            media_type_id,
+            name
+          FROM
+            tracks
+          ORDER BY
+            media_type_id,
+            name
+        ),
+        table_1 AS (
+          SELECT
+            table_0.track_id,
+            media_types.name,
+            table_0.media_type_id
+          FROM
+            table_0
+            INNER JOIN media_types ON table_0.media_type_id = media_types.media_type_id
+        )
+        SELECT
+          track_id,
+          name
+        FROM
+          table_1
+        ORDER BY
+          media_type_id
+        "
+        );
+    }
 }

prqlc/prqlc/src/sql/pq/postprocess.rs

@@ -27,6 +27,7 @@ pub(super) fn postprocess(query: SqlQuery, ctx: &mut Context) -> SqlQuery {
 fn infer_sorts(query: SqlQuery, ctx: &mut Context) -> SqlQuery {
     let mut s = SortingInference {
         last_sorting: Vec::new(),
+        last_sorting_from_distinct_on: false,
         ctes_sorting: HashMap::new(),
         main_relation: false,
         ctx,
@@ -37,6 +38,12 @@ fn infer_sorts(query: SqlQuery, ctx: &mut Context) -> SqlQuery {
 
 struct SortingInference<'a> {
     last_sorting: Sorting,
+    /// True if last_sorting originated from DISTINCT ON (used for row selection).
+    ///
+    /// Per the PRQL spec, `group` resets the order - any `sort` inside a group
+    /// is for internal row selection, not output ordering. This flag tracks such
+    /// internal sorting so it doesn't propagate past transforms like `join`.
+    last_sorting_from_distinct_on: bool,
     ctes_sorting: HashMap<TId, CteSorting>,
     main_relation: bool,
     ctx: &'a mut Context,
@@ -148,6 +155,12 @@ impl SortingInference<'_> {
 #[derive(Debug)]
 struct CteSorting {
     sorting: Sorting,
+    /// True if the CTE's sorting originated from DISTINCT ON (row selection).
+    ///
+    /// Per the PRQL spec, `group` resets the order. DISTINCT ON sorting is
+    /// internal to the group - it determines which row to keep, not output
+    /// ordering. This flag ensures such sorting doesn't leak to outer queries.
+    from_distinct_on: bool,
 }
 
 impl RqFold for SortingInference<'_> {}
@@ -239,7 +252,11 @@ impl PqFold for SortingInference<'_> {
             // store sorting to be used later in From references
             let sorting = self.last_sorting.drain(..).collect();
             log::debug!("--- sorting {sorting:?}");
-            let sorting = CteSorting { sorting };
+            let sorting = CteSorting {
+                sorting,
+                from_distinct_on: self.last_sorting_from_distinct_on,
+            };
+            self.last_sorting_from_distinct_on = false;
             self.ctes_sorting.insert(cte.tid, sorting);
 
             ctes.push(cte);
@@ -305,6 +322,9 @@ impl PqMapper<RelationExpr, RelationExpr, (), ()> for SortingInference<'_> {
         transforms: Vec<SqlTransform<RelationExpr, ()>>,
     ) -> Result<Vec<SqlTransform<RelationExpr, ()>>> {
         let mut sorting = Vec::new();
+        // Track whether sorting originated from DISTINCT ON (internal row selection).
+        // Per PRQL spec, `group` resets order - internal sorts don't define output order.
+        let mut sorting_from_distinct_on = false;
 
         let mut result = Vec::with_capacity(transforms.len() + 1);
 
@@ -314,17 +334,20 @@ impl PqMapper<RelationExpr, RelationExpr, (), ()> for SortingInference<'_> {
                     match expr.kind {
                         RelationExprKind::Ref(ref tid) => {
                             // infer sorting from referenced pipeline
-                            if let Some(cte_sorting) = self.ctes_sorting.get_mut(tid) {
+                            if let Some(cte_sorting) = self.ctes_sorting.get(tid) {
                                 sorting.clone_from(&cte_sorting.sorting);
+                                sorting_from_distinct_on = cte_sorting.from_distinct_on;
                             } else {
-                                sorting = Vec::new();
+                                sorting.clear();
+                                sorting_from_distinct_on = false;
                             };
                         }
                         RelationExprKind::SubQuery(rel) => {
                             let rel = self.fold_sql_relation(rel)?;
 
                             // infer sorting from sub-query
                             sorting = self.last_sorting.drain(..).collect();
+                            sorting_from_distinct_on = self.last_sorting_from_distinct_on;
 
                             expr.kind = RelationExprKind::SubQuery(rel);
                         }
@@ -337,16 +360,40 @@ impl PqMapper<RelationExpr, RelationExpr, (), ()> for SortingInference<'_> {
                 // just store sorting and don't emit Sort
                 SqlTransform::Sort(expr) => {
                     sorting.clone_from(&expr);
+                    // A new explicit Sort clears the DISTINCT ON flag - this is a
+                    // user-requested ordering, not an internal DISTINCT ON sort.
+                    sorting_from_distinct_on = false;
                     continue;
                 }
 
                 // clear sorting
                 SqlTransform::Distinct | SqlTransform::Aggregate { .. } => {
-                    sorting = Vec::new();
+                    sorting.clear();
+                    sorting_from_distinct_on = false;
+                }
+
+                // Per PRQL spec: `group` resets order, `join` retains left's order.
+                // DISTINCT ON sorting is internal to the group (for row selection),
+                // so it must not propagate past joins. Explicit user sorts are preserved.
+                // See issue #4633.
+                SqlTransform::Join { .. } => {
+                    if sorting_from_distinct_on {
+                        sorting.clear();
+                        sorting_from_distinct_on = false;
+                    }
                 }
 
                 // emit Sort before Take
-                SqlTransform::Take(_) | SqlTransform::DistinctOn(_) => {
+                SqlTransform::Take(_) => {
+                    result.push(SqlTransform::Sort(sorting.clone()));
+                }
+
+                SqlTransform::DistinctOn(_) => {
+                    // DISTINCT ON uses sorting for row selection within the group.
+                    // Mark it so this internal sorting doesn't leak to outer queries.
+                    // Note: ROW_NUMBER (used for take > 1 or non-Postgres) doesn't have
+                    // this issue because its sorting is embedded in the window function.
+                    sorting_from_distinct_on = true;
                     result.push(SqlTransform::Sort(sorting.clone()));
                 }
                 _ => {}
@@ -371,6 +418,7 @@ impl PqMapper<RelationExpr, RelationExpr, (), ()> for SortingInference<'_> {
 
         // remember sorting for this pipeline
         self.last_sorting = sorting;
+        self.last_sorting_from_distinct_on = sorting_from_distinct_on;
 
         Ok(result)
     }

prqlc/prqlc/tests/integration/sql.rs

@@ -3299,7 +3299,7 @@ fn test_prql_to_sql_table() {
     "#;
     let sql = compile(query).unwrap();
     assert_snapshot!(sql,
-        @r"
+        @"
     WITH newest_employees AS (
       SELECT
         *
@@ -6338,7 +6338,7 @@ fn test_sort_cast_filter_join_select() {
     join side:left artists (this.`artist_id` == that.`artist_id`)
     select {this.`artist_id`, this.`title`, this.`name`}
     "###
-    ).unwrap(), @r"
+    ).unwrap(), @"
     WITH table_1 AS (
       SELECT
         CAST(artist_id AS double precision) AS artist_id,
@@ -6394,7 +6394,7 @@ fn test_sort_filter_derive_join_select() {
     ) (this.`artist_id` == that.`artist_id_right`)
     select {this.`artist_id`, this.`title`}
     "###
-    ).unwrap(), @r"
+    ).unwrap(), @"
     WITH table_2 AS (
       SELECT
         CAST(artist_id AS double precision) AS artist_id,
@@ -6442,7 +6442,7 @@ fn test_sort_cast_filter_join_select_with_alias() {
     join side:left artists (this.`double_artist_id` == that.`artist_id`)
     select {this.`double_artist_id`, this.`title`, this.`name`}
     "###
-    ).unwrap(), @r"
+    ).unwrap(), @"
     WITH table_1 AS (
       SELECT
         CAST(artist_id AS double precision) AS double_artist_id,