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

Code before:
        ),
        ddf2,
        check_divisions=False,
        check_index=False,
    )


def test_read_inconsistent_schema_pyarrow(tmpdir):
    check_pyarrow()

    # Note: This is a proxy test for a cudf-related issue fix
    # (see cudf#5062 github issue).  The cause of that issue is
    # schema inconsistencies that do not actually correspond to
    # different types, but whether or not the file/column contains
    # null values.

    df1 = pd.DataFrame({"id": [0, 1], "val": [10, 20]})
    df2 = pd.DataFrame({"id": [2, 3], "val": [30, 40]})

    desired_type = "int64"
    other_type = "int32"
    df1.val = df1.val.astype(desired_type)
    df2.val = df2.val.astype(other_type)

    df_expect = pd.concat([df1, df2], ignore_index=True)
    df_expect['val'] = df_expect.val.astype(desired_type)

    df1.to_parquet(os.path.join(tmpdir, "0.parquet"))
    df2.to_parquet(os.path.join(tmpdir, "1.parquet"))

    # Read Directory
    check = dd.read_parquet(str(tmpdir), dataset={"validate_schema": False})
    assert_eq(check.compute(), df_expect, check_index=False)

    # Read List
    check = dd.read_parquet(
        os.path.join(tmpdir, "*.parquet"), dataset={"validate_schema": False}
    )
    assert_eq(check.compute(), df_expect, check_index=False)


def test_graph_size_pyarrow(tmpdir, engine):
    import pickle

    fn = str(tmpdir)

    ddf1 = dask.datasets.timeseries(
        start="2000-01-01", end="2000-01-02", freq="60S", partition_freq="1H"
    )

    ddf1.to_parquet(fn, engine=engine)

Review comment (suggestion):
Sorry, my suggestion introduced a lint error :) Let's try this.

```suggestion
    df_expect["val"] = df_expect.val.astype(desired_type)
```

Code after:
        ),
        ddf2,
        check_divisions=False,
        check_index=False,
    )


def test_read_inconsistent_schema_pyarrow(tmpdir):
    check_pyarrow()

    # Note: This is a proxy test for a cudf-related issue fix
    # (see cudf#5062 github issue).  The cause of that issue is
    # schema inconsistencies that do not actually correspond to
    # different types, but whether or not the file/column contains
    # null values.

    df1 = pd.DataFrame({"id": [0, 1], "val": [10, 20]})
    df2 = pd.DataFrame({"id": [2, 3], "val": [30, 40]})

    desired_type = "int64"
    other_type = "int32"
    df1.val = df1.val.astype(desired_type)
    df2.val = df2.val.astype(other_type)

    df_expect = pd.concat([df1, df2], ignore_index=True)
    df_expect["val"] = df_expect.val.astype(desired_type)

    df1.to_parquet(os.path.join(tmpdir, "0.parquet"))
    df2.to_parquet(os.path.join(tmpdir, "1.parquet"))

    # Read Directory
    check = dd.read_parquet(str(tmpdir), dataset={"validate_schema": False})
    assert_eq(check.compute(), df_expect, check_index=False)

    # Read List
    check = dd.read_parquet(
        os.path.join(tmpdir, "*.parquet"), dataset={"validate_schema": False}
    )
    assert_eq(check.compute(), df_expect, check_index=False)


def test_graph_size_pyarrow(tmpdir, engine):
    import pickle

    fn = str(tmpdir)

    ddf1 = dask.datasets.timeseries(
        start="2000-01-01", end="2000-01-02", freq="60S", partition_freq="1H"
    )

    ddf1.to_parquet(fn, engine=engine)