Skip to content

Commit b5de9e5

Browse files
BryanCutlerwesm
authored andcommitted
ARROW-369: [Python] Convert multiple record batches at once to Pandas
Modified Pandas adapter to handle columns with multiple chunks with `ConvertColumnToPandas`. This modifies the pyarrow public API by adding a class `RecordBatchList` and static method `toPandas` which takes a list of Arrow RecordBatches and outputs a Pandas DataFrame. Adds unit test in test_table.py to do the conversion for each column with typed specialization. Author: Bryan Cutler <cutlerb@gmail.com> Closes #216 from BryanCutler/multi-batch-toPandas-ARROW-369 and squashes the following commits: b6c9986 [Bryan Cutler] fixed formatting edf056e [Bryan Cutler] simplified with pyarrow.schema.Schema.equals 068bc1b [Bryan Cutler] Merge remote-tracking branch 'upstream/master' into multi-batch-toPandas-ARROW-369 da65345 [Bryan Cutler] fixed test case for schema checking 9edb0ba [Bryan Cutler] used auto keyword where some typecasting was done in ConvertValues bd2a720 [Bryan Cutler] added testcase for schema not equal, disabled now c3d7e8f [Bryan Cutler] Changed conversion to make Table from columns first, now conversion is now just a free function 3ee51e6 [Bryan Cutler] cleanup 398b18d [Bryan Cutler] Fixed case for Integer specialization without nulls 7b29a55 [Bryan Cutler] Initial working version of RecordBatch list to_pandas, need more tests and cleanup
1 parent ebe7dc8 commit b5de9e5

5 files changed

Lines changed: 219 additions & 79 deletions

File tree

python/pyarrow/__init__.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,5 +41,7 @@
4141
list_, struct, field,
4242
DataType, Field, Schema, schema)
4343

44-
from pyarrow.table import Column, RecordBatch, Table, from_pandas_dataframe
44+
from pyarrow.table import (Column, RecordBatch, dataframe_from_batches, Table,
45+
from_pandas_dataframe)
46+
4547
from pyarrow.version import version as __version__

python/pyarrow/includes/libarrow.pxd

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
158158
CColumn(const shared_ptr[CField]& field,
159159
const shared_ptr[CArray]& data)
160160

161+
CColumn(const shared_ptr[CField]& field,
162+
const vector[shared_ptr[CArray]]& chunks)
163+
161164
int64_t length()
162165
int64_t null_count()
163166
const c_string& name()

python/pyarrow/table.pyx

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ cimport pyarrow.includes.pyarrow as pyarrow
2828
import pyarrow.config
2929

3030
from pyarrow.array cimport Array, box_arrow_array
31+
from pyarrow.error import ArrowException
3132
from pyarrow.error cimport check_status
3233
from pyarrow.schema cimport box_data_type, box_schema
3334

@@ -414,6 +415,52 @@ cdef class RecordBatch:
414415
return result
415416

416417

418+
def dataframe_from_batches(batches):
419+
"""
420+
Convert a list of Arrow RecordBatches to a pandas.DataFrame
421+
422+
Parameters
423+
----------
424+
425+
batches: list of RecordBatch
426+
RecordBatch list to be converted, schemas must be equal
427+
"""
428+
429+
cdef:
430+
vector[shared_ptr[CArray]] c_array_chunks
431+
vector[shared_ptr[CColumn]] c_columns
432+
shared_ptr[CTable] c_table
433+
Array arr
434+
Schema schema
435+
436+
import pandas as pd
437+
438+
schema = batches[0].schema
439+
440+
# check schemas are equal
441+
if any((not schema.equals(other.schema) for other in batches[1:])):
442+
raise ArrowException("Error converting list of RecordBatches to "
443+
"DataFrame, not all schemas are equal")
444+
445+
cdef int K = batches[0].num_columns
446+
447+
# create chunked columns from the batches
448+
c_columns.resize(K)
449+
for i in range(K):
450+
for batch in batches:
451+
arr = batch[i]
452+
c_array_chunks.push_back(arr.sp_array)
453+
c_columns[i].reset(new CColumn(schema.sp_schema.get().field(i),
454+
c_array_chunks))
455+
c_array_chunks.clear()
456+
457+
# create a Table from columns and convert to DataFrame
458+
c_table.reset(new CTable('', schema.sp_schema, c_columns))
459+
table = Table()
460+
table.init(c_table)
461+
return table.to_pandas()
462+
463+
417464
cdef class Table:
418465
"""
419466
A collection of top-level named, equal length Arrow arrays.

python/pyarrow/tests/test_table.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
from pandas.util.testing import assert_frame_equal
2121
import pandas as pd
22+
import pytest
2223

2324
import pyarrow as pa
2425

@@ -50,6 +51,40 @@ def test_recordbatch_from_to_pandas():
5051
assert_frame_equal(data, result)
5152

5253

54+
def test_recordbatchlist_to_pandas():
55+
data1 = pd.DataFrame({
56+
'c1': np.array([1, 1, 2], dtype='uint32'),
57+
'c2': np.array([1.0, 2.0, 3.0], dtype='float64'),
58+
'c3': [True, None, False],
59+
'c4': ['foo', 'bar', None]
60+
})
61+
62+
data2 = pd.DataFrame({
63+
'c1': np.array([3, 5], dtype='uint32'),
64+
'c2': np.array([4.0, 5.0], dtype='float64'),
65+
'c3': [True, True],
66+
'c4': ['baz', 'qux']
67+
})
68+
69+
batch1 = pa.RecordBatch.from_pandas(data1)
70+
batch2 = pa.RecordBatch.from_pandas(data2)
71+
72+
result = pa.dataframe_from_batches([batch1, batch2])
73+
data = pd.concat([data1, data2], ignore_index=True)
74+
assert_frame_equal(data, result)
75+
76+
77+
def test_recordbatchlist_schema_equals():
78+
data1 = pd.DataFrame({'c1': np.array([1], dtype='uint32')})
79+
data2 = pd.DataFrame({'c1': np.array([4.0, 5.0], dtype='float64')})
80+
81+
batch1 = pa.RecordBatch.from_pandas(data1)
82+
batch2 = pa.RecordBatch.from_pandas(data2)
83+
84+
with pytest.raises(pa.ArrowException):
85+
pa.dataframe_from_batches([batch1, batch2])
86+
87+
5388
def test_table_basics():
5489
data = [
5590
pa.from_pylist(range(5)),

0 commit comments

Comments
 (0)