Repository: apache/arrow
File: python/pyarrow/tests/parquet/test_basic.py (Python)

Code before:
            "d": [True, False, True],
            "e": pd.date_range("20130101", periods=3),
            "f": pd.Categorical([5, 6, 7]),
            # fastparquet writes list as BYTE_ARRAY JSON, so no roundtrip
            # "g": [[1, 2], None, [1, 2, 3]],
        }
    )
    table = pa.table(df)

    # Arrow -> fastparquet
    file_arrow = str(tempdir / "cross_compat_arrow.parquet")
    pq.write_table(table, file_arrow, compression=None)

    fp_file = fp.ParquetFile(file_arrow)
    df_fp = fp_file.to_pandas()
    # TODO: once fastparquet supports pandas 3 dtypes revert string and categorical
    # tests by removing `check_dtype=False` and `check_categorical=False` so type
    # equality is asserted again
    tm.assert_frame_equal(df, df_fp, check_dtype=False, check_categorical=False)

    # Fastparquet -> arrow
    file_fastparquet = str(tempdir / "cross_compat_fastparquet.parquet")
    # fastparquet can't write pandas 3.0 StringDtype
    df_for_fp = df.copy()
    df_for_fp['a'] = df_for_fp['a'].astype(object)
    fp.write(file_fastparquet, df_for_fp)

    table_fp = pq.read_pandas(file_fastparquet)
    # for fastparquet written file, categoricals comes back as strings
    # (no arrow schema in parquet metadata)
    tm.assert_frame_equal(table_fp.to_pandas(), df_for_fp, check_dtype=False,
                          check_categorical=False)


@pytest.mark.parametrize('array_factory', [
    lambda: pa.array([0, None] * 10),
    lambda: pa.array([0, None] * 10).dictionary_encode(),
    lambda: pa.array(["", None] * 10),
    lambda: pa.array(["", None] * 10).dictionary_encode(),
])
@pytest.mark.parametrize('read_dictionary', [False, True])
def test_buffer_contents(
        array_factory, read_dictionary
):
    # Test that null values are deterministically initialized to zero
    # after a roundtrip through Parquet.
    # See ARROW-8006 and ARROW-8011.
    orig_table = pa.Table.from_pydict({"col": array_factory()})
    bio = io.BytesIO()
    pq.write_table(orig_table, bio, use_dictionary=True)
    bio.seek(0)

Review comment (suggestion):
Perhaps we can avoid creating a new dataframe and only non-destructively cast at write-time?
```suggestion
    # fastparquet can't write pandas 3.0 StringDtype
    fp_compatible_types =  {"a": object}
    fp.write(file_fastparquet, df.astype(fp_compatible_types))
```

Code after:
            "d": [True, False, True],
            "e": pd.date_range("20130101", periods=3),
            "f": pd.Categorical(["a", "b", "a"]),
            # fastparquet writes list as BYTE_ARRAY JSON, so no roundtrip
            # "g": [[1, 2], None, [1, 2, 3]],
        }
    )
    table = pa.table(df)

    # Arrow -> fastparquet
    file_arrow = str(tempdir / "cross_compat_arrow.parquet")
    pq.write_table(table, file_arrow, compression=None)

    fp_file = fp.ParquetFile(file_arrow)
    df_fp = fp_file.to_pandas()
    # pandas 3 defaults to StringDtype for strings, fastparquet still returns object
    # TODO: remove astype casts once fastparquet supports pandas 3 StringDtype
    tm.assert_frame_equal(df_fp, df.astype({"a": object}))

    # Fastparquet -> arrow
    file_fastparquet = str(tempdir / "cross_compat_fastparquet.parquet")
    # fastparquet doesn't support writing pandas 3 StringDtype yet
    fp.write(file_fastparquet, df.astype({"a": object}))

    table_fp = pq.read_pandas(file_fastparquet)
    # for fastparquet written file, categoricals comes back as strings
    # (no arrow schema in parquet metadata)
    tm.assert_frame_equal(table_fp.to_pandas(), df.astype({"f": object}))


@pytest.mark.parametrize('array_factory', [
    lambda: pa.array([0, None] * 10),
    lambda: pa.array([0, None] * 10).dictionary_encode(),
    lambda: pa.array(["", None] * 10),
    lambda: pa.array(["", None] * 10).dictionary_encode(),
])
@pytest.mark.parametrize('read_dictionary', [False, True])
def test_buffer_contents(
        array_factory, read_dictionary
):
    # Test that null values are deterministically initialized to zero
    # after a roundtrip through Parquet.
    # See ARROW-8006 and ARROW-8011.
    orig_table = pa.Table.from_pydict({"col": array_factory()})
    bio = io.BytesIO()
    pq.write_table(orig_table, bio, use_dictionary=True)
    bio.seek(0)
    read_dictionary = ['col'] if read_dictionary else None
    table = pq.read_table(bio, use_threads=False,
                          read_dictionary=read_dictionary)