fix(postprocess): when redirecting sort columns, look up the relation from the matching table_ref (#5486)

Signed-off-by: Luka Peschke <[email protected]>

Author: lukapeschke

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

@@ -209,15 +209,26 @@ impl PqFold for SortingInference<'_> {
                             cid_redirects_to_add.push((*old_cid, new_cid, col));
                         }
 
-                        let riid = RIId::from(cte.tid.get());
-                        let relation_instance =
-                            self.ctx.anchor.relation_instances.get_mut(&riid).unwrap();
+                        let (riid, relation_instance) = self
+                            .ctx
+                            .anchor
+                            .relation_instances
+                            .iter_mut()
+                            .find(|(_riid, rel_inst)| rel_inst.table_ref.source == cte.tid)
+                            .unwrap();
 
                         cid_redirects_to_add
                             .into_iter()
                             .for_each(|(old_cid, new_cid, col)| {
-                                let def = ColumnDecl::RelationColumn(riid, new_cid, col);
+                                let def = ColumnDecl::RelationColumn(*riid, new_cid, col);
+
                                 self.ctx.anchor.column_decls.insert(new_cid, def);
+                                log::debug!(
+                                    "-- redirecting {old_cid:?} to {new_cid:?} for CTE {cte:?} (RIId: {riid:?})",
+                                    cte = cte.tid,
+                                    riid = riid
+                                );
+
                                 relation_instance.cid_redirects.insert(old_cid, new_cid);
                             });
                     }

prqlc/prqlc/tests/integration/sql.rs

@@ -6115,6 +6115,55 @@ fn test_sort_cast_filter_join_select() {
     ")
 }
 
+/// Ensures that the sort happens on `table1`.`artist_id` and not on `table2`.`artist_id`
+#[test]
+fn test_sort_filter_derive_join_select() {
+    assert_snapshot!(compile(r###"
+    from albums
+    sort this.`artist_id`
+    filter (this.`artist_id` != null)
+    derive { `artist_id` = as `double precision` this.`artist_id` }
+    join side:left (from artists
+        select {`artist_id_right` = this.`artist_id`}
+    ) (this.`artist_id` == that.`artist_id_right`)
+    select {this.`artist_id`, this.`title`}
+    "###
+    ).unwrap(), @r"
+    WITH table_2 AS (
+      SELECT
+        CAST(artist_id AS double precision) AS artist_id,
+        title,
+        artist_id AS _expr_0
+      FROM
+        albums
+      WHERE
+        artist_id IS NOT NULL
+    ),
+    table_1 AS (
+      SELECT
+        artist_id,
+        title,
+        _expr_0
+      FROM
+        table_2
+    ),
+    table_0 AS (
+      SELECT
+        artist_id AS artist_id_right
+      FROM
+        artists
+    )
+    SELECT
+      table_1.artist_id,
+      table_1.title
+    FROM
+      table_1
+      LEFT OUTER JOIN table_0 ON table_1.artist_id = table_0.artist_id_right
+    ORDER BY
+      table_1._expr_0
+    ")
+}
+
 /// Ensures that the sort happens on `table0`.`artist_id` and not on `table0`.`double_artist_id`
 #[test]
 fn test_sort_cast_filter_join_select_with_alias() {