fix(prqlc, append): avoid type mismatch with postgres (#5343)

Author: Fanaen

prqlc/prqlc/src/ir/rq/fold.rs

@@ -129,6 +129,7 @@ pub fn fold_table_ref<F: ?Sized + RqFold>(fold: &mut F, table_ref: TableRef) ->
                 Ok((fold.fold_relation_column(col)?, fold.fold_cid(cid)?))
             })
             .try_collect()?,
+        prefer_cte: table_ref.prefer_cte,
     })
 }
 

prqlc/prqlc/src/ir/rq/mod.rs

@@ -98,4 +98,7 @@ pub struct TableRef {
 
     /// Name hint for relation within this pipeline (table alias)
     pub name: Option<String>,
+
+    /// We prefer CTEs for most syntaxes but some like UNION works best with subqueries.
+    pub prefer_cte: bool,
 }

prqlc/prqlc/src/semantic/lowering.rs

@@ -471,6 +471,7 @@ impl Lowerer {
             source: tid,
             name,
             columns,
+            prefer_cte: true,
         }
     }
 
@@ -586,7 +587,8 @@ impl Lowerer {
                 self.pipeline.push(transform);
             }
             pl::TransformKind::Append(bottom) => {
-                let bottom = self.lower_table_ref(*bottom)?;
+                let mut bottom = self.lower_table_ref(*bottom)?;
+                bottom.prefer_cte = false;
 
                 self.pipeline.push(Transform::Append(bottom));
             }

prqlc/prqlc/src/semantic/mod.rs

@@ -253,6 +253,7 @@ pub mod test {
                     - - Wildcard
                       - 0
                   name: employees
+                  prefer_cte: true
               - Select:
                   - 0
           columns:

prqlc/prqlc/src/semantic/resolver/transforms.rs

@@ -1133,6 +1133,7 @@ mod tests {
                     - - Wildcard
                       - 1
                   name: c_invoice
+                  prefer_cte: true
               - Select:
                   - 0
               - Take:
@@ -1227,6 +1228,7 @@ mod tests {
                     - - Wildcard
                       - 3
                   name: invoices
+                  prefer_cte: true
               - Sort:
                   - direction: Asc
                     column: 0

prqlc/prqlc/src/sql/dialect.rs

@@ -213,6 +213,12 @@ pub(super) trait DialectHandler: Any + Debug {
             named: false,
         }))
     }
+
+    /// Whether source and subqueries should be put between simple parentheses
+    /// for `UNION` and similar verbs.
+    fn prefers_subquery_parentheses_shorthand(&self) -> bool {
+        false
+    }
 }
 
 impl dyn DialectHandler {
@@ -282,6 +288,10 @@ impl DialectHandler for PostgresDialect {
     fn supports_zero_columns(&self) -> bool {
         true
     }
+
+    fn prefers_subquery_parentheses_shorthand(&self) -> bool {
+        true
+    }
 }
 
 impl DialectHandler for GlareDbDialect {
@@ -474,6 +484,10 @@ impl DialectHandler for BigQueryDialect {
         // https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#set_operators
         true
     }
+
+    fn prefers_subquery_parentheses_shorthand(&self) -> bool {
+        true
+    }
 }
 
 impl DialectHandler for SnowflakeDialect {

prqlc/prqlc/src/sql/gen_query.rs

@@ -265,21 +265,37 @@ fn translate_set_ops_pipeline(
             _ => unreachable!(),
         };
 
-        // prepare top
+        // Some engines (like SQLite) do not support subqueries between simple parentheses, so we keep
+        // the general `SELECT * FROM (query)` or `SELECT * FROM cte` by default for complex cases.
+        //
+        // Some engines (like Postgres) need the subquery directly to properly match column type on
+        // both sides. For those:
+        // 1. we need the subquery as-is or in parentheses
+        // 2. we need to avoid CTEs
+
+        // Left hand side (aka top)
         let left = query_to_set_expr(top, context);
 
-        top = default_query(SetExpr::SetOperation {
-            left,
-            right: Box::new(SetExpr::Select(Box::new(sql_ast::Select {
+        // Right hand side (aka bottom)
+        let right_rel = translate_relation_expr(bottom, context)?;
+        let right = if let TableFactor::Derived { subquery, .. } = right_rel {
+            query_to_set_expr(*subquery, context)
+        } else {
+            Box::new(SetExpr::Select(Box::new(sql_ast::Select {
                 projection: vec![SelectItem::Wildcard(
                     sql_ast::WildcardAdditionalOptions::default(),
                 )],
                 from: vec![TableWithJoins {
-                    relation: translate_relation_expr(bottom, context)?,
+                    relation: right_rel,
                     joins: vec![],
                 }],
                 ..default_select()
-            }))),
+            })))
+        };
+
+        top = default_query(SetExpr::SetOperation {
+            left,
+            right,
             set_quantifier: if distinct {
                 if context.dialect.set_ops_distinct() {
                     sql_ast::SetQuantifier::Distinct
@@ -635,24 +651,36 @@ fn query_to_set_expr(query: sql_ast::Query, context: &mut Context) -> Box<SetExp
         return query.body;
     }
 
-    // query is not simple, so we need to wrap it into
-    // `SELECT * FROM (query)`
-    Box::new(SetExpr::Select(Box::new(Select {
-        projection: vec![SelectItem::Wildcard(
-            sql_ast::WildcardAdditionalOptions::default(),
-        )],
-        from: vec![TableWithJoins {
-            relation: TableFactor::Derived {
-                lateral: false,
-                subquery: Box::new(query),
-                alias: Some(simple_table_alias(sql_ast::Ident::new(
-                    context.anchor.table_name.gen(),
-                ))),
-            },
-            joins: vec![],
-        }],
-        ..default_select()
-    })))
+    // Query is not simple, so we need to wrap it.
+    //
+    // Some engines (like SQLite) do not support subqueries between simple parentheses, so we keep
+    // the general `SELECT * FROM (query)` by default.
+    //
+    // Some engines (like Postgres) may need the subquery directly as-is or between parentheses
+    // to properly match columns for syntaxes like `UNION`, `EXCEPT`, and `INTERSECT`.
+    // Incidentally, the parenthesis syntax `(query)` allows for complex queries.
+    let set_expr = if context.dialect.prefers_subquery_parentheses_shorthand() {
+        SetExpr::Query(query.into())
+    } else {
+        SetExpr::Select(Box::new(Select {
+            projection: vec![SelectItem::Wildcard(
+                sql_ast::WildcardAdditionalOptions::default(),
+            )],
+            from: vec![TableWithJoins {
+                relation: TableFactor::Derived {
+                    lateral: false,
+                    subquery: Box::new(query),
+                    alias: Some(simple_table_alias(sql_ast::Ident::new(
+                        context.anchor.table_name.gen(),
+                    ))),
+                },
+                joins: vec![],
+            }],
+            ..default_select()
+        }))
+    };
+
+    Box::new(set_expr)
 }
 
 fn count_tables(transforms: &[Transform]) -> usize {

prqlc/prqlc/src/sql/pq/anchor.rs

@@ -291,6 +291,7 @@ pub(super) fn anchor_split(
             source: new_tid,
             name: None,
             columns: new_columns,
+            prefer_cte: true,
         },
         cid_redirects,
     );

prqlc/prqlc/src/sql/pq/gen_query.rs

@@ -189,7 +189,7 @@ pub(super) fn compile_relation_instance(riid: RIId, ctx: &mut Context) -> Result
     // ensure that the table is declared
     if let RelationStatus::NotYetDefined(sql_relation) = decl.relation.take_to_define() {
         // if we cannot use CTEs (probably because we are within RECURSIVE)
-        if !ctx.query.allow_ctes {
+        if !(ctx.query.allow_ctes && table_ref.prefer_cte) {
             // restore relation for other references
             decl.relation = RelationStatus::NotYetDefined(sql_relation.clone());
 

prqlc/prqlc/tests/integration/queries/append_select_nulls.prql

@@ -0,0 +1,8 @@
+from invoices
+select {an_id = invoice_id, name = null}
+take 2
+append (
+  from employees
+  select {an_id = null, name = first_name}
+  take 2
+)