def write_table(table: pyarrow.Table, compression_level=0) -> bytes:
table = table.combine_chunks()
doc = collections.OrderedDict()
for col in table.columns:
buf = write_array(col.data.chunks[0], compression_level)
doc[col.name] = bson.raw_bson.RawBSONDocument(buf)
return bson.encode(doc)
def make_table_one_chunk(table: pa.Table) -> pa.Table:
assert len(table.columns), "Workbench must not give a zero-column table"
if table.columns[0].num_chunks == 0:
return pa.table(
{
field.name: pa.array([], field.type)
for field in (table.schema.field(i)
for i in range(len(table.schema.names)))
},
schema=table.schema,
)
if table.columns[0].num_chunks == 1:
return table
return table.combine_chunks()
