fix(ir): make ibis.to_sql() parse join limits and projections; fixes #11105

Author: YStrauchP4

ibis/expr/sql.py

@@ -71,6 +71,25 @@ def overlay(self, step):
         return Catalog({**self, **updates})
 
 
+def apply_limit(table, step):
+    """Applies a LIMIT, if applicable."""
+
+    if not isinstance(step.limit, int):
+        return table
+
+    return table.limit(step.limit)
+
+
+def apply_projections(table, step, catalog):
+    """Applies a SELECT projection, if applicable."""
+
+    if not step.projections:
+        return table
+
+    projs = [convert(proj, catalog=catalog) for proj in step.projections]
+    return table.select(projs)
+
+
 @singledispatch
 def convert(step, catalog):
     raise TypeError(type(step))
@@ -86,12 +105,8 @@ def convert_scan(scan, catalog):
         pred = convert(scan.condition, catalog=catalog)
         table = table.filter(pred)
 
-    if scan.projections:
-        projs = [convert(proj, catalog=catalog) for proj in scan.projections]
-        table = table.select(projs)
-
-    if isinstance(scan.limit, int):
-        table = table.limit(scan.limit)
+    table = apply_projections(table, scan, catalog)
+    table = apply_limit(table, scan)
 
     return table
 
@@ -156,8 +171,7 @@ def convert_sort(sort, catalog):
         ]
         table = table.select(projs)
 
-    if isinstance(sort.limit, int):
-        table = table.limit(sort.limit)
+    table = apply_limit(table, sort)
 
     return table
 
@@ -203,6 +217,9 @@ def convert_join(join, catalog):
         predicate = convert(join.condition, catalog=catalog)
         left_table = left_table.filter(predicate)
 
+    left_table = apply_projections(left_table, join, catalog)
+    left_table = apply_limit(left_table, join)
+
     catalog[left_name] = left_table
 
     return left_table

ibis/expr/tests/snapshots/test_sql/test_parse_sql_join_subquery/decompiled.py

@@ -0,0 +1,25 @@
+import ibis
+
+
+call = ibis.table(
+    name="call",
+    schema={
+        "start_time": "timestamp",
+        "end_time": "timestamp",
+        "employee_id": "int64",
+        "call_outcome_id": "int64",
+        "call_attempts": "int64",
+    },
+)
+agg = call.aggregate([call.call_attempts.mean().name("mean")])
+
+result = call.inner_join(
+    agg, [(call.call_attempts > agg.mean), ibis.literal(True)]
+).select(
+    call.start_time,
+    call.end_time,
+    call.employee_id,
+    call.call_outcome_id,
+    call.call_attempts,
+    agg.mean,
+)

ibis/expr/tests/snapshots/test_sql/test_parse_sql_limited_join/decompiled.py

@@ -0,0 +1,33 @@
+import ibis
+
+
+employee = ibis.table(
+    name="employee",
+    schema={"first_name": "string", "last_name": "string", "id": "int64"},
+)
+call = ibis.table(
+    name="call",
+    schema={
+        "start_time": "timestamp",
+        "end_time": "timestamp",
+        "employee_id": "int64",
+        "call_outcome_id": "int64",
+        "call_attempts": "int64",
+    },
+)
+
+result = (
+    employee.inner_join(call, [(employee.id == call.employee_id), ibis.literal(True)])
+    .select(
+        employee.first_name,
+        employee.last_name,
+        employee.id,
+        call.start_time,
+        call.end_time,
+        call.employee_id,
+        call.call_outcome_id,
+        call.call_attempts,
+        employee.first_name.name("first"),
+    )
+    .limit(3)
+)

ibis/expr/tests/snapshots/test_sql/test_parse_sql_scalar_subquery/decompiled.py

@@ -21,5 +21,4 @@
     call.employee_id,
     call.call_outcome_id,
     call.call_attempts,
-    agg.mean,
 )

ibis/expr/tests/test_sql.py

@@ -44,9 +44,23 @@ def test_parse_sql_basic_join(how, snapshot):
     snapshot.assert_match(code, "decompiled.py")
 
 
+def test_parse_sql_limited_join(snapshot):
+    sql = """
+SELECT
+  *,
+  first_name as first
+FROM employee
+JOIN call ON
+  employee.id = call.employee_id
+LIMIT 3"""
+    expr = ibis.parse_sql(sql, catalog)
+    code = ibis.decompile(expr, format=True)
+    snapshot.assert_match(code, "decompiled.py")
+
+
 def test_parse_sql_multiple_joins(snapshot):
     sql = """
-SELECT *
+SELECT employee.*, call.*, call_outcome.outcome_text, call_outcome.id as id_right
 FROM employee
 JOIN call
   ON employee.id = call.employee_id
@@ -118,6 +132,21 @@ def test_parse_sql_scalar_subquery(snapshot):
     snapshot.assert_match(code, "decompiled.py")
 
 
+def test_parse_sql_join_subquery(snapshot):
+    sql = """
+SELECT *
+FROM call
+INNER JOIN (
+  SELECT
+    AVG(call.call_attempts) AS mean
+  FROM call
+) AS subq
+  ON subq.mean < call.call_attempts"""
+    expr = ibis.parse_sql(sql, catalog)
+    code = ibis.decompile(expr, format=True)
+    snapshot.assert_match(code, "decompiled.py")
+
+
 def test_parse_sql_simple_select_count(snapshot):
     sql = """SELECT COUNT(first_name) FROM employee"""
     expr = ibis.parse_sql(sql, catalog)