Repository: ibis-project/ibis

Issue:
bug(api): cannot chain `.union()` calls
Using the `geography.db` used in the tutorial:

```python
In [1]: import ibis

In [2]: con = ibis.sqlite.connect("geography.db")

In [3]: a = con.tables.countries.filter(lambda _: _.continent == "EU").limit(1)

In [4]: b = con.tables.countries.filter(lambda _: _.continent == "AS").limit(1)

In [5]: c = con.tables.countries.filter(lambda _: _.continent == "AF").limit(1)

In [6]: o = a.union(b).union(c)

In [7]: o.execute()
...
OperationalError: (sqlite3.OperationalError) near "(": syntax error
[SQL: WITH anon_1 AS 
(SELECT t0.iso_alpha2 AS iso_alpha2, t0.iso_alpha3 AS iso_alpha3, t0.iso_numeric AS iso_numeric, t0.fips AS fips, t0.name AS name, t0.capital AS capital, t0.area_km2 AS area_km2, t0.population AS population, t0.continent AS continent 
FROM main.countries AS t0 
WHERE t0.continent = ?
 LIMIT ? OFFSET ?), 
anon_2 AS 
(SELECT t0.iso_alpha2 AS iso_alpha2, t0.iso_alpha3 AS iso_alpha3, t0.iso_numeric AS iso_numeric, t0.fips AS fips, t0.name AS name, t0.capital AS capital, t0.area_km2 AS area_km2, t0.population AS population, t0.continent AS continent 
FROM main.countries AS t0 
WHERE t0.continent = ?
 LIMIT ? OFFSET ?), 
anon_3 AS 
(SELECT t0.iso_alpha2 AS iso_alpha2, t0.iso_alpha3 AS iso_alpha3, t0.iso_numeric AS iso_numeric, t0.fips AS fips, t0.name AS name, t0.capital AS capital, t0.area_km2 AS area_km2, t0.population AS population, t0.continent AS continent 
FROM main.countries AS t0 
WHERE t0.continent = ?
 LIMIT ? OFFSET ?)
 (SELECT anon_1.iso_alpha2, anon_1.iso_alpha3, anon_1.iso_numeric, anon_1.fips, anon_1.name, anon_1.capital, anon_1.area_km2, anon_1.population, anon_1.continent 
FROM anon_1 UNION ALL SELECT anon_2.iso_alpha2, anon_2.iso_alpha3, anon_2.iso_numeric, anon_2.fips, anon_2.name, anon_2.capital, anon_2.area_km2, anon_2.population, anon_2.continent 
FROM anon_2) UNION ALL SELECT anon_3.iso_alpha2, anon_3.iso_alpha3, anon_3.iso_numeric, anon_3.fips, anon_3.name, anon_3.capital, anon_3.area_km2, anon_3.population, anon_3.continent 
FROM anon_3]
[parameters: ('EU', 1, 0, 'AS', 1, 0, 'AF', 1, 0)]
```

Running `a.union(b)` alone works fine, the error seems to be caused by the chaining alone.


Discussion:

Resolving patch:
diff --git a/ibis/backends/base/sql/alchemy/query_builder.py b/ibis/backends/base/sql/alchemy/query_builder.py
--- a/ibis/backends/base/sql/alchemy/query_builder.py
+++ b/ibis/backends/base/sql/alchemy/query_builder.py
@@ -344,39 +344,49 @@ def _convert_group_by(self, exprs):
 
 
 class AlchemySetOp(SetOp):
-    @classmethod
-    def reduce(cls, left, right, distincts):
-        distinct = next(distincts)
-        sa_func = cls.distinct_func if distinct else cls.non_distinct_func
-        return sa_func(left, right)
-
     def compile(self):
         context = self.context
         selects = []
 
+        def call(distinct, *args):
+            return (
+                self.distinct_func(*args)
+                if distinct
+                else self.non_distinct_func(*args)
+            )
+
         for table in self.tables:
             table_set = context.get_compiled_expr(table)
             selects.append(table_set.cte().select())
 
-        return functools.reduce(
-            functools.partial(self.reduce, distincts=iter(self.distincts)),
-            selects,
-        )
+        if len(set(self.distincts)) == 1:
+            # distinct is either all True or all False, handle with a single
+            # call. This generates much more concise SQL.
+            return call(self.distincts[0], *selects)
+        else:
+            # We need to iteratively apply the set operations to handle
+            # disparate `distinct` values. Subqueries _must_ be converted using
+            # `.subquery().select()` to get sqlalchemy to put parenthesis in
+            # the proper places.
+            result = selects[0]
+            for select, distinct in zip(selects[1:], self.distincts):
+                result = call(distinct, result.subquery().select(), select)
+            return result
 
 
 class AlchemyUnion(AlchemySetOp):
-    distinct_func = sa.union
-    non_distinct_func = sa.union_all
+    distinct_func = staticmethod(sa.union)
+    non_distinct_func = staticmethod(sa.union_all)
 
 
 class AlchemyIntersection(AlchemySetOp):
-    distinct_func = sa.intersect
-    non_distinct_func = sa.intersect_all
+    distinct_func = staticmethod(sa.intersect)
+    non_distinct_func = staticmethod(sa.intersect_all)
 
 
 class AlchemyDifference(AlchemySetOp):
-    distinct_func = sa.except_
-    non_distinct_func = sa.except_all
+    distinct_func = staticmethod(sa.except_)
+    non_distinct_func = staticmethod(sa.except_all)
 
 
 class AlchemyCompiler(Compiler):
diff --git a/ibis/backends/pandas/execution/generic.py b/ibis/backends/pandas/execution/generic.py
--- a/ibis/backends/pandas/execution/generic.py
+++ b/ibis/backends/pandas/execution/generic.py
@@ -953,7 +953,7 @@ def execute_difference_dataframe_dataframe(
     merged = left.merge(
         right, on=list(left.columns), how="outer", indicator=True
     )
-    result = merged[merged["_merge"] != "both"].drop("_merge", axis=1)
+    result = merged[merged["_merge"] == "left_only"].drop("_merge", axis=1)
     return result
 
 


Test patch:
diff --git a/ibis/backends/postgres/tests/test_functions.py b/ibis/backends/postgres/tests/test_functions.py
--- a/ibis/backends/postgres/tests/test_functions.py
+++ b/ibis/backends/postgres/tests/test_functions.py
@@ -544,45 +544,37 @@ def test_category_label(alltypes, df):
 
 
 @pytest.mark.parametrize(
-    ('distinct1', 'distinct2', 'expected1', 'expected2'),
-    [
-        (True, True, 'UNION', 'UNION'),
-        (True, False, 'UNION', 'UNION ALL'),
-        (False, True, 'UNION ALL', 'UNION'),
-        (False, False, 'UNION ALL', 'UNION ALL'),
-    ],
+    ('distinct', 'union'),
+    [(True, 'UNION'), (False, 'UNION ALL')],
 )
-def test_union_cte(alltypes, distinct1, distinct2, expected1, expected2):
+def test_union_cte(alltypes, distinct, union):
     t = alltypes
     expr1 = t.group_by(t.string_col).aggregate(metric=t.double_col.sum())
     expr2 = expr1.view()
     expr3 = expr1.view()
-    expr = expr1.union(expr2, distinct=distinct1).union(
-        expr3, distinct=distinct2
+    expr = expr1.union(expr2, distinct=distinct).union(
+        expr3, distinct=distinct
     )
-    result = '\n'.join(
-        map(
-            lambda line: line.rstrip(),  # strip trailing whitespace
-            str(
-                expr.compile().compile(compile_kwargs={'literal_binds': True})
-            ).splitlines(),
-        )
+    result = ' '.join(
+        line.strip()
+        for line in str(
+            expr.compile().compile(compile_kwargs={'literal_binds': True})
+        ).splitlines()
     )
-    expected = """\
-WITH anon_1 AS
-(SELECT t0.string_col AS string_col, sum(t0.double_col) AS metric
-FROM functional_alltypes AS t0 GROUP BY t0.string_col),
-anon_2 AS
-(SELECT t0.string_col AS string_col, sum(t0.double_col) AS metric
-FROM functional_alltypes AS t0 GROUP BY t0.string_col),
-anon_3 AS
-(SELECT t0.string_col AS string_col, sum(t0.double_col) AS metric
-FROM functional_alltypes AS t0 GROUP BY t0.string_col)
- (SELECT anon_1.string_col, anon_1.metric
-FROM anon_1 {} SELECT anon_2.string_col, anon_2.metric
-FROM anon_2) {} SELECT anon_3.string_col, anon_3.metric
-FROM anon_3""".format(
-        expected1, expected2
+    expected = (
+        "WITH anon_1 AS "
+        "(SELECT t0.string_col AS string_col, sum(t0.double_col) AS metric "
+        "FROM functional_alltypes AS t0 GROUP BY t0.string_col), "
+        "anon_2 AS "
+        "(SELECT t0.string_col AS string_col, sum(t0.double_col) AS metric "
+        "FROM functional_alltypes AS t0 GROUP BY t0.string_col), "
+        "anon_3 AS "
+        "(SELECT t0.string_col AS string_col, sum(t0.double_col) AS metric "
+        "FROM functional_alltypes AS t0 GROUP BY t0.string_col) "
+        "SELECT anon_1.string_col, anon_1.metric "
+        f"FROM anon_1 {union} SELECT anon_2.string_col, anon_2.metric "
+        f"FROM anon_2 {union} SELECT anon_3.string_col, anon_3.metric "
+        "FROM anon_3"
     )
     assert str(result) == expected
 
diff --git a/ibis/backends/tests/test_set_ops.py b/ibis/backends/tests/test_set_ops.py
--- a/ibis/backends/tests/test_set_ops.py
+++ b/ibis/backends/tests/test_set_ops.py
@@ -2,7 +2,20 @@
 import pytest
 from pytest import param
 
-import ibis
+from ibis import _
+
+
+@pytest.fixture
+def union_subsets(alltypes, df):
+    a = alltypes.filter((5200 <= _.id) & (_.id <= 5210))
+    b = alltypes.filter((5205 <= _.id) & (_.id <= 5215))
+    c = alltypes.filter((5213 <= _.id) & (_.id <= 5220))
+
+    da = df[(5200 <= df.id) & (df.id <= 5210)]
+    db = df[(5205 <= df.id) & (df.id <= 5215)]
+    dc = df[(5213 <= df.id) & (df.id <= 5220)]
+
+    return (a, b, c), (da, db, dc)
 
 
 @pytest.mark.parametrize(
@@ -10,40 +23,35 @@
     [param(False, id="all"), param(True, id="distinct")],
 )
 @pytest.mark.notimpl(["datafusion"])
-def test_union(backend, alltypes, df, distinct):
-    default_limit = ibis.options.sql.default_limit
-    ibis.options.sql.default_limit = 15000
-    try:
-        expr = alltypes.union(alltypes, distinct=distinct).sort_by("id")
-        result = expr.execute().reset_index(drop=True)
-        expected = df if distinct else pd.concat([df, df], axis=0)
-
-        backend.assert_frame_equal(
-            result,
-            expected.sort_values("id").reset_index(drop=True),
-        )
-    finally:
-        ibis.options.sql.default_limit = default_limit
+def test_union(backend, union_subsets, distinct):
+    (a, b, c), (da, db, dc) = union_subsets
+
+    expr = (
+        a.union(b, distinct=distinct).union(c, distinct=distinct).sort_by("id")
+    )
+    result = expr.execute()
+
+    expected = (
+        pd.concat([da, db, dc], axis=0)
+        .sort_values("id")
+        .reset_index(drop=True)
+    )
+    if distinct:
+        expected = expected.drop_duplicates("id")
+
+    backend.assert_frame_equal(result, expected)
 
 
-@pytest.mark.parametrize(
-    "distinct",
-    [
-        param(False, id="all"),
-        param(
-            True,
-            marks=pytest.mark.notimpl(
-                ["duckdb", "postgres"],
-                reason="Result order not guaranteed when distinct=True",
-            ),
-            id="distinct",
-        ),
-    ],
-)
 @pytest.mark.notimpl(["datafusion"])
-def test_union_no_sort(backend, alltypes, df, distinct):
-    result = alltypes.union(alltypes, distinct=distinct).execute()
-    expected = df if distinct else pd.concat([df, df], axis=0)
+def test_union_mixed_distinct(backend, union_subsets):
+    (a, b, c), (da, db, dc) = union_subsets
+
+    expr = a.union(b, distinct=True).union(c, distinct=False).sort_by("id")
+    result = expr.execute()
+    expected = pd.concat(
+        [pd.concat([da, db], axis=0).drop_duplicates("id"), dc], axis=0
+    ).sort_values("id")
+
     backend.assert_frame_equal(result, expected)
 
 
@@ -64,13 +72,29 @@ def test_union_no_sort(backend, alltypes, df, distinct):
 @pytest.mark.notimpl(["datafusion"])
 @pytest.mark.notyet(["impala"])
 def test_intersect(backend, alltypes, df, distinct):
-    expr = alltypes.intersect(alltypes, distinct=distinct).sort_by("id")
-    result = expr.execute()
-    expected = df
-    backend.assert_frame_equal(
-        result,
-        expected.sort_values("id").reset_index(drop=True),
+    a = alltypes.filter((5200 <= _.id) & (_.id <= 5210))
+    b = alltypes.filter((5205 <= _.id) & (_.id <= 5215))
+    c = alltypes.filter((5195 <= _.id) & (_.id <= 5208))
+
+    # Reset index to ensure simple RangeIndex, needed for computing `expected`
+    df = df.reset_index(drop=True)
+    da = df[(5200 <= df.id) & (df.id <= 5210)]
+    db = df[(5205 <= df.id) & (df.id <= 5215)]
+    dc = df[(5195 <= df.id) & (df.id <= 5208)]
+
+    expr = (
+        a.intersect(b, distinct=distinct)
+        .intersect(c, distinct=distinct)
+        .sort_by("id")
     )
+    result = expr.execute()
+
+    index = da.index.intersection(db.index).intersection(dc.index)
+    expected = df.iloc[index].sort_values("id").reset_index(drop=True)
+    if distinct:
+        expected = expected.drop_duplicates()
+
+    backend.assert_frame_equal(result, expected)
 
 
 @pytest.mark.parametrize(
@@ -89,12 +113,27 @@ def test_intersect(backend, alltypes, df, distinct):
 )
 @pytest.mark.notimpl(["datafusion"])
 @pytest.mark.notyet(["impala"])
-def test_difference(backend, alltypes, distinct):
-    expr = alltypes.difference(alltypes, distinct=distinct)
+def test_difference(backend, alltypes, df, distinct):
+    a = alltypes.filter((5200 <= _.id) & (_.id <= 5210))
+    b = alltypes.filter((5205 <= _.id) & (_.id <= 5215))
+    c = alltypes.filter((5195 <= _.id) & (_.id <= 5202))
+
+    # Reset index to ensure simple RangeIndex, needed for computing `expected`
+    df = df.reset_index(drop=True)
+    da = df[(5200 <= df.id) & (df.id <= 5210)]
+    db = df[(5205 <= df.id) & (df.id <= 5215)]
+    dc = df[(5195 <= df.id) & (df.id <= 5202)]
+
+    expr = (
+        a.difference(b, distinct=distinct)
+        .difference(c, distinct=distinct)
+        .sort_by("id")
+    )
     result = expr.execute()
 
-    dtypes = {
-        column: dtype.to_pandas() for column, dtype in expr.schema().items()
-    }
-    expected = pd.DataFrame(columns=alltypes.columns).astype(dtypes)
+    index = da.index.difference(db.index).difference(dc.index)
+    expected = df.iloc[index].sort_values("id").reset_index(drop=True)
+    if distinct:
+        expected = expected.drop_duplicates()
+
     backend.assert_frame_equal(result, expected)