Skip to content

Commit

Permalink
Allow read_csv(header=None) to return int column labels in `mode.pand…
Browse files Browse the repository at this point in the history
…as_compatible` (rapidsai#16769)

closes rapidsai#16766

If the cudf `read_csv` behavior of always returning string column labels is long standing behavior, we can match the pandas behavior of returning integer column labels in `mode.pandas_compatible`

Authors:
  - Matthew Roeschke (https://github.com/mroeschke)

Approvers:
  - GALI PREM SAGAR (https://github.com/galipremsagar)

URL: rapidsai#16769
  • Loading branch information
mroeschke authored Sep 9, 2024
1 parent 4784067 commit 26a81b6
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 1 deletion.
4 changes: 3 additions & 1 deletion python/cudf/cudf/_lib/csv.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -276,8 +276,10 @@ def read_csv(
col_name = df._data.names[index]
df._data[col_name] = df._data[col_name].astype(col_dtype)

if names is not None and len(names) and isinstance(names[0], (int)):
if names is not None and len(names) and isinstance(names[0], int):
df.columns = [int(x) for x in df._data]
elif names is None and header == -1 and cudf.get_option("mode.pandas_compatible"):
df.columns = [int(x) for x in df._column_names]

# Set index if the index_col parameter is passed
if index_col is not None and index_col is not False:
Expand Down
8 changes: 8 additions & 0 deletions python/cudf/cudf/tests/test_csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -2269,3 +2269,11 @@ def test_read_compressed_BOM(tmpdir):
f.write(buffer)

assert_eq(pd.read_csv(fname), cudf.read_csv(fname))


def test_read_header_none_pandas_compat_column_type():
data = "1\n2\n"
with cudf.option_context("mode.pandas_compatible", True):
result = cudf.read_csv(StringIO(data), header=None).columns
expected = pd.read_csv(StringIO(data), header=None).columns
pd.testing.assert_index_equal(result, expected, exact=True)

0 comments on commit 26a81b6

Please sign in to comment.