def test_to_from_buffer(df: pl.DataFrame) -> None: df = df.drop("strings_nulls") for to_fn, from_fn, text_based in zip( [df.to_parquet, df.to_csv, df.to_ipc, df.to_json], [ pl.read_parquet, partial(pl.read_csv, parse_dates=True), pl.read_ipc, pl.read_json, ], [False, True, False, True], ): f = io.BytesIO() to_fn(f) # type: ignore f.seek(0) df_1 = from_fn(f) # type: ignore # some type information is lost due to text conversion if text_based: df_1 = df_1.with_columns([ pl.col("cat").cast(pl.Categorical), pl.col("time").cast(pl.Time) ]) assert df.frame_equal(df_1)
def test_drop(): df = DataFrame({"a": [2, 1, 3], "b": ["a", "b", "c"], "c": [1, 2, 3]}) df = df.drop("a") assert df.shape == (3, 2) df = DataFrame({"a": [2, 1, 3], "b": ["a", "b", "c"], "c": [1, 2, 3]}) s = df.drop_in_place("a") assert s.name == "a"
def test_to_from_file(io_test_dir: str, df: pl.DataFrame) -> None: 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 df.frame_equal(read_df)
def test_multiple_columns_drop(): df = DataFrame({"a": [2, 1, 3], "b": [1, 2, 3], "c": [1, 2, 3]}) out = df.drop(["a", "b"]) assert out.columns == ["c"]