From 373668634526f7204662dc0e5d48309b91ca74ee Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Wed, 21 Feb 2024 15:10:28 -0600 Subject: [PATCH] fix(sql): support set operations wrapping subqueries --- .../test_union_cte/False/out.sql | 2 +- .../test_union_cte/True/out.sql | 2 +- .../test_union_cte/False/out.sql | 2 +- .../test_union_cte/True/out.sql | 2 +- ibis/backends/sql/compiler.py | 12 +++---- .../test_union_aliasing/clickhouse/out.sql | 36 +++++++++++-------- .../test_union_aliasing/duckdb/out.sql | 36 +++++++++++-------- .../test_union_aliasing/postgres/out.sql | 36 +++++++++++-------- .../test_union_aliasing/pyspark/out.sql | 36 +++++++++++-------- .../test_union_aliasing/risingwave/out.sql | 36 +++++++++++-------- .../test_union_aliasing/trino/out.sql | 36 +++++++++++-------- .../test_difference_project_column/out.sql | 28 +++++++++------ .../test_intersect_project_column/out.sql | 28 +++++++++------ .../test_table_difference/out.sql | 28 +++++++++------ .../test_table_intersect/out.sql | 28 +++++++++------ .../test_compiler/test_union/out.sql | 28 +++++++++------ .../test_union_project_column/out.sql | 28 +++++++++------ ibis/backends/tests/test_set_ops.py | 34 ++++++++++++++++-- 18 files changed, 282 insertions(+), 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 index f0366d83444de..0f8848f4142e2 100644 --- 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" \ No newline at end of file +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" \ No newline at end of file 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 index 5a873785e92bb..59496e7819dce 100644 --- 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" \ No newline at end of file +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" \ No newline at end of file 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 index f0366d83444de..0f8848f4142e2 100644 --- 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" \ No newline at end of file +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" \ No newline at end of file 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 index 5a873785e92bb..59496e7819dce 100644 --- 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" \ No newline at end of file +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" \ No newline at end of file diff --git a/ibis/backends/sql/compiler.py b/ibis/backends/sql/compiler.py index d4568d768c744..654d926c4ac51 100644 --- 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 index 3078c2ff27e37..fd5060ef3e99a 100644 --- 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" \ No newline at end of file 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 index c86e858ede00b..147f962504157 100644 --- 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" \ No newline at end of file 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 index b7508b9ef535f..bfbba92e1f982 100644 --- 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" \ No newline at end of file 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 index 42c725d213617..61bb0090ec9b4 100644 --- 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` \ No newline at end of file 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 index b7508b9ef535f..bfbba92e1f982 100644 --- 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" \ No newline at end of file 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 index ad38eea501b96..c1dca5e222011 100644 --- 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" \ No newline at end of file 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 index a3eb8d979926f..3ac830ba92b30 100644 --- 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" \ No newline at end of file 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 index f9f8dec986987..07c91b224cd0b 100644 --- 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" \ No newline at end of file diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_difference/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_difference/out.sql index fff41a520ec48..fd1d141034256 100644 --- a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_difference/out.sql +++ b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_difference/out.sql @@ -3,16 +3,24 @@ SELECT "t3"."value" 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" \ No newline at end of file diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql index d7b9f61d48dcc..28a01f0da0cb5 100644 --- a/ibis/backends/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql +++ b/ibis/backends/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql @@ -3,16 +3,24 @@ SELECT "t3"."value" 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" \ No newline at end of file diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_union/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_union/out.sql index cf49db40dbeae..ed2ff926dd274 100644 --- a/ibis/backends/tests/sql/snapshots/test_compiler/test_union/out.sql +++ b/ibis/backends/tests/sql/snapshots/test_compiler/test_union/out.sql @@ -3,16 +3,24 @@ SELECT "t3"."value" 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" UNION 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" \ No newline at end of file diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql index b57144bd31c98..4efd945170795 100644 --- a/ibis/backends/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql +++ b/ibis/backends/tests/sql/snapshots/test_compiler/test_union_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" UNION ALL 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" \ No newline at end of file diff --git a/ibis/backends/tests/test_set_ops.py b/ibis/backends/tests/test_set_ops.py index 4df076da7f978..77d9b87112dfa 100644 --- a/ibis/backends/tests/test_set_ops.py +++ b/ibis/backends/tests/test_set_ops.py @@ -192,9 +192,24 @@ def test_table_set_operations_api(alltypes, method): False, ], ) -def test_top_level_union(backend, con, alltypes, distinct): +@pytest.mark.parametrize( + "ordered", + [ + False, + param( + True, + marks=pytest.mark.notyet( + ["mssql"], reason="ORDER BY not supported in subquery" + ), + ), + ], +) +def test_top_level_union(backend, con, alltypes, distinct, ordered): t1 = alltypes.select(a="bigint_col").filter(lambda t: t.a == 10).distinct() t2 = alltypes.select(a="bigint_col").filter(lambda t: t.a == 20).distinct() + if ordered: + t1 = t1.order_by("a") + t2 = t2.order_by("a") expr = t1.union(t2, distinct=distinct).limit(2) result = con.execute(expr) expected = pd.DataFrame({"a": [10, 20]}) @@ -237,10 +252,22 @@ def test_top_level_union(backend, con, alltypes, distinct): ], ids=["intersect", "difference"], ) +@pytest.mark.parametrize( + "ordered", + [ + False, + param( + True, + marks=pytest.mark.notyet( + ["mssql"], reason="ORDER BY not supported in subquery" + ), + ), + ], +) @pytest.mark.notimpl(["polars"], raises=com.OperationNotDefinedError) @pytest.mark.broken(["druid"], raises=PyDruidProgrammingError) def test_top_level_intersect_difference( - backend, con, alltypes, distinct, opname, expected + backend, con, alltypes, distinct, opname, expected, ordered ): t1 = ( alltypes.select(a="bigint_col") @@ -252,6 +279,9 @@ def test_top_level_intersect_difference( .filter(lambda t: (t.a == 20) | (t.a == 30)) .distinct() ) + if ordered: + t1 = t1.order_by("a") + t2 = t2.order_by("a") op = getattr(t1, opname) expr = op(t2, distinct=distinct).limit(2) result = con.execute(expr)