def test_from_to_buffer(df: pl.DataFrame, compressions: list[str]) -> None:
for compression in compressions:
buf = io.BytesIO()
df.write_ipc(buf, compression=compression) # type: ignore[arg-type]
buf.seek(0)
read_df = pl.read_ipc(buf)
assert_frame_equal_local_categoricals(df, 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_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_buffer(df: pl.DataFrame) -> None:
for buf in (io.BytesIO(), io.StringIO()):
df.write_json(buf)
buf.seek(0)
read_df = pl.read_json(buf)
read_df = read_df.with_columns(
[pl.col("cat").cast(pl.Categorical),
pl.col("time").cast(pl.Time)])
assert_frame_equal_local_categoricals(df, read_df)
def test_from_to_file(io_test_dir: str, df: pl.DataFrame,
compressions: list[str]) -> None:
f_ipc = os.path.join(io_test_dir, "small.ipc")
# does not yet work on windows because we hold an mmap?
if os.name != "nt":
for compression in compressions:
for f in (str(f_ipc), Path(f_ipc)):
df.write_ipc(f,
compression=compression) # type: ignore[arg-type]
df_read = pl.read_ipc(f) # type: ignore[arg-type]
assert_frame_equal_local_categoricals(df, df_read)
def test_to_from_buffer(df_no_lists: pl.DataFrame) -> None:
df = df_no_lists
buf = io.BytesIO()
df.write_csv(buf)
buf.seek(0)
read_df = pl.read_csv(buf, parse_dates=True)
read_df = read_df.with_columns(
[pl.col("cat").cast(pl.Categorical),
pl.col("time").cast(pl.Time)])
assert_frame_equal_local_categoricals(df, read_df)
def test_to_from_file(io_test_dir: str, df_no_lists: pl.DataFrame) -> None:
df = df_no_lists
df = df.drop("strings_nulls")
f = os.path.join(io_test_dir, "small.csv")
df.write_csv(f)
read_df = pl.read_csv(f, parse_dates=True)
read_df = read_df.with_columns(
[pl.col("cat").cast(pl.Categorical),
pl.col("time").cast(pl.Time)])
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_frame_equal_local_categoricals(df, read_df)