fix(sql): support set operations wrapping subqueries

jcrist · Mar 6, 2024 · 3736686 · 3736686
1 parent 6ed2e39
commit 3736686
Show file tree

Hide file tree

Showing 18 changed files with 282 additions and 156 deletions.
diff --git a/ibis/backends/postgres/tests/snapshots/test_functions/test_union_cte/False/out.sql b/ibis/backends/postgres/tests/snapshots/test_functions/test_union_cte/False/out.sql
@@ -1 +1 @@
-WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION ALL SELECT * FROM "t1" AS "t4" ) AS "t5" UNION ALL SELECT * FROM "t1" AS "t3" ) AS "t7"
+WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT * FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION ALL SELECT * FROM "t1" AS "t4" ) AS "t5" ) AS "t6" UNION ALL SELECT * FROM "t1" AS "t3" ) AS "t7"
diff --git a/ibis/backends/postgres/tests/snapshots/test_functions/test_union_cte/True/out.sql b/ibis/backends/postgres/tests/snapshots/test_functions/test_union_cte/True/out.sql
@@ -1 +1 @@
-WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION SELECT * FROM "t1" AS "t4" ) AS "t5" UNION SELECT * FROM "t1" AS "t3" ) AS "t7"
+WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT * FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION SELECT * FROM "t1" AS "t4" ) AS "t5" ) AS "t6" UNION SELECT * FROM "t1" AS "t3" ) AS "t7"
diff --git a/ibis/backends/risingwave/tests/snapshots/test_functions/test_union_cte/False/out.sql b/ibis/backends/risingwave/tests/snapshots/test_functions/test_union_cte/False/out.sql
@@ -1 +1 @@
-WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION ALL SELECT * FROM "t1" AS "t4" ) AS "t5" UNION ALL SELECT * FROM "t1" AS "t3" ) AS "t7"
+WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT * FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION ALL SELECT * FROM "t1" AS "t4" ) AS "t5" ) AS "t6" UNION ALL SELECT * FROM "t1" AS "t3" ) AS "t7"
diff --git a/ibis/backends/risingwave/tests/snapshots/test_functions/test_union_cte/True/out.sql b/ibis/backends/risingwave/tests/snapshots/test_functions/test_union_cte/True/out.sql
@@ -1 +1 @@
-WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION SELECT * FROM "t1" AS "t4" ) AS "t5" UNION SELECT * FROM "t1" AS "t3" ) AS "t7"
+WITH "t1" AS ( SELECT "t0"."string_col", SUM("t0"."double_col") AS "metric" FROM "functional_alltypes" AS "t0" GROUP BY 1 ) SELECT "t7"."string_col", "t7"."metric" FROM ( SELECT * FROM ( SELECT "t5"."string_col", "t5"."metric" FROM ( SELECT * FROM "t1" AS "t2" UNION SELECT * FROM "t1" AS "t4" ) AS "t5" ) AS "t6" UNION SELECT * FROM "t1" AS "t3" ) AS "t7"
diff --git a/ibis/backends/sql/compiler.py b/ibis/backends/sql/compiler.py
@@ -1217,10 +1217,10 @@ def visit_Sort(self, op, *, parent, keys):
             return sg.select(STAR).from_(parent).order_by(*keys)
 
     def visit_Union(self, op, *, left, right, distinct):
-        if isinstance(left, sge.Table):
+        if isinstance(left, (sge.Table, sge.Subquery)):
             left = sg.select(STAR).from_(left)
 
-        if isinstance(right, sge.Table):
+        if isinstance(right, (sge.Table, sge.Subquery)):
             right = sg.select(STAR).from_(right)
 
         return sg.union(
@@ -1230,10 +1230,10 @@ def visit_Union(self, op, *, left, right, distinct):
         )
 
     def visit_Intersection(self, op, *, left, right, distinct):
-        if isinstance(left, sge.Table):
+        if isinstance(left, (sge.Table, sge.Subquery)):
             left = sg.select(STAR).from_(left)
 
-        if isinstance(right, sge.Table):
+        if isinstance(right, (sge.Table, sge.Subquery)):
             right = sg.select(STAR).from_(right)
 
         return sg.intersect(
@@ -1243,10 +1243,10 @@ def visit_Intersection(self, op, *, left, right, distinct):
         )
 
     def visit_Difference(self, op, *, left, right, distinct):
-        if isinstance(left, sge.Table):
+        if isinstance(left, (sge.Table, sge.Subquery)):
             left = sg.select(STAR).from_(left)
 
-        if isinstance(right, sge.Table):
+        if isinstance(right, (sge.Table, sge.Subquery)):
             right = sg.select(STAR).from_(right)
 
         return sg.except_(

diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/clickhouse/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/clickhouse/out.sql
@@ -41,20 +41,28 @@ SELECT
   "t11"."diff"
 FROM (
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  ORDER BY
-    "t6"."diff" DESC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    ORDER BY
+      "t6"."diff" DESC
+    LIMIT 10
+  ) AS "t9"
   UNION ALL
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  WHERE
-    "t6"."diff" < 0
-  ORDER BY
-    "t6"."diff" ASC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    WHERE
+      "t6"."diff" < 0
+    ORDER BY
+      "t6"."diff" ASC
+    LIMIT 10
+  ) AS "t10"
 ) AS "t11"
diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/duckdb/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/duckdb/out.sql
@@ -41,20 +41,28 @@ SELECT
   "t11"."diff"
 FROM (
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  ORDER BY
-    "t6"."diff" DESC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    ORDER BY
+      "t6"."diff" DESC
+    LIMIT 10
+  ) AS "t9"
   UNION ALL
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  WHERE
-    "t6"."diff" < CAST(0 AS TINYINT)
-  ORDER BY
-    "t6"."diff" ASC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    WHERE
+      "t6"."diff" < CAST(0 AS TINYINT)
+    ORDER BY
+      "t6"."diff" ASC
+    LIMIT 10
+  ) AS "t10"
 ) AS "t11"
diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/postgres/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/postgres/out.sql
@@ -41,20 +41,28 @@ SELECT
   "t11"."diff"
 FROM (
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  ORDER BY
-    "t6"."diff" DESC NULLS LAST
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    ORDER BY
+      "t6"."diff" DESC NULLS LAST
+    LIMIT 10
+  ) AS "t9"
   UNION ALL
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  WHERE
-    "t6"."diff" < 0
-  ORDER BY
-    "t6"."diff" ASC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    WHERE
+      "t6"."diff" < 0
+    ORDER BY
+      "t6"."diff" ASC
+    LIMIT 10
+  ) AS "t10"
 ) AS "t11"
diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/pyspark/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/pyspark/out.sql
@@ -60,20 +60,28 @@ SELECT
   `t11`.`diff`
 FROM (
   SELECT
-    `t6`.`field_of_study`,
-    `t6`.`diff`
-  FROM `t5` AS `t6`
-  ORDER BY
-    `t6`.`diff` DESC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      `t6`.`field_of_study`,
+      `t6`.`diff`
+    FROM `t5` AS `t6`
+    ORDER BY
+      `t6`.`diff` DESC
+    LIMIT 10
+  ) AS `t9`
   UNION ALL
   SELECT
-    `t6`.`field_of_study`,
-    `t6`.`diff`
-  FROM `t5` AS `t6`
-  WHERE
-    `t6`.`diff` < 0
-  ORDER BY
-    `t6`.`diff` ASC NULLS LAST
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      `t6`.`field_of_study`,
+      `t6`.`diff`
+    FROM `t5` AS `t6`
+    WHERE
+      `t6`.`diff` < 0
+    ORDER BY
+      `t6`.`diff` ASC NULLS LAST
+    LIMIT 10
+  ) AS `t10`
 ) AS `t11`
diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/risingwave/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/risingwave/out.sql
@@ -41,20 +41,28 @@ SELECT
   "t11"."diff"
 FROM (
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  ORDER BY
-    "t6"."diff" DESC NULLS LAST
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    ORDER BY
+      "t6"."diff" DESC NULLS LAST
+    LIMIT 10
+  ) AS "t9"
   UNION ALL
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  WHERE
-    "t6"."diff" < 0
-  ORDER BY
-    "t6"."diff" ASC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    WHERE
+      "t6"."diff" < 0
+    ORDER BY
+      "t6"."diff" ASC
+    LIMIT 10
+  ) AS "t10"
 ) AS "t11"
diff --git a/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/trino/out.sql b/ibis/backends/tests/snapshots/test_sql/test_union_aliasing/trino/out.sql
@@ -58,20 +58,28 @@ SELECT
   "t11"."diff"
 FROM (
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  ORDER BY
-    "t6"."diff" DESC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    ORDER BY
+      "t6"."diff" DESC
+    LIMIT 10
+  ) AS "t9"
   UNION ALL
   SELECT
-    "t6"."field_of_study",
-    "t6"."diff"
-  FROM "t5" AS "t6"
-  WHERE
-    "t6"."diff" < 0
-  ORDER BY
-    "t6"."diff" ASC
-  LIMIT 10
+    *
+  FROM (
+    SELECT
+      "t6"."field_of_study",
+      "t6"."diff"
+    FROM "t5" AS "t6"
+    WHERE
+      "t6"."diff" < 0
+    ORDER BY
+      "t6"."diff" ASC
+    LIMIT 10
+  ) AS "t10"
 ) AS "t11"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql
@@ -2,16 +2,24 @@ SELECT
   "t3"."key"
 FROM (
   SELECT
-    "t0"."string_col" AS "key",
-    CAST("t0"."float_col" AS DOUBLE) AS "value"
-  FROM "functional_alltypes" AS "t0"
-  WHERE
-    "t0"."int_col" > CAST(0 AS TINYINT)
+    *
+  FROM (
+    SELECT
+      "t0"."string_col" AS "key",
+      CAST("t0"."float_col" AS DOUBLE) AS "value"
+    FROM "functional_alltypes" AS "t0"
+    WHERE
+      "t0"."int_col" > CAST(0 AS TINYINT)
+  ) AS "t1"
   EXCEPT
   SELECT
-    "t0"."string_col" AS "key",
-    "t0"."double_col" AS "value"
-  FROM "functional_alltypes" AS "t0"
-  WHERE
-    "t0"."int_col" <= CAST(0 AS TINYINT)
+    *
+  FROM (
+    SELECT
+      "t0"."string_col" AS "key",
+      "t0"."double_col" AS "value"
+    FROM "functional_alltypes" AS "t0"
+    WHERE
+      "t0"."int_col" <= CAST(0 AS TINYINT)
+  ) AS "t2"
 ) AS "t3"
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql
@@ -2,16 +2,24 @@ SELECT
   "t3"."key"
 FROM (
   SELECT
-    "t0"."string_col" AS "key",
-    CAST("t0"."float_col" AS DOUBLE) AS "value"
-  FROM "functional_alltypes" AS "t0"
-  WHERE
-    "t0"."int_col" > CAST(0 AS TINYINT)
+    *
+  FROM (
+    SELECT
+      "t0"."string_col" AS "key",
+      CAST("t0"."float_col" AS DOUBLE) AS "value"
+    FROM "functional_alltypes" AS "t0"
+    WHERE
+      "t0"."int_col" > CAST(0 AS TINYINT)
+  ) AS "t1"
   INTERSECT
   SELECT
-    "t0"."string_col" AS "key",
-    "t0"."double_col" AS "value"
-  FROM "functional_alltypes" AS "t0"
-  WHERE
-    "t0"."int_col" <= CAST(0 AS TINYINT)
+    *
+  FROM (
+    SELECT
+      "t0"."string_col" AS "key",
+      "t0"."double_col" AS "value"
+    FROM "functional_alltypes" AS "t0"
+    WHERE
+      "t0"."int_col" <= CAST(0 AS TINYINT)
+  ) AS "t2"
 ) AS "t3"