Spaces:
Sleeping
Sleeping
| """A Parquet file writer that wraps the pyarrow writer.""" | |
| from typing import IO, Optional | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| from .schema import Item, Schema, schema_to_arrow_schema | |
| class ParquetWriter: | |
| """A writer to parquet.""" | |
| def __init__(self, | |
| schema: Schema, | |
| codec: str = 'snappy', | |
| row_group_buffer_size: int = 128 * 1024 * 1024, | |
| record_batch_size: int = 10_000): | |
| self._schema = schema_to_arrow_schema(schema) | |
| self._codec = codec | |
| self._row_group_buffer_size = row_group_buffer_size | |
| self._buffer: list[list[Optional[Item]]] = [[] for _ in range(len(self._schema.names))] | |
| self._buffer_size = record_batch_size | |
| self._record_batches: list[pa.RecordBatch] = [] | |
| self._record_batches_byte_size = 0 | |
| self.writer: pq.ParquetWriter = None | |
| def open(self, file_handle: IO) -> None: | |
| """Open the destination file for writing.""" | |
| self.writer = pq.ParquetWriter(file_handle, self._schema, compression=self._codec) | |
| def write(self, record: Item) -> None: | |
| """Write the record to the destination file.""" | |
| if len(self._buffer[0]) >= self._buffer_size: | |
| self._flush_buffer() | |
| if self._record_batches_byte_size >= self._row_group_buffer_size: | |
| self._write_batches() | |
| # reorder the data in columnar format. | |
| for i, n in enumerate(self._schema.names): | |
| self._buffer[i].append(record.get(n)) | |
| def close(self) -> None: | |
| """Flushes the write buffer and closes the destination file.""" | |
| if len(self._buffer[0]) > 0: | |
| self._flush_buffer() | |
| if self._record_batches_byte_size > 0: | |
| self._write_batches() | |
| self.writer.close() | |
| def _write_batches(self) -> None: | |
| table = pa.Table.from_batches(self._record_batches, schema=self._schema) | |
| self._record_batches = [] | |
| self._record_batches_byte_size = 0 | |
| self.writer.write_table(table) | |
| def _flush_buffer(self) -> None: | |
| arrays: list[pa.array] = [[] for _ in range(len(self._schema.names))] | |
| for x, y in enumerate(self._buffer): | |
| arrays[x] = pa.array(y, type=self._schema.types[x]) | |
| self._buffer[x] = [] | |
| rb = pa.RecordBatch.from_arrays(arrays, schema=self._schema) | |
| self._record_batches.append(rb) | |
| size = 0 | |
| for x in arrays: | |
| for b in x.buffers(): # type: ignore | |
| if b is not None: | |
| size = size + b.size | |
| self._record_batches_byte_size = self._record_batches_byte_size + size | |