mirror of
https://github.com/apache/arrow.git
synced 2026-10-29 00:56:31 +08:00
changed batches conversion to Table instead
This commit is contained in:
1 parent
4b72329fe2
commit
72ea875255
3 files changed
+50
-50
No files matched your search
@@ -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
@@ -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
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in new issue
Block a user