test: add partitioned table creation test and improve existing table assertions

Author: kosiew

python/tests/test_dataframe.py

@@ -1608,6 +1608,22 @@ def test_table_from_batches_dataframe(df, fail_collect):
         assert isinstance(batch, pa.RecordBatch)
 
 
+def test_table_from_batches_partitions(fail_collect):
+    ctx = SessionContext()
+
+    batch1 = pa.record_batch([pa.array([1])], names=["a"])
+    batch2 = pa.record_batch([pa.array([2])], names=["a"])
+    df = ctx.create_dataframe([[batch1], [batch2]])
+
+    table = pa.Table.from_batches(df)
+    expected = pa.Table.from_batches([batch1, batch2])
+
+    assert table.equals(expected)
+
+    for batch in df:
+        assert isinstance(batch, RecordBatch)
+
+
 def test_arrow_c_stream_to_table_and_reader(fail_collect):
     ctx = SessionContext()
 

python/tests/test_io.py

@@ -129,8 +129,9 @@ def test_table_from_batches_stream(ctx, fail_collect):
     df = range_table(ctx, 0, 10)
 
     table = pa.Table.from_batches(df)
-    assert table.shape == (10, 1)
-    assert table.column_names == ["value"]
+    expected = pa.Table.from_arrays([pa.array(range(10))], names=["value"])
+
+    assert table.equals(expected)
 
     for batch in df:
         assert isinstance(batch, pa.RecordBatch)