feat(ir): more flexible dereferencing support for join right hand side

ibis-project · Apr 18, 2024 · 6d6eed3 · 6d6eed3
1 parent 98ef69c
commit 6d6eed3
Show file tree

Hide file tree

Showing 27 changed files with 414 additions and 317 deletions.
diff --git a/ibis/backends/clickhouse/tests/snapshots/test_select/test_join_self_reference/out.sql b/ibis/backends/clickhouse/tests/snapshots/test_select/test_join_self_reference/out.sql
@@ -13,5 +13,5 @@ SELECT
   "t1"."year",
   "t1"."month"
 FROM "functional_alltypes" AS "t1"
-INNER JOIN "functional_alltypes" AS "t3"
-  ON "t1"."id" = "t3"."id"
+INNER JOIN "functional_alltypes" AS "t2"
+  ON "t1"."id" = "t2"."id"
diff --git a/ibis/backends/sql/compiler.py b/ibis/backends/sql/compiler.py
@@ -1155,6 +1155,8 @@ def visit_DatabaseTable(
     def visit_SelfReference(self, op, *, parent, identifier):
         return parent
 
+    visit_JoinReference = visit_SelfReference
+
     def visit_JoinChain(self, op, *, first, rest, values):
         result = sg.select(*self._cleanup_names(values), copy=False).from_(
             first, copy=False
@@ -1385,9 +1387,6 @@ def visit_SQLStringView(self, op, *, query: str, child, schema):
     def visit_SQLQueryResult(self, op, *, query, schema, source):
         return sg.parse_one(query, dialect=self.dialect).subquery(copy=False)
 
-    def visit_JoinTable(self, op, *, parent, index):
-        return parent
-
     def visit_RegexExtract(self, op, *, arg, pattern, index):
         return self.f.regexp_extract(arg, pattern, index, dialect=self.dialect)
 

diff --git a/ibis/backends/tests/sql/snapshots/test_select_sql/test_limit_cte_extract/out.sql b/ibis/backends/tests/sql/snapshots/test_select_sql/test_limit_cte_extract/out.sql
@@ -19,5 +19,5 @@ SELECT
   "t3"."year",
   "t3"."month"
 FROM "t1" AS "t3"
-INNER JOIN "t1" AS "t5"
+INNER JOIN "t1" AS "t4"
   ON TRUE
diff --git a/ibis/backends/tests/sql/snapshots/test_select_sql/test_limit_with_self_join/out.sql b/ibis/backends/tests/sql/snapshots/test_select_sql/test_limit_with_self_join/out.sql
@@ -15,20 +15,20 @@ FROM (
     "t1"."timestamp_col",
     "t1"."year",
     "t1"."month",
-    "t3"."id" AS "id_right",
-    "t3"."bool_col" AS "bool_col_right",
-    "t3"."tinyint_col" AS "tinyint_col_right",
-    "t3"."smallint_col" AS "smallint_col_right",
-    "t3"."int_col" AS "int_col_right",
-    "t3"."bigint_col" AS "bigint_col_right",
-    "t3"."float_col" AS "float_col_right",
-    "t3"."double_col" AS "double_col_right",
-    "t3"."date_string_col" AS "date_string_col_right",
-    "t3"."string_col" AS "string_col_right",
-    "t3"."timestamp_col" AS "timestamp_col_right",
-    "t3"."year" AS "year_right",
-    "t3"."month" AS "month_right"
+    "t2"."id" AS "id_right",
+    "t2"."bool_col" AS "bool_col_right",
+    "t2"."tinyint_col" AS "tinyint_col_right",
+    "t2"."smallint_col" AS "smallint_col_right",
+    "t2"."int_col" AS "int_col_right",
+    "t2"."bigint_col" AS "bigint_col_right",
+    "t2"."float_col" AS "float_col_right",
+    "t2"."double_col" AS "double_col_right",
+    "t2"."date_string_col" AS "date_string_col_right",
+    "t2"."string_col" AS "string_col_right",
+    "t2"."timestamp_col" AS "timestamp_col_right",
+    "t2"."year" AS "year_right",
+    "t2"."month" AS "month_right"
   FROM "functional_alltypes" AS "t1"
-  INNER JOIN "functional_alltypes" AS "t3"
-    ON "t1"."tinyint_col" < EXTRACT(minute FROM "t3"."timestamp_col")
-) AS "t4"
+  INNER JOIN "functional_alltypes" AS "t2"
+    ON "t1"."tinyint_col" < EXTRACT(minute FROM "t2"."timestamp_col")
+) AS "t3"
diff --git a/...ckends/tests/sql/snapshots/test_select_sql/test_self_join_subquery_distinct_equal/out.sql b/...ckends/tests/sql/snapshots/test_select_sql/test_self_join_subquery_distinct_equal/out.sql
@@ -1,6 +1,6 @@
 SELECT
   "t2"."r_name",
-  "t6"."n_name"
+  "t5"."n_name"
 FROM "tpch_region" AS "t2"
 INNER JOIN "tpch_nation" AS "t3"
   ON "t2"."r_regionkey" = "t3"."n_regionkey"
@@ -16,5 +16,5 @@ INNER JOIN (
   FROM "tpch_region" AS "t2"
   INNER JOIN "tpch_nation" AS "t3"
     ON "t2"."r_regionkey" = "t3"."n_regionkey"
-) AS "t6"
-  ON "t2"."r_regionkey" = "t6"."r_regionkey"
+) AS "t5"
+  ON "t2"."r_regionkey" = "t5"."r_regionkey"
diff --git a/ibis/backends/tests/sql/snapshots/test_select_sql/test_subquery_in_union/out.sql b/ibis/backends/tests/sql/snapshots/test_select_sql/test_subquery_in_union/out.sql
@@ -7,25 +7,25 @@ WITH "t1" AS (
   GROUP BY
     1,
     2
-), "t6" AS (
+), "t5" AS (
   SELECT
     "t3"."a",
     "t3"."g",
     "t3"."metric"
   FROM "t1" AS "t3"
-  INNER JOIN "t1" AS "t5"
-    ON "t3"."g" = "t5"."g"
+  INNER JOIN "t1" AS "t4"
+    ON "t3"."g" = "t4"."g"
 )
 SELECT
-  "t9"."a",
-  "t9"."g",
-  "t9"."metric"
+  "t8"."a",
+  "t8"."g",
+  "t8"."metric"
 FROM (
   SELECT
     *
-  FROM "t6" AS "t7"
+  FROM "t5" AS "t6"
   UNION ALL
   SELECT
     *
-  FROM "t6" AS "t8"
-) AS "t9"
+  FROM "t5" AS "t7"
+) AS "t8"
diff --git a/ibis/backends/tests/sql/snapshots/test_select_sql/test_subquery_used_for_self_join/out.sql b/ibis/backends/tests/sql/snapshots/test_select_sql/test_subquery_used_for_self_join/out.sql
@@ -11,21 +11,21 @@ WITH "t1" AS (
     3
 )
 SELECT
-  "t6"."g",
-  MAX("t6"."total" - "t6"."total_right") AS "metric"
+  "t5"."g",
+  MAX("t5"."total" - "t5"."total_right") AS "metric"
 FROM (
   SELECT
     "t3"."g",
     "t3"."a",
     "t3"."b",
     "t3"."total",
-    "t5"."g" AS "g_right",
-    "t5"."a" AS "a_right",
-    "t5"."b" AS "b_right",
-    "t5"."total" AS "total_right"
+    "t4"."g" AS "g_right",
+    "t4"."a" AS "a_right",
+    "t4"."b" AS "b_right",
+    "t4"."total" AS "total_right"
   FROM "t1" AS "t3"
-  INNER JOIN "t1" AS "t5"
-    ON "t3"."a" = "t5"."b"
-) AS "t6"
+  INNER JOIN "t1" AS "t4"
+    ON "t3"."a" = "t4"."b"
+) AS "t5"
 GROUP BY
   1
diff --git a/ibis/backends/tests/sql/snapshots/test_select_sql/test_tpch_self_join_failure/out.sql b/ibis/backends/tests/sql/snapshots/test_select_sql/test_tpch_self_join_failure/out.sql
@@ -24,9 +24,9 @@ WITH "t9" AS (
 SELECT
   "t11"."region",
   "t11"."year",
-  "t11"."total" - "t13"."total" AS "yoy_change"
+  "t11"."total" - "t12"."total" AS "yoy_change"
 FROM "t9" AS "t11"
-INNER JOIN "t9" AS "t13"
+INNER JOIN "t9" AS "t12"
   ON "t11"."year" = (
-    "t13"."year" - CAST(1 AS TINYINT)
+    "t12"."year" - CAST(1 AS TINYINT)
   )
diff --git a/ibis/backends/tests/sql/snapshots/test_sql/test_cte_factor_distinct_but_equal/out.sql b/ibis/backends/tests/sql/snapshots/test_sql/test_cte_factor_distinct_but_equal/out.sql
@@ -16,5 +16,5 @@ INNER JOIN (
   FROM "alltypes" AS "t1"
   GROUP BY
     1
-) AS "t6"
-  ON "t3"."g" = "t6"."g"
+) AS "t5"
+  ON "t3"."g" = "t5"."g"
diff --git a/ibis/backends/tests/sql/snapshots/test_sql/test_self_reference_join/out.sql b/ibis/backends/tests/sql/snapshots/test_sql/test_self_reference_join/out.sql
@@ -4,5 +4,5 @@ SELECT
   "t1"."foo_id",
   "t1"."bar_id"
 FROM "star1" AS "t1"
-INNER JOIN "star1" AS "t3"
-  ON "t1"."foo_id" = "t3"."bar_id"
+INNER JOIN "star1" AS "t2"
+  ON "t1"."foo_id" = "t2"."bar_id"
diff --git a/ibis/backends/tests/tpch/snapshots/test_h07/test_tpc_h07/duckdb/h07.sql b/ibis/backends/tests/tpch/snapshots/test_h07/test_tpc_h07/duckdb/h07.sql
@@ -1,27 +1,27 @@
 SELECT
-  "t14"."supp_nation",
-  "t14"."cust_nation",
-  "t14"."l_year",
-  "t14"."revenue"
+  "t13"."supp_nation",
+  "t13"."cust_nation",
+  "t13"."l_year",
+  "t13"."revenue"
 FROM (
   SELECT
-    "t13"."supp_nation",
-    "t13"."cust_nation",
-    "t13"."l_year",
-    SUM("t13"."volume") AS "revenue"
+    "t12"."supp_nation",
+    "t12"."cust_nation",
+    "t12"."l_year",
+    SUM("t12"."volume") AS "revenue"
   FROM (
     SELECT
-      "t12"."supp_nation",
-      "t12"."cust_nation",
-      "t12"."l_shipdate",
-      "t12"."l_extendedprice",
-      "t12"."l_discount",
-      "t12"."l_year",
-      "t12"."volume"
+      "t11"."supp_nation",
+      "t11"."cust_nation",
+      "t11"."l_shipdate",
+      "t11"."l_extendedprice",
+      "t11"."l_discount",
+      "t11"."l_year",
+      "t11"."volume"
     FROM (
       SELECT
         "t9"."n_name" AS "supp_nation",
-        "t11"."n_name" AS "cust_nation",
+        "t10"."n_name" AS "cust_nation",
         "t6"."l_shipdate",
         "t6"."l_extendedprice",
         "t6"."l_discount",
@@ -38,34 +38,34 @@ FROM (
         ON "t8"."c_custkey" = "t7"."o_custkey"
       INNER JOIN "nation" AS "t9"
         ON "t5"."s_nationkey" = "t9"."n_nationkey"
-      INNER JOIN "nation" AS "t11"
-        ON "t8"."c_nationkey" = "t11"."n_nationkey"
-    ) AS "t12"
+      INNER JOIN "nation" AS "t10"
+        ON "t8"."c_nationkey" = "t10"."n_nationkey"
+    ) AS "t11"
     WHERE
       (
         (
           (
-            "t12"."cust_nation" = 'FRANCE'
+            "t11"."cust_nation" = 'FRANCE'
           ) AND (
-            "t12"."supp_nation" = 'GERMANY'
+            "t11"."supp_nation" = 'GERMANY'
           )
         )
         OR (
           (
-            "t12"."cust_nation" = 'GERMANY'
+            "t11"."cust_nation" = 'GERMANY'
           ) AND (
-            "t12"."supp_nation" = 'FRANCE'
+            "t11"."supp_nation" = 'FRANCE'
           )
         )
       )
-      AND "t12"."l_shipdate" BETWEEN MAKE_DATE(1995, 1, 1) AND MAKE_DATE(1996, 12, 31)
-  ) AS "t13"
+      AND "t11"."l_shipdate" BETWEEN MAKE_DATE(1995, 1, 1) AND MAKE_DATE(1996, 12, 31)
+  ) AS "t12"
   GROUP BY
     1,
     2,
     3
-) AS "t14"
+) AS "t13"
 ORDER BY
-  "t14"."supp_nation" ASC,
-  "t14"."cust_nation" ASC,
-  "t14"."l_year" ASC
+  "t13"."supp_nation" ASC,
+  "t13"."cust_nation" ASC,
+  "t13"."l_year" ASC
diff --git a/ibis/backends/tests/tpch/snapshots/test_h08/test_tpc_h08/duckdb/h08.sql b/ibis/backends/tests/tpch/snapshots/test_h08/test_tpc_h08/duckdb/h08.sql
@@ -1,26 +1,26 @@
 SELECT
-  "t18"."o_year",
-  "t18"."mkt_share"
+  "t17"."o_year",
+  "t17"."mkt_share"
 FROM (
   SELECT
-    "t17"."o_year",
-    SUM("t17"."nation_volume") / SUM("t17"."volume") AS "mkt_share"
+    "t16"."o_year",
+    SUM("t16"."nation_volume") / SUM("t16"."volume") AS "mkt_share"
   FROM (
     SELECT
-      "t16"."o_year",
-      "t16"."volume",
-      "t16"."nation",
-      "t16"."r_name",
-      "t16"."o_orderdate",
-      "t16"."p_type",
-      CASE WHEN "t16"."nation" = 'BRAZIL' THEN "t16"."volume" ELSE CAST(0 AS TINYINT) END AS "nation_volume"
+      "t15"."o_year",
+      "t15"."volume",
+      "t15"."nation",
+      "t15"."r_name",
+      "t15"."o_orderdate",
+      "t15"."p_type",
+      CASE WHEN "t15"."nation" = 'BRAZIL' THEN "t15"."volume" ELSE CAST(0 AS TINYINT) END AS "nation_volume"
     FROM (
       SELECT
         EXTRACT(year FROM "t10"."o_orderdate") AS "o_year",
         "t8"."l_extendedprice" * (
           CAST(1 AS TINYINT) - "t8"."l_discount"
         ) AS "volume",
-        "t15"."n_name" AS "nation",
+        "t13"."n_name" AS "nation",
         "t14"."r_name",
         "t10"."o_orderdate",
         "t7"."p_type"
@@ -37,16 +37,16 @@ FROM (
         ON "t11"."c_nationkey" = "t12"."n_nationkey"
       INNER JOIN "region" AS "t14"
         ON "t12"."n_regionkey" = "t14"."r_regionkey"
-      INNER JOIN "nation" AS "t15"
-        ON "t9"."s_nationkey" = "t15"."n_nationkey"
-    ) AS "t16"
+      INNER JOIN "nation" AS "t13"
+        ON "t9"."s_nationkey" = "t13"."n_nationkey"
+    ) AS "t15"
     WHERE
-      "t16"."r_name" = 'AMERICA'
-      AND "t16"."o_orderdate" BETWEEN MAKE_DATE(1995, 1, 1) AND MAKE_DATE(1996, 12, 31)
-      AND "t16"."p_type" = 'ECONOMY ANODIZED STEEL'
-  ) AS "t17"
+      "t15"."r_name" = 'AMERICA'
+      AND "t15"."o_orderdate" BETWEEN MAKE_DATE(1995, 1, 1) AND MAKE_DATE(1996, 12, 31)
+      AND "t15"."p_type" = 'ECONOMY ANODIZED STEEL'
+  ) AS "t16"
   GROUP BY
     1
-) AS "t18"
+) AS "t17"
 ORDER BY
-  "t18"."o_year" ASC
+  "t17"."o_year" ASC
diff --git a/ibis/common/graph.py b/ibis/common/graph.py
@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
+import itertools
 from abc import abstractmethod
 from collections import deque
 from collections.abc import Iterable, Iterator, KeysView, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 
 from ibis.common.bases import Hashable
-from ibis.common.collections import frozendict
 from ibis.common.patterns import NoMatch, Pattern
 from ibis.common.typing import _ClassInfo
 from ibis.util import experimental, promote_list
@@ -66,8 +66,9 @@ def _flatten_collections(node: Any) -> Iterator[N]:
             yield item
         elif isinstance(item, (tuple, list)):
             yield from _flatten_collections(item)
-        elif isinstance(item, (dict, frozendict)):
-            yield from _flatten_collections(item.values())
+        elif isinstance(item, dict):
+            items = itertools.chain.from_iterable(item.items())
+            yield from _flatten_collections(items)
 
 
 def _recursive_lookup(obj: Any, dct: dict) -> Any:
@@ -117,8 +118,10 @@ def _recursive_lookup(obj: Any, dct: dict) -> Any:
         return dct.get(obj, obj)
     elif isinstance(obj, (tuple, list)):
         return tuple(_recursive_lookup(o, dct) for o in obj)
-    elif isinstance(obj, (dict, frozendict)):
-        return {k: _recursive_lookup(v, dct) for k, v in obj.items()}
+    elif isinstance(obj, dict):
+        return {
+            _recursive_lookup(k, dct): _recursive_lookup(v, dct) for k, v in obj.items()
+        }
     else:
         return obj