Esempio n. 1
0
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)
Esempio n. 2
0
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)
Esempio n. 3
0
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)