Add test of rapidsai#14159

wence- · Mar 22, 2024 · 3cc608b · 3cc608b
1 parent 39a2839
commit 3cc608b
Showing 1 changed file with 10 additions and 0 deletions.
diff --git 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