fix(bigquery): apply unnest transformation in other methods that exec…

…ute SQL
ibis-project · Oct 19, 2023 · 2cc9d0e · 2cc9d0e
1 parent 940ed21
commit 2cc9d0e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/ibis/backends/bigquery/__init__.py b/ibis/backends/bigquery/__init__.py
@@ -473,8 +473,8 @@ def to_pyarrow(
         **kwargs: Any,
     ) -> pa.Table:
         self._import_pyarrow()
-        query_ast = self.compiler.to_ast_ensure_limit(expr, limit, params=params)
-        sql = query_ast.compile()
+        sql = self.compile(expr, limit=limit, params=params, **kwargs)
+        self._log(sql)
         cursor = self.raw_sql(sql, params=params, **kwargs)
         table = self._cursor_to_arrow(cursor)
         return expr.__pyarrow_result__(table)
@@ -492,8 +492,8 @@ def to_pyarrow_batches(
 
         schema = expr.as_table().schema()
 
-        query_ast = self.compiler.to_ast_ensure_limit(expr, limit, params=params)
-        sql = query_ast.compile()
+        sql = self.compile(expr, limit=limit, params=params, **kwargs)
+        self._log(sql)
         cursor = self.raw_sql(sql, params=params, **kwargs)
         batch_iter = self._cursor_to_arrow(
             cursor,

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -373,6 +373,9 @@ def test_unnest_complex(backend):
     result = expr.execute()
     tm.assert_frame_equal(result, expected)
 
+    # test that unnest works with to_pyarrow
+    assert len(expr.to_pyarrow()) == len(result)
+
 
 @builtin_array
 @pytest.mark.never(