Function test_parquet

dask/bytes/tests/test_s3.py:449–542

(s3, engine, s3so, metadata_file)

@pytest.mark.filterwarnings("ignore:Dask annotations")
@pytest.mark.parametrize("metadata_file", [True, False])
def test_parquet(s3, engine, s3so, metadata_file):
    dd = pytest.importorskip("dask.dataframe")
    pd = pytest.importorskip("pandas")
    np = pytest.importorskip("numpy")

    url = f"s3://{test_bucket_name}/test.parquet"
    data = pd.DataFrame(
        {
            "i32": np.arange(1000, dtype=np.int32),
            "i64": np.arange(1000, dtype=np.int64),
            "f": np.arange(1000, dtype=np.float64),
            "bhello": np.random.choice(["hello", "you", "people"], size=1000).astype(
                "O"
            ),
        },
        index=pd.Index(np.arange(1000), name="foo"),
    )
    df = dd.from_pandas(data, chunksize=500)
    df.to_parquet(
        url, engine=engine, storage_options=s3so, write_metadata_file=metadata_file
    )

    files = [f.split("/")[-1] for f in s3.ls(url)]
    if metadata_file:
        assert "_common_metadata" in files
        assert "_metadata" in files
    assert "part.0.parquet" in files

    df2 = dd.read_parquet(
        url, index="foo", calculate_divisions=True, engine=engine, storage_options=s3so
    )
    assert len(df2.divisions) > 1

    dd.utils.assert_eq(data, df2)

    # Check that `open_file_options` arguments are
    # really passed through to fsspec
    if fsspec_parquet:
        # Passing `open_file_options` kwargs will fail
        # if you set an unsupported engine
        with pytest.raises(ValueError):
            dd.read_parquet(
                url,
                engine=engine,
                storage_options=s3so,
                open_file_options={
                    "precache_options": {"method": "parquet", "engine": "foo"},
                },
            ).compute()

        # ...but should work fine if you modify the
        # maximum block-transfer size (max_block)
        dd.read_parquet(
            url,
            engine=engine,
            storage_options=s3so,
            open_file_options={
                "precache_options": {"method": "parquet", "max_block": 8_000},