fix: datafusion: ensure that to_pyarrow_batches does do compute

ibis-project · Feb 27, 2024 · d1a62d0 · d1a62d0
1 parent 99dda5b
commit d1a62d0
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/ibis/backends/datafusion/__init__.py b/ibis/backends/datafusion/__init__.py
@@ -509,9 +509,8 @@ def to_pyarrow_batches(
 
         struct_schema = schema.as_struct().to_pyarrow()
 
-        return pa.ipc.RecordBatchReader.from_batches(
-            schema.to_pyarrow(),
-            (
+        def make_gen():
+            yield from (
                 # convert the renamed + casted columns into a record batch
                 pa.RecordBatch.from_struct_array(
                     # rename columns to match schema because datafusion lowercases things
@@ -522,7 +521,11 @@ def to_pyarrow_batches(
                     .cast(struct_schema)
                 )
                 for batch in frame.collect()
-            ),
+            )
+
+        return pa.ipc.RecordBatchReader.from_batches(
+            schema.to_pyarrow(),
+            make_gen(),
         )
 
     def to_pyarrow(self, expr: ir.Expr, **kwargs: Any) -> pa.Table: