Function test_partition_pruning

dask/dataframe/dask_expr/io/tests/test_parquet.py:189–216

(tmpdir)

def test_partition_pruning(tmpdir):
    with dask.config.set({"dataframe.parquet.minimum-partition-size": 1}):
        filesystem = fs.LocalFileSystem()
        df = from_pandas(
            pd.DataFrame(
                {
                    "a": [1, 2, 3, 4, 5] * 10,
                    "b": range(50),
                }
            ),
            npartitions=2,
        )
        df.to_parquet(tmpdir, partition_on=["a"])
        ddf = read_parquet(tmpdir, filesystem=filesystem)
        ddf_filtered = read_parquet(
            tmpdir, filters=[[("a", "==", 1)]], filesystem=filesystem
        )
        assert ddf_filtered.npartitions == ddf.npartitions // 5

        ddf_optimize = read_parquet(tmpdir, filesystem=filesystem)
        ddf_optimize = ddf_optimize[ddf_optimize.a == 1].optimize()
        assert ddf_optimize.npartitions == ddf.npartitions // 5
        assert_eq(
            ddf_filtered,
            ddf_optimize,
            # FIXME ?
            check_names=False,
        )


def test_predicate_pushdown(tmpdir):