feat(polars): reuse existing ctx for set operations

ibis-project · Dec 28, 2024 · 62fa479 · 62fa479
1 parent 2742f44
commit 62fa479
Showing 1 changed file with 18 additions and 16 deletions.
diff --git a/ibis/backends/polars/compiler.py b/ibis/backends/polars/compiler.py
@@ -1258,37 +1258,39 @@ def execute_union(op, **kw):
 
 @translate.register(ops.Intersection)
 def execute_intersection(op, **kw):
-    with pl.SQLContext(
+    ctx = kw.get("ctx")
+    ctx.register_many(
         frames={
             "left": translate(op.left, **kw),
             "right": translate(op.right, **kw),
         }
-    ) as ctx:
-        sql = (
-            sg.select(STAR)
-            .from_(sg.to_identifier("left", quoted=True))
-            .intersect(sg.select(STAR).from_(sg.to_identifier("right", quoted=True)))
-        )
-        result = ctx.execute(sql.sql())
+    )
+    sql = (
+        sg.select(STAR)
+        .from_(sg.to_identifier("left", quoted=True))
+        .intersect(sg.select(STAR).from_(sg.to_identifier("right", quoted=True)))
+    )
+    result = ctx.execute(sql.sql(Polars), eager=False)
     if op.distinct is True:
         return result.unique()
     return result
 
 
 @translate.register(ops.Difference)
 def execute_difference(op, **kw):
-    with pl.SQLContext(
+    ctx = kw.get("ctx")
+    ctx.register_many(
         frames={
             "left": translate(op.left, **kw),
             "right": translate(op.right, **kw),
         }
-    ) as ctx:
-        sql = (
-            sg.select(STAR)
-            .from_(sg.to_identifier("left", quoted=True))
-            .except_(sg.select(STAR).from_(sg.to_identifier("right", quoted=True)))
-        )
-        result = ctx.execute(sql.sql())
+    )
+    sql = (
+        sg.select(STAR)
+        .from_(sg.to_identifier("left", quoted=True))
+        .except_(sg.select(STAR).from_(sg.to_identifier("right", quoted=True)))
+    )
+    result = ctx.execute(sql.sql(Polars), eager=False)
     if op.distinct is True:
         return result.unique()
     return result