def test_to_from_buffer(df: pl.DataFrame, compressions: list[str]) -> None:
for compression in compressions:
if compression == "lzo":
# lzo compression is not supported now
with pytest.raises(pl.ArrowError):
buf = io.BytesIO()
df.write_parquet(buf, compression=compression)
buf.seek(0)
_ = pl.read_parquet(buf)
with pytest.raises(OSError):
buf = io.BytesIO()
df.write_parquet(buf,
compression=compression,
use_pyarrow=True)
buf.seek(0)
_ = pl.read_parquet(buf)
else:
buf = io.BytesIO()
df.write_parquet(buf, compression=compression)
buf.seek(0)
read_df = pl.read_parquet(buf)
assert_frame_equal_local_categoricals(df, read_df)
for use_pyarrow in [True, False]:
buf = io.BytesIO()
df.write_parquet(buf, use_pyarrow=use_pyarrow)
buf.seek(0)
read_df = pl.read_parquet(buf, use_pyarrow=use_pyarrow)
assert_frame_equal_local_categoricals(df, read_df)
def test_to_from_file(io_test_dir: str, df: pl.DataFrame,
compressions: List[str]) -> None:
f = os.path.join(io_test_dir, "small.parquet")
for compression in compressions:
if compression == "lzo":
# lzo compression is not supported now
with pytest.raises(pl.ArrowError):
df.write_parquet(f, compression=compression)
_ = pl.read_parquet(f)
with pytest.raises(OSError):
df.write_parquet(f, compression=compression, use_pyarrow=True)
_ = pl.read_parquet(f)
else:
df.write_parquet(f, compression=compression)
read_df = pl.read_parquet(f)
assert df.frame_equal(read_df)
def test_to_from_buffer(df: pl.DataFrame, compressions: List[str]) -> None:
for compression in compressions:
if compression == "lzo":
# lzo compression is not supported now
with pytest.raises(pl.ArrowError):
buf = io.BytesIO()
df.write_parquet(buf, compression=compression)
buf.seek(0)
_ = pl.read_parquet(buf)
with pytest.raises(OSError):
buf = io.BytesIO()
df.write_parquet(buf,
compression=compression,
use_pyarrow=True)
buf.seek(0)
_ = pl.read_parquet(buf)
else:
buf = io.BytesIO()
df.write_parquet(buf, compression=compression)
buf.seek(0)
read_df = pl.read_parquet(buf)
assert df.frame_equal(read_df, null_equal=True)