feat(sql): add option to enable/disable select merging (ibis-project#9065)

Add a global option to enable or disable merging select statements where
possible. This can serve as an escape route if the optimization fails.
It is currently enabled by default, we may change to make it opt-in.

xref ibis-project#9064 and ibis-project#9058

Closes ibis-project#9058.

Author: kszucs

ibis/backends/sql/compiler.py

@@ -27,6 +27,7 @@
     rewrite_capitalize,
     sqlize,
 )
+from ibis.config import options
 from ibis.expr.operations.udf import InputType
 from ibis.expr.rewrites import replace_bucket
 
@@ -453,7 +454,12 @@ def translate(self, op, *, params: Mapping[ir.Value, Any]) -> sge.Expression:
         # substitute parameters immediately to avoid having to define a
         # ScalarParameter translation rule
         params = self._prepare_params(params)
-        op, ctes = sqlize(op, params=params, rewrites=self.rewrites)
+        op, ctes = sqlize(
+            op,
+            params=params,
+            rewrites=self.rewrites,
+            fuse_selects=options.sql.fuse_selects,
+        )
 
         aliases = {}
         counter = itertools.count()

ibis/backends/sql/rewrites.py

@@ -208,6 +208,7 @@ def sqlize(
     node: ops.Node,
     params: Mapping[ops.ScalarParameter, Any],
     rewrites: Sequence[Pattern] = (),
+    fuse_selects: bool = True,
 ) -> tuple[ops.Node, list[ops.Node]]:
     """Lower the ibis expression graph to a SQL-like relational algebra.
 
@@ -219,6 +220,8 @@ def sqlize(
         A mapping of scalar parameters to their values.
     rewrites
         Supplementary rewrites to apply to the expression graph.
+    fuse_selects
+        Whether to merge subsequent Select nodes into one where possible.
 
     Returns
     -------
@@ -243,7 +246,10 @@ def sqlize(
     )
 
     # squash subsequent Select nodes into one
-    simplified = sqlized.replace(merge_select_select)
+    if fuse_selects:
+        simplified = sqlized.replace(merge_select_select)
+    else:
+        simplified = sqlized
 
     # extract common table expressions while wrapping them in a CTE node
     ctes = frozenset(extract_ctes(simplified))

ibis/config.py

@@ -47,6 +47,9 @@ class SQL(Config):
 
     Attributes
     ----------
+    fuse_selects : bool
+        Whether to fuse consecutive select queries into a single query where
+        possible.
     default_limit : int | None
         Number of rows to be retrieved for a table expression without an
         explicit limit. [](`None`) means no limit.
@@ -55,6 +58,7 @@ class SQL(Config):
 
     """
 
+    fuse_selects: bool = True
     default_limit: Optional[PosInt] = None
     default_dialect: str = "duckdb"