Repository: ibis-project/ibis

Issue:
bug: Combining union and order_by generates invalid BigQuery SQL
### What happened?

Hi Ibis team,

When applying union operation on table expression with order_by, it generates bad SQL.

A simple code piece can reproduce the issue:
```
import ibis

conn = ibis.bigquery.connect(
    project_id='garrettwu-test-project-2',
    dataset_id='bigquery-public-data.stackoverflow')

table = conn.table('posts_questions')

t = table.order_by("id")

unioned = ibis.union(t, t)

print(unioned.compile())

unioned.execute()
```
Generated SQL:
```
SELECT t0.*
FROM `bigquery-public-data.stackoverflow.posts_questions` t0
ORDER BY t0.`id` ASC
UNION ALL
SELECT t0.*
FROM `bigquery-public-data.stackoverflow.posts_questions` t0
ORDER BY t0.`id` ASC
```
Error:
```
BadRequest: 400 Syntax error: Expected end of input but got keyword UNION at [4:1]
```
(Full message in log output)

Same operation used to work for some previous commits.

### What version of ibis are you using?

master

Since the operation worked for versions sometime ago, we tried to run "git bisect" to locate the bad commit. It looks like https://github.com/ibis-project/ibis/pull/5571 is the one.

### What backend(s) are you using, if any?

BigQuery

### Relevant log output

```sh
# Error Message
---------------------------------------------------------------------------
BadRequest                                Traceback (most recent call last)
Cell In[11], line 1
----> 1 unioned.execute()

File ~/src/ibis/ibis/expr/types/core.py:303, in Expr.execute(self, limit, timecontext, params, **kwargs)
    276 def execute(
    277     self,
    278     limit: int | str | None = 'default',
   (...)
    281     **kwargs: Any,
    282 ):
    283     """Execute an expression against its backend if one exists.
    284 
    285     Parameters
   (...)
    301         Keyword arguments
    302     """
--> 303     return self._find_backend(use_default=True).execute(
    304         self, limit=limit, timecontext=timecontext, params=params, **kwargs
    305     )

File ~/src/ibis/ibis/backends/bigquery/__init__.py:298, in Backend.execute(self, expr, params, limit, **kwargs)
    296 sql = query_ast.compile()
    297 self._log(sql)
--> 298 cursor = self.raw_sql(sql, params=params, **kwargs)
    299 schema = self.ast_schema(query_ast, **kwargs)
    300 result = self.fetch_from_cursor(cursor, schema)

File ~/src/ibis/ibis/backends/bigquery/__init__.py:255, in Backend.raw_sql(self, query, results, params)
    242 def raw_sql(self, query: str, results=False, params=None):
    243     query_parameters = [
    244         bigquery_param(
    245             param.type(),
   (...)
    253         for param, value in (params or {}).items()
    254     ]
--> 255     return self._execute(query, results=results, query_parameters=query_parameters)

File ~/src/ibis/ibis/backends/bigquery/__init__.py:239, in Backend._execute(self, stmt, results, query_parameters)
    235 job_config.use_legacy_sql = False  # False by default in >=0.28
    236 query = self.client.query(
    237     stmt, job_config=job_config, project=self.billing_project
    238 )
--> 239 query.result()  # blocks until finished
    240 return BigQueryCursor(query)

File ~/src/bigframes/venv/lib/python3.10/site-packages/google/cloud/bigquery/job/query.py:1499, in QueryJob.result(self, page_size, max_results, retry, timeout, start_index, job_retry)
   1496     if retry_do_query is not None and job_retry is not None:
   1497         do_get_result = job_retry(do_get_result)
-> 1499     do_get_result()
   1501 except exceptions.GoogleAPICallError as exc:
   1502     exc.message = _EXCEPTION_FOOTER_TEMPLATE.format(
   1503         message=exc.message, location=self.location, job_id=self.job_id
   1504     )

File ~/src/bigframes/venv/lib/python3.10/site-packages/google/api_core/retry.py:349, in Retry.__call__.<locals>.retry_wrapped_func(*args, **kwargs)
    345 target = functools.partial(func, *args, **kwargs)
    346 sleep_generator = exponential_sleep_generator(
    347     self._initial, self._maximum, multiplier=self._multiplier
    348 )
--> 349 return retry_target(
    350     target,
    351     self._predicate,
    352     sleep_generator,
    353     self._timeout,
    354     on_error=on_error,
    355 )

File ~/src/bigframes/venv/lib/python3.10/site-packages/google/api_core/retry.py:191, in retry_target(target, predicate, sleep_generator, timeout, on_error, **kwargs)
    189 for sleep in sleep_generator:
    190     try:
--> 191         return target()
    193     # pylint: disable=broad-except
    194     # This function explicitly must deal with broad exceptions.
    195     except Exception as exc:

File ~/src/bigframes/venv/lib/python3.10/site-packages/google/cloud/bigquery/job/query.py:1489, in QueryJob.result.<locals>.do_get_result()
   1486     self._retry_do_query = retry_do_query
   1487     self._job_retry = job_retry
-> 1489 super(QueryJob, self).result(retry=retry, timeout=timeout)
   1491 # Since the job could already be "done" (e.g. got a finished job
   1492 # via client.get_job), the superclass call to done() might not
   1493 # set the self._query_results cache.
   1494 self._reload_query_results(retry=retry, timeout=timeout)

File ~/src/bigframes/venv/lib/python3.10/site-packages/google/cloud/bigquery/job/base.py:728, in _AsyncJob.result(self, retry, timeout)
    725     self._begin(retry=retry, timeout=timeout)
    727 kwargs = {} if retry is DEFAULT_RETRY else {"retry": retry}
--> 728 return super(_AsyncJob, self).result(timeout=timeout, **kwargs)

File ~/src/bigframes/venv/lib/python3.10/site-packages/google/api_core/future/polling.py:261, in PollingFuture.result(self, timeout, retry, polling)
    256 self._blocking_poll(timeout=timeout, retry=retry, polling=polling)
    258 if self._exception is not None:
    259     # pylint: disable=raising-bad-type
    260     # Pylint doesn't recognize that this is valid in this case.
--> 261     raise self._exception
    263 return self._result

BadRequest: 400 Syntax error: Expected end of input but got keyword UNION at [4:1]

Location: US
Job ID: 7d6ccc8d-f948-4d60-b681-7a23eb5179da
```


### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:

Resolving patch:
diff --git a/ibis/backends/base/sql/compiler/base.py b/ibis/backends/base/sql/compiler/base.py
--- a/ibis/backends/base/sql/compiler/base.py
+++ b/ibis/backends/base/sql/compiler/base.py
@@ -68,8 +68,12 @@ def _get_keyword_list(self):
         return map(self.keyword, self.distincts)
 
     def _extract_subqueries(self):
+        # extract any subquery to avoid generating incorrect sql when at least
+        # one of the set operands is invalid outside of being a subquery
+        #
+        # for example: SELECT * FROM t ORDER BY x UNION ...
         self.subqueries = an.find_subqueries(
-            [self.table_set, *self.filters], min_dependents=2
+            [self.table_set, *self.filters], min_dependents=1
         )
         for subquery in self.subqueries:
             self.context.set_extracted(subquery)


Test patch:
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql
--- a/ibis/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql
+++ b/ibis/tests/sql/snapshots/test_compiler/test_difference_project_column/out.sql
@@ -1,10 +1,18 @@
 SELECT t0.`key`
 FROM (
-  SELECT t1.`string_col` AS `key`, CAST(t1.`float_col` AS double) AS `value`
-  FROM functional_alltypes t1
-  WHERE t1.`int_col` > 0
+  WITH t1 AS (
+    SELECT t3.`string_col` AS `key`, t3.`double_col` AS `value`
+    FROM functional_alltypes t3
+    WHERE t3.`int_col` <= 0
+  ),
+  t2 AS (
+    SELECT t3.`string_col` AS `key`, CAST(t3.`float_col` AS double) AS `value`
+    FROM functional_alltypes t3
+    WHERE t3.`int_col` > 0
+  )
+  SELECT *
+  FROM t2
   EXCEPT
-  SELECT t1.`string_col` AS `key`, t1.`double_col` AS `value`
-  FROM functional_alltypes t1
-  WHERE t1.`int_col` <= 0
+  SELECT *
+  FROM t1
 ) t0
\ No newline at end of file
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql
--- a/ibis/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql
+++ b/ibis/tests/sql/snapshots/test_compiler/test_intersect_project_column/out.sql
@@ -1,10 +1,18 @@
 SELECT t0.`key`
 FROM (
-  SELECT t1.`string_col` AS `key`, CAST(t1.`float_col` AS double) AS `value`
-  FROM functional_alltypes t1
-  WHERE t1.`int_col` > 0
+  WITH t1 AS (
+    SELECT t3.`string_col` AS `key`, t3.`double_col` AS `value`
+    FROM functional_alltypes t3
+    WHERE t3.`int_col` <= 0
+  ),
+  t2 AS (
+    SELECT t3.`string_col` AS `key`, CAST(t3.`float_col` AS double) AS `value`
+    FROM functional_alltypes t3
+    WHERE t3.`int_col` > 0
+  )
+  SELECT *
+  FROM t2
   INTERSECT
-  SELECT t1.`string_col` AS `key`, t1.`double_col` AS `value`
-  FROM functional_alltypes t1
-  WHERE t1.`int_col` <= 0
+  SELECT *
+  FROM t1
 ) t0
\ No newline at end of file
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_table_difference/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_table_difference/out.sql
--- a/ibis/tests/sql/snapshots/test_compiler/test_table_difference/out.sql
+++ b/ibis/tests/sql/snapshots/test_compiler/test_table_difference/out.sql
@@ -1,7 +1,15 @@
-SELECT t0.`string_col` AS `key`, CAST(t0.`float_col` AS double) AS `value`
-FROM functional_alltypes t0
-WHERE t0.`int_col` > 0
+WITH t0 AS (
+  SELECT t2.`string_col` AS `key`, t2.`double_col` AS `value`
+  FROM functional_alltypes t2
+  WHERE t2.`int_col` <= 0
+),
+t1 AS (
+  SELECT t2.`string_col` AS `key`, CAST(t2.`float_col` AS double) AS `value`
+  FROM functional_alltypes t2
+  WHERE t2.`int_col` > 0
+)
+SELECT *
+FROM t1
 EXCEPT
-SELECT t0.`string_col` AS `key`, t0.`double_col` AS `value`
-FROM functional_alltypes t0
-WHERE t0.`int_col` <= 0
\ No newline at end of file
+SELECT *
+FROM t0
\ No newline at end of file
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql
--- a/ibis/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql
+++ b/ibis/tests/sql/snapshots/test_compiler/test_table_intersect/out.sql
@@ -1,7 +1,15 @@
-SELECT t0.`string_col` AS `key`, CAST(t0.`float_col` AS double) AS `value`
-FROM functional_alltypes t0
-WHERE t0.`int_col` > 0
+WITH t0 AS (
+  SELECT t2.`string_col` AS `key`, t2.`double_col` AS `value`
+  FROM functional_alltypes t2
+  WHERE t2.`int_col` <= 0
+),
+t1 AS (
+  SELECT t2.`string_col` AS `key`, CAST(t2.`float_col` AS double) AS `value`
+  FROM functional_alltypes t2
+  WHERE t2.`int_col` > 0
+)
+SELECT *
+FROM t1
 INTERSECT
-SELECT t0.`string_col` AS `key`, t0.`double_col` AS `value`
-FROM functional_alltypes t0
-WHERE t0.`int_col` <= 0
\ No newline at end of file
+SELECT *
+FROM t0
\ No newline at end of file
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_union/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_union/out.sql
--- a/ibis/tests/sql/snapshots/test_compiler/test_union/out.sql
+++ b/ibis/tests/sql/snapshots/test_compiler/test_union/out.sql
@@ -1,7 +1,15 @@
-SELECT t0.`string_col` AS `key`, CAST(t0.`float_col` AS double) AS `value`
-FROM functional_alltypes t0
-WHERE t0.`int_col` > 0
+WITH t0 AS (
+  SELECT t2.`string_col` AS `key`, t2.`double_col` AS `value`
+  FROM functional_alltypes t2
+  WHERE t2.`int_col` <= 0
+),
+t1 AS (
+  SELECT t2.`string_col` AS `key`, CAST(t2.`float_col` AS double) AS `value`
+  FROM functional_alltypes t2
+  WHERE t2.`int_col` > 0
+)
+SELECT *
+FROM t1
 UNION
-SELECT t0.`string_col` AS `key`, t0.`double_col` AS `value`
-FROM functional_alltypes t0
-WHERE t0.`int_col` <= 0
\ No newline at end of file
+SELECT *
+FROM t0
\ No newline at end of file
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_union_order_by/decompiled.py b/ibis/tests/sql/snapshots/test_compiler/test_union_order_by/decompiled.py
new file mode 100644
--- /dev/null
+++ b/ibis/tests/sql/snapshots/test_compiler/test_union_order_by/decompiled.py
@@ -0,0 +1,7 @@
+import ibis
+
+
+t = ibis.table(name="t", schema={"a": "int64", "b": "string"})
+proj = t.order_by(t.b.asc())
+
+result = proj.union(proj)
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_union_order_by/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_union_order_by/out.sql
new file mode 100644
--- /dev/null
+++ b/ibis/tests/sql/snapshots/test_compiler/test_union_order_by/out.sql
@@ -0,0 +1,10 @@
+WITH t0 AS (
+  SELECT t1.*
+  FROM t t1
+  ORDER BY t1.`b` ASC
+)
+SELECT *
+FROM t0
+UNION ALL
+SELECT *
+FROM t0
\ No newline at end of file
diff --git a/ibis/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql b/ibis/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql
--- a/ibis/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql
+++ b/ibis/tests/sql/snapshots/test_compiler/test_union_project_column/out.sql
@@ -1,10 +1,18 @@
 SELECT t0.`key`
 FROM (
-  SELECT t1.`string_col` AS `key`, CAST(t1.`float_col` AS double) AS `value`
-  FROM functional_alltypes t1
-  WHERE t1.`int_col` > 0
+  WITH t1 AS (
+    SELECT t3.`string_col` AS `key`, t3.`double_col` AS `value`
+    FROM functional_alltypes t3
+    WHERE t3.`int_col` <= 0
+  ),
+  t2 AS (
+    SELECT t3.`string_col` AS `key`, CAST(t3.`float_col` AS double) AS `value`
+    FROM functional_alltypes t3
+    WHERE t3.`int_col` > 0
+  )
+  SELECT *
+  FROM t2
   UNION ALL
-  SELECT t1.`string_col` AS `key`, t1.`double_col` AS `value`
-  FROM functional_alltypes t1
-  WHERE t1.`int_col` <= 0
+  SELECT *
+  FROM t1
 ) t0
\ No newline at end of file
diff --git a/ibis/tests/sql/test_compiler.py b/ibis/tests/sql/test_compiler.py
--- a/ibis/tests/sql/test_compiler.py
+++ b/ibis/tests/sql/test_compiler.py
@@ -218,3 +218,10 @@ def test_column_expr_default_name(snapshot):
     expr = t.int_col + 4
     snapshot.assert_match(to_sql(expr), "out.sql")
     assert_decompile_roundtrip(expr, snapshot)
+
+
+def test_union_order_by(snapshot):
+    t = ibis.table(dict(a="int", b="string"), name="t")
+    expr = t.order_by("b").union(t.order_by("b"))
+    snapshot.assert_match(to_sql(expr), "out.sql")
+    assert_decompile_roundtrip(expr, snapshot)