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"]
