From 999bb193349af6457e739fe404102d7de1c5cf87 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 22 Mar 2024 12:41:41 +0000 Subject: [PATCH] Add test of #14159 --- .../cudf/cudf/tests/dataframe/test_io_serialization.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/cudf/cudf/tests/dataframe/test_io_serialization.py b/python/cudf/cudf/tests/dataframe/test_io_serialization.py index 911a7f9e865..3f086612654 100644 --- a/python/cudf/cudf/tests/dataframe/test_io_serialization.py +++ b/python/cudf/cudf/tests/dataframe/test_io_serialization.py @@ -2,6 +2,7 @@ from io import BytesIO import pandas as pd +import pyarrow as pa import pyarrow.parquet as pq import pytest @@ -35,3 +36,12 @@ def test_dataframe_parquet_roundtrip(index, write_index, empty): gpu_read = cudf.read_parquet(gpu_buf) cpu_read = cudf.read_parquet(cpu_buf) assert_eq(gpu_read, cpu_read) + + +@pytest.mark.parametrize("preserve_index", [False, True, None]) +def test_dataframe_to_arrow_preserve_index(preserve_index): + df = cudf.DataFrame({"x": ["cat", "dog"] * 5}) + pf = df.to_pandas() + expect = pa.Table.from_pandas(pf, preserve_index=preserve_index).schema + got = df.to_arrow(preserve_index=preserve_index).schema + assert expect == got