xlnt/python/xlntpyarrow/__init__.py

87 lines
3.0 KiB
Python

import pyarrow as pa
import xlntpyarrow.lib as xpa
COLUMN_TYPE_FIELD = {
xpa.Cell.Type.Number: pa.float64,
xpa.Cell.Type.SharedString: pa.string,
xpa.Cell.Type.InlineString: pa.string,
xpa.Cell.Type.FormulaString: pa.string,
xpa.Cell.Type.Error: pa.string,
xpa.Cell.Type.Boolean: pa.bool_,
xpa.Cell.Type.Date: pa.date32,
xpa.Cell.Type.Empty: pa.string,
}
def cell_to_pyarrow_array(cell, type):
if cell.data_type() == xpa.Cell.Type.Number:
return pa.array([cell.value_double()], type)
elif cell.data_type() == xpa.Cell.Type.SharedString:
return pa.array([cell.value_string()], type)
elif cell.data_type() == xpa.Cell.Type.InlineString:
return pa.array([cell.value_string()], type)
elif cell.data_type() == xpa.Cell.Type.FormulaString:
return pa.array([cell.value_string()], type)
elif cell.data_type() == xpa.Cell.Type.Error:
return pa.array([cell.value_string()], type)
elif cell.data_type() == xpa.Cell.Type.Boolean:
return pa.array([cell.value_bool()], type)
elif cell.data_type() == xpa.Cell.Type.Date:
return pa.array([cell.value_unsigned_int()], type)
elif cell.data_type() == xpa.Cell.Type.Empty:
return pa.array([cell.value_string()], type)
def xlsx2arrow(io, sheetname):
reader = xpa.StreamingWorkbookReader()
reader.open(io)
sheet_titles = reader.sheet_titles()
sheet_title = sheet_titles[0]
if sheetname is not None:
if isinstance(sheetname, int):
sheet_title = sheet_titles[sheetname]
elif isinstance(sheetname, str):
sheet_title = sheetname
reader.begin_worksheet(sheet_title)
column_names = []
fields = []
batches = []
schema = None
first_batch = []
max_column = 0
while reader.has_cell():
if schema is None:
cell = reader.read_cell()
type = cell.data_type()
if cell.row() == 1:
column_names.append(cell.value_string())
max_column = max(max_column, cell.column())
continue
elif cell.row() == 2:
column_name = column_names[cell.column() - 1]
if type == xpa.Cell.Type.Number and cell.format_is_date():
fields.append(pa.field(column_name, pa.date32))
else:
fields.append(pa.field(column_name, COLUMN_TYPE_FIELD[type]()))
first_batch.append(cell_to_pyarrow_array(cell, fields[-1].type))
if cell.column() == max_column:
schema = pa.schema(fields)
print(schema)
batches.append(pa.RecordBatch.from_arrays(first_batch, column_names))
continue
batches.append(reader.read_batch(schema, 10000))
reader.end_worksheet()
return pa.Table.from_batches(batches)
if __name__ == '__main__':
file = open('tmp.xlsx', 'rb')
table = xlsx2arrow(file, 'Sheet1')
print(table.to_pandas())