fix: update DataFrame iteration to use num_rows and correct RecordBatch type

Author: kosiew

python/tests/test_dataframe.py

@@ -1302,7 +1302,7 @@ def test_execution_plan(aggregate_df):
         try:
             batch = stream.next()
             assert batch is not None
-            rows_returned += len(batch.to_pyarrow()[0])
+            rows_returned += batch.num_rows
         except StopIteration:
             # This is one of the partitions with no values
             pass
@@ -1317,7 +1317,7 @@ async def test_async_iteration_of_df(aggregate_df):
     rows_returned = 0
     async for batch in aggregate_df:
         assert batch is not None
-        rows_returned += len(batch.to_pyarrow()[0])
+        rows_returned += batch.num_rows
 
     assert rows_returned == 5
 
@@ -1728,6 +1728,7 @@ def test_arrow_c_stream_capsule_released(ctx):
         get_ptr(capsule, b"arrow_array_stream")
     pyerr_clear()
 
+
 def test_to_pylist(df):
     # Convert datafusion dataframe to Python list
     pylist = df.to_pylist()

python/tests/test_dataframe_iter.py

@@ -15,26 +15,26 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import pyarrow as pa
 import pytest
-import datafusion
 
 
 def test_iter_dataframe(ctx):
     df = ctx.from_pydict({"a": [1, 2]})
-    batches = [batch.to_pyarrow() for batch in df]
+    batches = list(df)
     assert len(batches) == 1
     assert batches[0].column(0).to_pylist() == [1, 2]
 
 
 def test_iter_returns_record_batch(ctx):
     df = ctx.from_pydict({"a": [1, 2]})
     batch = next(iter(df))
-    assert isinstance(batch, datafusion.RecordBatch)
+    assert isinstance(batch, pa.RecordBatch)
 
 
 @pytest.mark.asyncio
 async def test_async_iter_dataframe(ctx):
     df = ctx.from_pydict({"a": [1, 2]})
     batches = [batch async for batch in df]
     assert len(batches) == 1
-    assert batches[0].to_pyarrow().column(0).to_pylist() == [1, 2]
+    assert batches[0].column(0).to_pylist() == [1, 2]