Skip to content

Commit

Permalink
Restore column type metadata with dropna to fix factorize API (#1…
Browse files Browse the repository at this point in the history
…3980)

closes #13979 

This PR restores column type metadata for `dropna` call, absense of this restoration was causing an issue with the `CategoricalColumn.dropna` that was necessary for `factorize` API.

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

Approvers:
  - Bradley Dice (https://github.com/bdice)

URL: #13980
  • Loading branch information
galipremsagar authored Aug 28, 2023
1 parent aba001c commit d138dd0
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 2 deletions.
2 changes: 1 addition & 1 deletion python/cudf/cudf/core/column/column.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,7 @@ def any(self, skipna: bool = True) -> bool:

def dropna(self, drop_nan: bool = False) -> ColumnBase:
# The drop_nan argument is only used for numerical columns.
return drop_nulls([self])[0]
return drop_nulls([self])[0]._with_type_metadata(self.dtype)

def to_arrow(self) -> pa.Array:
"""Convert to PyArrow Array
Expand Down
20 changes: 19 additions & 1 deletion python/cudf/cudf/tests/test_factorize.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright (c) 2018-2022, NVIDIA CORPORATION.
# Copyright (c) 2018-2023, NVIDIA CORPORATION.

import cupy as cp
import numpy as np
Expand Down Expand Up @@ -139,3 +139,21 @@ def test_factorize_result_classes():

assert isinstance(labels, cp.ndarray)
assert isinstance(cats, cp.ndarray)


@pytest.mark.parametrize(
"data",
[
["abc", "def", "abc", "a", "def", None],
[10, 20, 100, -10, 0, 1, None, 10, 100],
],
)
def test_category_dtype_factorize(data):
gs = cudf.Series(data, dtype="category")
ps = gs.to_pandas()

actual_codes, actual_uniques = gs.factorize()
expected_codes, expected_uniques = ps.factorize()

assert_eq(actual_codes, expected_codes)
assert_eq(actual_uniques, expected_uniques)

0 comments on commit d138dd0

Please sign in to comment.