fix(sqlalchemy): handle correlated exists sanely

Also fix an issue where exists queries on projections didn't work at all

Author: cpcloud

ibis/backends/base/sql/alchemy/query_builder.py

@@ -42,7 +42,7 @@ def get_result(self):
         for jtype, table, preds in zip(
             self.join_types, self.join_tables[1:], self.join_predicates
         ):
-            if len(preds):
+            if preds:
                 sqla_preds = [self._translate(pred) for pred in preds]
                 onclause = functools.reduce(sql.and_, sqla_preds)
             else:
@@ -59,9 +59,23 @@ def get_result(self):
             elif jtype is ops.OuterJoin:
                 result = result.outerjoin(table, onclause, full=True)
             elif jtype is ops.LeftSemiJoin:
-                result = result.select().where(sa.exists(sa.select(1).where(onclause)))
+                # subquery is required for semi and anti joins done using
+                # sqlalchemy, otherwise multiple references to the original
+                # select are treated as distinct tables
+                #
+                # with a subquery, the result is a distinct table and so there's only one
+                # thing for subsequent expressions to reference
+                result = (
+                    result.select()
+                    .where(sa.exists(sa.select(1).where(onclause)))
+                    .subquery()
+                )
             elif jtype is ops.LeftAntiJoin:
-                result = result.select().where(~sa.exists(sa.select(1).where(onclause)))
+                result = (
+                    result.select()
+                    .where(~sa.exists(sa.select(1).where(onclause)))
+                    .subquery()
+                )
             else:
                 raise NotImplementedError(jtype)
 
@@ -227,32 +241,7 @@ def _add_select(self, table_set):
         if has_select_star or table_set is None:
             return result
 
-        # if we're selecting from something that isn't a subquery e.g., Select,
-        # Alias, Table
-        if not isinstance(table_set, sa.sql.Subquery):
-            return result.select_from(table_set)
-
-        final_froms = result.get_final_froms()
-        num_froms = len(final_froms)
-
-        # if the result subquery has no FROMs then we can select from the
-        # table_set since there's only a single possibility for FROM
-        if not num_froms:
-            return result.select_from(table_set)
-
-        # we need to replace every occurrence of `result`'s `FROM`
-        # with `table_set` to handle correlated EXISTs coming from
-        # semi/anti-join
-        #
-        # previously this was `replace_selectable`, but that's deprecated so we
-        # inline its implementation here
-        #
-        # sqlalchemy suggests using the functionality in sa.sql.visitors, but
-        # that would effectively require reimplementing ClauseAdapter
-        replaced = sa.sql.util.ClauseAdapter(table_set).traverse(result)
-        num_froms = len(replaced.get_final_froms())
-        assert num_froms == 1, f"num_froms == {num_froms:d}"
-        return replaced
+        return result.select_from(table_set)
 
     def _add_group_by(self, fragment):
         # GROUP BY and HAVING

ibis/backends/base/sql/alchemy/registry.py

@@ -98,7 +98,16 @@ def get_col_or_deferred_col(sa_table, colname):
     try:
         return sa_table.exported_columns[colname]
     except KeyError:
-        return sa.column(colname)
+        # cols is a sqlalchemy column collection which contains column
+        # names that are secretly prefixed by their containing table
+        #
+        # sqlalchemy doesn't let you select by the *un*prefixed column name
+        # despite the uniqueness of `colname`
+        #
+        # however, in ibis we have already deduplicated column names so we can
+        # refer to the name by position
+        colindex = op.table.schema._name_locs[colname]
+        return cols[colindex]
 
 
 def _table_column(t, op):
@@ -123,9 +132,18 @@ def _table_column(t, op):
 
 
 def _table_array_view(t, op):
+    # the table that the TableArrayView op contains (op.table) has
+    # one or more input relations that we need to "pin" for sqlalchemy's
+    # auto correlation functionality -- this is what `.correlate_except` does
+    #
+    # every relation that is NOT passed to `correlate_except` is considered an
+    # outer-query table
     ctx = t.context
     table = ctx.get_compiled_expr(op.table)
-    return table
+    # TODO: handle the case of `op.table` being a join
+    first, *_ = an.find_immediate_parent_tables(op.table, keep_input=False)
+    ref = ctx.get_ref(first)
+    return table.correlate_except(ref)
 
 
 def _exists_subquery(t, op):

ibis/backends/clickhouse/compiler/values.py

@@ -548,11 +548,20 @@ def _truncate(op, **kw):
 @translate_val.register(ops.ExistsSubquery)
 @translate_val.register(ops.NotExistsSubquery)
 def _exists_subquery(op, **kw):
-    foreign_table = translate_val(op.foreign_table, **kw)
+    # https://github.com/ClickHouse/ClickHouse/issues/6697
+    #
+    # this would work, if clickhouse supported correlated subqueries
+    from ibis.backends.clickhouse.compiler.relations import translate_rel
+
+    foreign_table = translate_rel(op.foreign_table, **kw)
     predicates = translate_val(op.predicates, **kw)
-    subq = sg.subquery(foreign_table.where(predicates, dialect="clickhouse").select(1))
+    subq = (
+        sg.select(1)
+        .from_(foreign_table, dialect="clickhouse")
+        .where(sg.condition(predicates), dialect="clickhouse")
+    )
     prefix = "NOT " * isinstance(op, ops.NotExistsSubquery)
-    return f"{prefix}EXISTS {subq}"
+    return f"{prefix}EXISTS ({subq})"
 
 
 @translate_val.register(ops.StringSplit)

ibis/backends/tests/test_generic.py

@@ -1,7 +1,7 @@
 import decimal
 import io
 from contextlib import redirect_stdout
-from operator import invert, neg
+from operator import invert, methodcaller, neg
 
 import numpy as np
 import pandas as pd
@@ -745,3 +745,27 @@ def test_int_scalar(alltypes):
     result = expr.execute()
     assert expr.type() == dt.int16
     assert result.dtype == np.int16
+
+
+@pytest.mark.notimpl(["dask", "datafusion", "pandas", "polars"])
+@pytest.mark.notyet(
+    ["clickhouse"], reason="https://github.com/ClickHouse/ClickHouse/issues/6697"
+)
+@pytest.mark.notyet(["pyspark"])
+@pytest.mark.parametrize(
+    "method_name",
+    [
+        "any",
+        param(
+            "notany",
+            marks=pytest.mark.broken(
+                ["impala"], reason="aliases are incorrectly elided"
+            ),
+        ),
+    ],
+)
+def test_exists(batting, awards_players, method_name):
+    method = methodcaller(method_name)
+    expr = batting[method(batting.yearID == awards_players.yearID)]
+    result = expr.execute()
+    assert not result.empty