changed batches conversion to Table instead

This commit is contained in:
Bryan Cutler committed 2016-12-06 14:24:01 -08:00
1 parent 4b72329fe2
commit 72ea875255
3 files changed
+50 -50

No files matched your search

+1 -2
View File
@@ -49,5 +49,4 @@ from pyarrow.schema import (null, bool_,
list_, struct, field,
DataType, Field, Schema, schema)
from pyarrow.table import (Column, RecordBatch, dataframe_from_batches, Table,
from_pandas_dataframe)
from pyarrow.table import Column, RecordBatch, Table, from_pandas_dataframe
+46 -46
View File
@@ -415,52 +415,6 @@ cdef class RecordBatch:
return result
def dataframe_from_batches(batches):
"""
Convert a list of Arrow RecordBatches to a pandas.DataFrame
Parameters
----------
batches: list of RecordBatch
RecordBatch list to be converted, schemas must be equal
"""
cdef:
vector[shared_ptr[CArray]] c_array_chunks
vector[shared_ptr[CColumn]] c_columns
shared_ptr[CTable] c_table
Array arr
Schema schema
import pandas as pd
schema = batches[0].schema
# check schemas are equal
if any((not schema.equals(other.schema) for other in batches[1:])):
raise ArrowException("Error converting list of RecordBatches to "
"DataFrame, not all schemas are equal")
cdef int K = batches[0].num_columns
# create chunked columns from the batches
c_columns.resize(K)
for i in range(K):
for batch in batches:
arr = batch[i]
c_array_chunks.push_back(arr.sp_array)
c_columns[i].reset(new CColumn(schema.sp_schema.get().field(i),
c_array_chunks))
c_array_chunks.clear()
# create a Table from columns and convert to DataFrame
c_table.reset(new CTable('', schema.sp_schema, c_columns))
table = Table()
table.init(c_table)
return table.to_pandas()
cdef class Table:
"""
A collection of top-level named, equal length Arrow arrays.
@@ -567,6 +521,52 @@ cdef class Table:
return result
@staticmethod
def from_batches(batches):
"""
Construct a Table from a list of Arrow RecordBatches
Parameters
----------
batches: list of RecordBatch
RecordBatch list to be converted, schemas must be equal
"""
cdef:
vector[shared_ptr[CArray]] c_array_chunks
vector[shared_ptr[CColumn]] c_columns
shared_ptr[CTable] c_table
Array arr
Schema schema
import pandas as pd
schema = batches[0].schema
# check schemas are equal
if any((not schema.equals(other.schema) for other in batches[1:])):
raise ArrowException("Error converting list of RecordBatches to "
"DataFrame, not all schemas are equal")
cdef int K = batches[0].num_columns
# create chunked columns from the batches
c_columns.resize(K)
for i in range(K):
for batch in batches:
arr = batch[i]
c_array_chunks.push_back(arr.sp_array)
c_columns[i].reset(new CColumn(schema.sp_schema.get().field(i),
c_array_chunks))
c_array_chunks.clear()
# create a Table from columns and convert to DataFrame
c_table.reset(new CTable('', schema.sp_schema, c_columns))
table = Table()
table.init(c_table)
return table
def to_pandas(self):
"""
Convert the arrow::Table to a pandas DataFrame
+3 -2
View File
@@ -69,7 +69,8 @@ def test_recordbatchlist_to_pandas():
batch1 = pa.RecordBatch.from_pandas(data1)
batch2 = pa.RecordBatch.from_pandas(data2)
result = pa.dataframe_from_batches([batch1, batch2])
table = pa.Table.from_batches([batch1, batch2])
result = table.to_pandas()
data = pd.concat([data1, data2], ignore_index=True)
assert_frame_equal(data, result)
@@ -82,7 +83,7 @@ def test_recordbatchlist_schema_equals():
batch2 = pa.RecordBatch.from_pandas(data2)
with pytest.raises(pa.ArrowException):
pa.dataframe_from_batches([batch1, batch2])
pa.Table.from_batches([batch1, batch2])
def test_table_basics():