Skip to content

Commit 1ebd3c1

Browse files
committed
test: improve test_iter_batches_dataframe to validate RecordBatch conversion
1 parent 610aed3 commit 1ebd3c1

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

python/tests/test_dataframe.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1629,9 +1629,11 @@ def test_iter_batches_dataframe(fail_collect):
16291629
df = ctx.create_dataframe([[batch1], [batch2]])
16301630

16311631
expected = [batch1, batch2]
1632-
for got, exp in zip(df, expected):
1633-
assert isinstance(got, RecordBatch)
1634-
assert got.to_pyarrow().equals(exp)
1632+
results = [b.to_pyarrow() for b in df]
1633+
1634+
assert len(results) == len(expected)
1635+
for exp in expected:
1636+
assert any(got.equals(exp) for got in results)
16351637

16361638

16371639
def test_arrow_c_stream_to_table_and_reader(fail_collect):

0 commit comments

Comments
 (0)