def test_parquet_stats(io_test_dir: str) -> None: file = path.join(io_test_dir, "binary_stats.parquet") df1 = pd.DataFrame({"a": [None, 1, None, 2, 3, 3, 4, 4, 5, 5]}) df1.to_parquet(file, engine="pyarrow") df = ( pl.scan_parquet(file) .filter(pl.col("a").is_not_null() & (pl.col("a") > 4)) .collect() ) assert df["a"].to_list() == [5.0, 5.0] assert ( pl.scan_parquet(file).filter(pl.col("a") > 4).select(pl.col("a").sum()) ).collect()[0, "a"] == 10.0 assert ( pl.scan_parquet(file).filter(pl.col("a") < 4).select(pl.col("a").sum()) ).collect()[0, "a"] == 9.0 assert ( pl.scan_parquet(file).filter(4 > pl.col("a")).select(pl.col("a").sum()) ).collect()[0, "a"] == 9.0 assert ( pl.scan_parquet(file).filter(4 < pl.col("a")).select(pl.col("a").sum()) ).collect()[0, "a"] == 10.0
def test_row_count(foods_parquet: str) -> None: df = pl.read_parquet(foods_parquet, row_count_name="row_count") assert df["row_count"].to_list() == list(range(27)) df = (pl.scan_parquet(foods_parquet, row_count_name="row_count").filter( pl.col("category") == pl.lit("vegetables")).collect()) assert df["row_count"].to_list() == [0, 6, 11, 13, 14, 20, 25] df = (pl.scan_parquet( foods_parquet, row_count_name="row_count").with_row_count( "foo", 10).filter(pl.col("category") == pl.lit("vegetables")).collect()) assert df["foo"].to_list() == [10, 16, 21, 23, 24, 30, 35]
def test_lazy_self_join_file_cache_prop_3979(io_test_dir: str) -> None: path = os.path.join(io_test_dir, "small.parquet") a = pl.scan_parquet(path) b = pl.DataFrame({"a": [1]}).lazy() assert a.join(b, how="cross").collect().shape == (3, 17) assert b.join(a, how="cross").collect().shape == (3, 17)
def test_categorical_parquet_statistics(io_test_dir: str) -> None: file = path.join(io_test_dir, "books.parquet") ( pl.DataFrame( { "book": [ "bookA", "bookA", "bookB", "bookA", "bookA", "bookC", "bookC", "bookC", ], "transaction_id": [1, 2, 3, 4, 5, 6, 7, 8], "user": ["bob", "bob", "bob", "tim", "lucy", "lucy", "lucy", "lucy"], } ) .with_column(pl.col("book").cast(pl.Categorical)) .write_parquet(file, statistics=True) ) for par in [True, False]: df = ( pl.scan_parquet(file, parallel=par) .filter(pl.col("book") == "bookA") .collect() ) assert df.shape == (4, 3)
def test_scan_parquet() -> None: df = pl.scan_parquet(Path(__file__).parent / "files" / "small.parquet") assert df.collect().shape == (4, 3)
def test_glob_parquet(io_test_dir: str) -> None: path = os.path.join(io_test_dir, "small*.parquet") assert pl.read_parquet(path).shape == (3, 16) assert pl.scan_parquet(path).collect().shape == (3, 16)