Repository: pola-rs/polars
File: py-polars/tests/unit/io/test_pyarrow_dataset.py (Python)

Code before:
def test_pyarrow_dataset_partial_predicate_pushdown(
    tmp_path: Path,
    plmonkeypatch: PlMonkeyPatch,
    capfd: pytest.CaptureFixture[str],
) -> None:
    plmonkeypatch.setenv("POLARS_VERBOSE_SENSITIVE", "1")

    df = pl.DataFrame({"a": [1, 2, 3], "b": [10.0, 20.0, 30.0]})
    file_path = tmp_path / "0"
    df.write_parquet(file_path)
    dset = ds.dataset(file_path, format="parquet")

    # col("a") > 1 is convertible; col("a") * col("b") > 25 is not (arithmetic
    # on two columns cannot be expressed as a pyarrow compute expression).
    # The optimizer pushes both terms into the scan's SELECTION, so our
    # MintermIter-based partial conversion should push the convertible part.
    q = pl.scan_pyarrow_dataset(dset).filter(
        (pl.col("a") > 1) & (pl.col("a") * pl.col("b") > 25)
    )

    capfd.readouterr()
    result = q.collect()
    capture = capfd.readouterr().err

    # Verify: partial predicate was pushed to pyarrow
    assert "(pa.compute.field('a') > 1)" in capture
    # Verify: correctness
    expected = (
        df.lazy().filter((pl.col("a") > 1) & (pl.col("a") * pl.col("b") > 25)).collect()
    )
    assert_frame_equal(result, expected)


def test_pyarrow_dataset_comm_subplan_elim(tmp_path: Path) -> None:
    df0 = pl.DataFrame({"a": [1, 2, 3]})

    df1 = pl.DataFrame({"a": [1, 2]})

    file_path_0 = tmp_path / "0.parquet"
    file_path_1 = tmp_path / "1.parquet"

    df0.write_parquet(file_path_0)
    df1.write_parquet(file_path_1)

    ds0 = ds.dataset(file_path_0, format="parquet")
    ds1 = ds.dataset(file_path_1, format="parquet")

    lf0 = pl.scan_pyarrow_dataset(ds0)
    lf1 = pl.scan_pyarrow_dataset(ds1)

    assert_frame_equal(

Review comment (suggestion):
After you address https://github.com/pola-rs/polars/pull/26567/changes#r2815863226, please also add an assertion here to verify that you see `is_full_predicate: false` in the output.


Code after:
def test_pyarrow_dataset_partial_predicate_pushdown(
    tmp_path: Path,
    plmonkeypatch: PlMonkeyPatch,
    capfd: pytest.CaptureFixture[str],
) -> None:
    plmonkeypatch.setenv("POLARS_VERBOSE_SENSITIVE", "1")

    df = pl.DataFrame({"a": [1, 2, 3], "b": [10.0, 20.0, 30.0]})
    file_path = tmp_path / "0"
    df.write_parquet(file_path)
    dset = ds.dataset(file_path, format="parquet")

    # col("a") > 1 is convertible; col("a") * col("b") > 25 is not (arithmetic
    # on two columns cannot be expressed as a pyarrow compute expression).
    # The optimizer pushes both terms into the scan's SELECTION, so our
    # MintermIter-based partial conversion should push the convertible part.
    q = pl.scan_pyarrow_dataset(dset).filter(
        (pl.col("a") > 1) & (pl.col("a") * pl.col("b") > 25)
    )

    capfd.readouterr()
    result = q.collect()
    capture = capfd.readouterr().err

    # Verify: partial predicate was pushed to pyarrow
    assert "(pa.compute.field('a') > 1)" in capture
    assert (
        'residual predicate: Some([([(col("a").cast(Float64)) * (col("b"))]) > (25.0)])'
        in capture
    )
    # Verify: correctness
    expected = (
        df.lazy().filter((pl.col("a") > 1) & (pl.col("a") * pl.col("b") > 25)).collect()
    )
    assert_frame_equal(result, expected)


def test_pyarrow_dataset_comm_subplan_elim(tmp_path: Path) -> None:
    df0 = pl.DataFrame({"a": [1, 2, 3]})

    df1 = pl.DataFrame({"a": [1, 2]})

    file_path_0 = tmp_path / "0.parquet"
    file_path_1 = tmp_path / "1.parquet"

    df0.write_parquet(file_path_0)
    df1.write_parquet(file_path_1)

    ds0 = ds.dataset(file_path_0, format="parquet")
    ds1 = ds.dataset(file_path_1, format="parquet")