fix(ir): avoid traversing into any subqueries

ibis-project · Jan 1, 2024 · c625487 · c625487
1 parent 506296e
commit c625487
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 54 deletions.
diff --git a/ibis/backends/base/sqlglot/rewrites.py b/ibis/backends/base/sqlglot/rewrites.py
@@ -104,7 +104,7 @@ def merge_select_select(_):
         if v.find((Window, ops.Unnest), filter=ops.Value):
             return _
     for v in _.predicates:
-        if v.find(ops.ExistsSubquery, filter=ops.Value):
+        if v.find(ops.Subquery, filter=ops.Value):
             return _
 
     subs = {ops.Field(_.parent, k): v for k, v in _.parent.values.items()}

diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_and_non_agg_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_and_non_agg_filter/out.sql
@@ -1,10 +1,16 @@
 SELECT
-  t0.a,
-  t0.b
-FROM my_table AS t0
+  t1.a,
+  t1.b
+FROM (
+  SELECT
+    t0.a,
+    t0.b
+  FROM my_table AS t0
+  WHERE
+    t0.a < CAST(100 AS TINYINT)
+) AS t1
 WHERE
-  t0.a < CAST(100 AS TINYINT)
-  AND t0.a = (
+  t1.a = (
     SELECT
       MAX(t1.a) AS "Max(a)"
     FROM (
@@ -16,4 +22,4 @@ WHERE
         t0.a < CAST(100 AS TINYINT)
     ) AS t1
   )
-  AND t0.b = 'a'
+  AND t1.b = 'a'
diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter/out.sql
@@ -1,10 +1,16 @@
 SELECT
-  t0.a,
-  t0.b * CAST(2 AS TINYINT) AS b2
-FROM my_table AS t0
+  t1.a,
+  t1.b2
+FROM (
+  SELECT
+    t0.a,
+    t0.b * CAST(2 AS TINYINT) AS b2
+  FROM my_table AS t0
+  WHERE
+    t0.a < CAST(100 AS TINYINT)
+) AS t1
 WHERE
-  t0.a < CAST(100 AS TINYINT)
-  AND t0.a = (
+  t1.a = (
     SELECT
       MAX(t1.a) AS "Max(a)"
     FROM (

diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter_with_alias/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_agg_filter_with_alias/out.sql
@@ -1,10 +1,16 @@
 SELECT
-  t0.a,
-  t0.b * CAST(2 AS TINYINT) AS b2
-FROM my_table AS t0
+  t1.a,
+  t1.b2
+FROM (
+  SELECT
+    t0.a,
+    t0.b * CAST(2 AS TINYINT) AS b2
+  FROM my_table AS t0
+  WHERE
+    t0.a < CAST(100 AS TINYINT)
+) AS t1
 WHERE
-  t0.a < CAST(100 AS TINYINT)
-  AND t0.a = (
+  t1.a = (
     SELECT
       MAX(t1.a) AS "Max(a)"
     FROM (

diff --git a/ibis/backends/tests/sql/snapshots/test_compiler/test_simple_agg_filter/out.sql b/ibis/backends/tests/sql/snapshots/test_compiler/test_simple_agg_filter/out.sql
@@ -1,10 +1,16 @@
 SELECT
-  t0.a,
-  t0.b
-FROM my_table AS t0
+  t1.a,
+  t1.b
+FROM (
+  SELECT
+    t0.a,
+    t0.b
+  FROM my_table AS t0
+  WHERE
+    t0.a < CAST(100 AS TINYINT)
+) AS t1
 WHERE
-  t0.a < CAST(100 AS TINYINT)
-  AND t0.a = (
+  t1.a = (
     SELECT
       MAX(t1.a) AS "Max(a)"
     FROM (

diff --git a/ibis/backends/tests/tpch/snapshots/test_h20/test_tpc_h20/trino/h20.sql b/ibis/backends/tests/tpch/snapshots/test_h20/test_tpc_h20/trino/h20.sql
@@ -1,51 +1,76 @@
 SELECT
-  "t10"."s_name",
-  "t10"."s_address"
+  "t13"."s_name",
+  "t13"."s_address"
 FROM (
   SELECT
-    "t5"."s_suppkey",
-    "t5"."s_name",
-    "t5"."s_address",
-    "t5"."s_nationkey",
-    "t5"."s_phone",
-    "t5"."s_acctbal",
-    "t5"."s_comment",
-    "t6"."n_nationkey",
-    "t6"."n_name",
-    "t6"."n_regionkey",
-    "t6"."n_comment"
-  FROM "supplier" AS "t5"
-  INNER JOIN "nation" AS "t6"
-    ON "t5"."s_nationkey" = "t6"."n_nationkey"
-) AS "t10"
+    "t10"."s_suppkey",
+    "t10"."s_name",
+    "t10"."s_address",
+    "t10"."s_nationkey",
+    "t10"."s_phone",
+    "t10"."s_acctbal",
+    "t10"."s_comment",
+    "t8"."n_nationkey",
+    "t8"."n_name",
+    "t8"."n_regionkey",
+    "t8"."n_comment"
+  FROM (
+    SELECT
+      "t0"."s_suppkey",
+      "t0"."s_name",
+      "t0"."s_address",
+      "t0"."s_nationkey",
+      "t0"."s_phone",
+      CAST("t0"."s_acctbal" AS DECIMAL(15, 2)) AS "s_acctbal",
+      "t0"."s_comment"
+    FROM "supplier" AS "t0"
+  ) AS "t10"
+  INNER JOIN (
+    SELECT
+      "t2"."n_nationkey",
+      "t2"."n_name",
+      "t2"."n_regionkey",
+      "t2"."n_comment"
+    FROM "nation" AS "t2"
+  ) AS "t8"
+    ON "t10"."s_nationkey" = "t8"."n_nationkey"
+) AS "t13"
 WHERE
-  "t10"."n_name" = 'CANADA'
-  AND "t10"."s_suppkey" IN (
+  "t13"."n_name" = 'CANADA'
+  AND "t13"."s_suppkey" IN (
     SELECT
-      "t1"."ps_suppkey"
-    FROM "partsupp" AS "t1"
+      "t7"."ps_suppkey"
+    FROM (
+      SELECT
+        "t1"."ps_partkey",
+        "t1"."ps_suppkey",
+        "t1"."ps_availqty",
+        CAST("t1"."ps_supplycost" AS DECIMAL(15, 2)) AS "ps_supplycost",
+        "t1"."ps_comment"
+      FROM "partsupp" AS "t1"
+    ) AS "t7"
     WHERE
-      "t1"."ps_partkey" IN (
+      "t7"."ps_partkey" IN (
         SELECT
           "t3"."p_partkey"
         FROM "part" AS "t3"
         WHERE
           "t3"."p_name" LIKE 'forest%'
       )
-      AND "t1"."ps_availqty" > (
+      AND "t7"."ps_availqty" > (
         (
           SELECT
-            SUM("t8"."l_quantity") AS "Sum(l_quantity)"
+            SUM("t11"."l_quantity") AS "Sum(l_quantity)"
           FROM (
             SELECT
               "t4"."l_orderkey",
               "t4"."l_partkey",
               "t4"."l_suppkey",
               "t4"."l_linenumber",
-              "t4"."l_quantity",
-              "t4"."l_extendedprice",
-              "t4"."l_discount",
-              "t4"."l_tax",
+              CAST("t4"."l_quantity" AS DECIMAL(15, 2)) AS "l_quantity",
+              CAST("t4"."l_extendedprice" AS DECIMAL(15, 2)) AS "l_extendedprice",
+              CAST("t4"."l_discount" AS DECIMAL(15, 2)) AS "l_discount",
+              CAST("t4"."l_tax" AS DECIMAL(15, 2)) AS "l_tax",
               "t4"."l_returnflag",
               "t4"."l_linestatus",
               "t4"."l_shipdate",
@@ -56,13 +81,13 @@ WHERE
               "t4"."l_comment"
             FROM "lineitem" AS "t4"
             WHERE
-              "t4"."l_partkey" = "t1"."ps_partkey"
-              AND "t4"."l_suppkey" = "t1"."ps_suppkey"
+              "t4"."l_partkey" = "t7"."ps_partkey"
+              AND "t4"."l_suppkey" = "t7"."ps_suppkey"
               AND "t4"."l_shipdate" >= FROM_ISO8601_DATE('1994-01-01')
               AND "t4"."l_shipdate" < FROM_ISO8601_DATE('1995-01-01')
-          ) AS "t8"
+          ) AS "t11"
         ) * CAST(0.5 AS DOUBLE)
       )
   )
 ORDER BY
-  "t10"."s_name" ASC
+  "t13"."s_name" ASC