59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
import os
|
|
from typing import Optional
|
|
import polars as pl
|
|
import pyarrow.parquet as pq
|
|
|
|
|
|
class StreamingWriter:
|
|
def write_chunk(self, df: pl.DataFrame) -> None:
|
|
raise NotImplementedError
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
|
|
class StreamingCsvWriter(StreamingWriter):
|
|
def __init__(self, filepath: str):
|
|
self.filepath = filepath
|
|
self.is_first_chunk = True
|
|
os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
|
|
if os.path.exists(filepath):
|
|
os.remove(filepath)
|
|
|
|
def write_chunk(self, df: pl.DataFrame) -> None:
|
|
if df.is_empty():
|
|
return
|
|
with open(self.filepath, "ab") as f:
|
|
df.write_csv(f, include_header=self.is_first_chunk)
|
|
self.is_first_chunk = False
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
|
|
class StreamingParquetWriter(StreamingWriter):
|
|
def __init__(self, filepath: str, compression: str = "snappy"):
|
|
self.filepath = filepath
|
|
self.compression = compression
|
|
self.pq_writer: Optional[pq.ParquetWriter] = None
|
|
os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
|
|
if os.path.exists(filepath):
|
|
os.remove(filepath)
|
|
|
|
def write_chunk(self, df: pl.DataFrame) -> None:
|
|
if df.is_empty():
|
|
return
|
|
arrow_table = df.to_arrow()
|
|
if self.pq_writer is None:
|
|
self.pq_writer = pq.ParquetWriter(
|
|
self.filepath,
|
|
arrow_table.schema,
|
|
compression=self.compression,
|
|
)
|
|
self.pq_writer.write_table(arrow_table)
|
|
|
|
def close(self) -> None:
|
|
if self.pq_writer is not None:
|
|
self.pq_writer.close()
|
|
self.pq_writer = None
|