Repository: dask/dask
File: dask/dataframe/io/tests/test_parquet.py (Python)

Code before:
    ddf = dd.read_parquet(fn, index="index", engine=read_engine)
    assert_eq(dfp, ddf)


@PYARROW_MARK
def test_roundtrip_nullable_dtypes(tmp_path):
    """Test round-tripping nullable extension dtypes. Parquet engines will
    typically add dtype metadata for this.

    fastparquet is unsupported: https://github.com/dask/fastparquet/issues/465
    """
    df = pd.DataFrame(
        {
            "a": pd.Series([1, 2, pd.NA, 3, 4], dtype="Int64"),
            "b": pd.Series([True, pd.NA, False, True, False], dtype="boolean"),
            "c": pd.Series([0.1, 0.2, 0.3, pd.NA, 0.4], dtype="Float64"),
            "d": pd.Series(["a", "b", "c", "d", pd.NA], dtype="string"),
        }
    )
    ddf = dd.from_pandas(df, npartitions=2)
    ddf.to_parquet(tmp_path)
    ddf2 = dd.read_parquet(tmp_path)
    assert_eq(df, ddf2)


@pytest.mark.skipif(dd._dask_expr_enabled(), reason="not supported")
@PYARROW_MARK
@pytest.mark.parametrize(
    "dtype_backend",
    [
        "pandas",
        pytest.param(
            "pyarrow",
            marks=pytest.mark.skipif(
                not PANDAS_GE_150, reason="Requires pyarrow-backed nullable dtypes"
            ),
        ),
    ],
)
def test_use_nullable_dtypes(tmp_path, dtype_backend, engine):
    """
    Test reading a parquet file without pandas metadata,
    but forcing use of nullable dtypes where appropriate
    """
    dtype_extra = "" if dtype_backend == "pandas" else "[pyarrow]"
    df = pd.DataFrame(
        {
            "a": pd.Series([1, 2, pd.NA, 3, 4], dtype=f"Int64{dtype_extra}"),
            "b": pd.Series(
                [True, pd.NA, False, True, False], dtype=f"boolean{dtype_extra}"
            ),

Review comment (suggestion):
```suggestion
@pytest.mark.skipif(DASK_EXPR_ENABLED, reason="not supported")
```

Code after:
    ddf = dd.read_parquet(fn, index="index", engine=read_engine)
    assert_eq(dfp, ddf)


@PYARROW_MARK
def test_roundtrip_nullable_dtypes(tmp_path):
    """Test round-tripping nullable extension dtypes. Parquet engines will
    typically add dtype metadata for this.

    fastparquet is unsupported: https://github.com/dask/fastparquet/issues/465
    """
    df = pd.DataFrame(
        {
            "a": pd.Series([1, 2, pd.NA, 3, 4], dtype="Int64"),
            "b": pd.Series([True, pd.NA, False, True, False], dtype="boolean"),
            "c": pd.Series([0.1, 0.2, 0.3, pd.NA, 0.4], dtype="Float64"),
            "d": pd.Series(["a", "b", "c", "d", pd.NA], dtype="string"),
        }
    )
    ddf = dd.from_pandas(df, npartitions=2)
    ddf.to_parquet(tmp_path)
    ddf2 = dd.read_parquet(tmp_path)
    assert_eq(df, ddf2)


@pytest.mark.skipif(DASK_EXPR_ENABLED, reason="not supported")
@PYARROW_MARK
@pytest.mark.parametrize(
    "dtype_backend",
    [
        "pandas",
        pytest.param(
            "pyarrow",
            marks=pytest.mark.skipif(
                not PANDAS_GE_150, reason="Requires pyarrow-backed nullable dtypes"
            ),
        ),
    ],
)
def test_use_nullable_dtypes(tmp_path, dtype_backend, engine):
    """
    Test reading a parquet file without pandas metadata,
    but forcing use of nullable dtypes where appropriate
    """
    dtype_extra = "" if dtype_backend == "pandas" else "[pyarrow]"
    df = pd.DataFrame(
        {
            "a": pd.Series([1, 2, pd.NA, 3, 4], dtype=f"Int64{dtype_extra}"),
            "b": pd.Series(
                [True, pd.NA, False, True, False], dtype=f"boolean{dtype_extra}"
            ),