Skip to content

Implement factorize for MultiIndex #23139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 5 commits into from
Jul 8, 2026
Merged

Implement factorize for MultiIndex #23139

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ca37851
Implement factorize for MultiIndex
galipremsagar Jul 7, 2026
a1570f7
Simplify position-column assignment in MultiIndex._factorize
galipremsagar Jul 7, 2026
bf33cd4
Merge branch 'main' into mi_factorize
galipremsagar Jul 7, 2026
8b8dd84
Merge remote-tracking branch 'upstream/main' into mi_factorize
galipremsagar Jul 7, 2026
4d6588e
Merge remote-tracking branch 'upstream/main' into mi_factorize
galipremsagar Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions python/cudf/cudf/core/multiindex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1793,6 +1793,55 @@ def unique(self, level: int | None = None) -> Self | Index:
else:
return self.get_level_values(level).unique()

def _factorize(
self, sort: bool, use_na_sentinel: bool
) -> tuple[cp.ndarray, MultiIndex]:
if any(col.has_nulls() for col in self._columns):
raise NotImplementedError(
"factorize on a MultiIndex with missing values is not yet "
"supported"
)
if len(self) == 0:
return cp.empty(0, dtype=np.intp), self.copy()
Comment on lines +1804 to +1805

@coderabbitai coderabbitai Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Empty MultiIndex factorize doesn't strip level names, unlike the non-empty path.

For the non-empty path, uniques_mi.names is explicitly reset to [None] * self.nlevels to match pandas (per the comment at line 1838). But the empty branch returns self.copy(), which preserves the original names. For a named empty MultiIndex, this makes uniques inconsistent with the non-empty behavior (and, per the code's own rationale, with pandas). The current empty test doesn't catch this because it constructs an unnamed MultiIndex.

🐛 Proposed fix 
         if len(self) == 0:
-            return cp.empty(0, dtype=np.intp), self.copy()
+            empty = self.copy()
+            empty.names = [None] * self.nlevels
+            return cp.empty(0, dtype=np.intp), empty
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if len(self) == 0:
return cp.empty(0, dtype=np.intp), self.copy()
if len(self) == 0:
empty = self.copy()
empty.names = [None] * self.nlevels
return cp.empty(0, dtype=np.intp), empty
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudf/cudf/core/multiindex.py` around lines 1804 - 1805, The empty
MultiIndex path in `MultiIndex.factorize` returns `self.copy()`, which preserves
level names and makes the `uniques` result inconsistent with the non-empty
branch. Update the `len(self) == 0` branch so the returned empty `MultiIndex`
mirrors the non-empty behavior by clearing names to `[None] * self.nlevels`,
using the existing `MultiIndex.factorize`/`self.copy()` flow as the place to
apply the fix.


level_labels = list(range(self.nlevels))
df = cudf.DataFrame._from_data(
dict(zip(level_labels, self._columns, strict=True))
)
pos_label = self.nlevels
df[pos_label] = as_column(range(len(df)))
Comment thread
galipremsagar marked this conversation as resolved.
Outdated

# the uniques are the distinct rows in order of first appearance
# (or sorted lexicographically when requested)
uniques = (
df.groupby(level_labels, sort=False)
.agg({pos_label: "min"})
.reset_index()
)
uniques = uniques.sort_values(
by=level_labels if sort else pos_label, ignore_index=True
)
uid_label = self.nlevels + 1
uniques[uid_label] = as_column(range(len(uniques)))

# codes: map each row to its unique-row id, in the original order
merged = df.merge(
uniques[[*level_labels, uid_label]],
on=level_labels,
how="left",
)
codes = (
merged.sort_values(by=pos_label)[uid_label]
.astype(np.dtype(np.intp))
.values
)
# pandas does not propagate the level names to the uniques
uniques_mi = MultiIndex._from_data(
{label: uniques._data[label] for label in level_labels}
)
uniques_mi.names = [None] * self.nlevels
return codes, uniques_mi

@_performance_tracking
def nunique(self, dropna: bool = True) -> int:
mi = self.dropna(how="all") if dropna else self
Expand Down
7 changes: 1 addition & 6 deletions python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4381,13 +4381,8 @@ def pytest_unconfigure(config):
"tests/strings/test_split_partition.py::test_split_nan_expand[string=object]": "AssertionError: DataFrame.iloc[:, 0] (column name='0') are different",
"tests/strings/test_strings.py::test_index_str_accessor_multiindex_raises": "TODO: Add a reason for failure",
"tests/strings/test_strings.py::test_string_slice_out_of_bounds[string=object]": "AssertionError: Series are different",
"tests/test_algos.py::TestFactorize::test_factorize[multi-False]": "NotImplementedError: Fast implementation not available. Falling back to the slow implementation",
"tests/test_algos.py::TestFactorize::test_factorize[multi-True]": "NotImplementedError: Fast implementation not available. Falling back to the slow implementation",
"tests/test_algos.py::TestFactorize::test_factorize[tuples-False]": "AssertionError: Index are different",
"tests/test_algos.py::TestFactorize::test_factorize[tuples-True]": "AssertionError: Index are different",
"tests/test_algos.py::TestFactorize::test_int_factorize_use_na_sentinel_false[data0-expected_codes0-expected_uniques0]": "TODO: Add a reason for failure",
"tests/test_algos.py::TestFactorize::test_int_factorize_use_na_sentinel_false[data0-expected_codes0-expected_uniques0]": "cudf cannot represent an object column of Python ints with NaN; the uniques come back float64",
"tests/test_algos.py::TestUnique::test_datetime64_dtype_array_returned": "TODO: Add a reason for failure",
"tests/test_algos.py::TestUnique::test_factorize_multiindex_empty": "AssertionError: Index are different",
"tests/test_algos.py::TestUnique::test_order_of_appearance_dt64tz[ms]": "AssertionError: Index are different",
"tests/test_algos.py::TestUnique::test_order_of_appearance_dt64tz[ns]": "AssertionError: Index are different",
"tests/test_algos.py::TestUnique::test_order_of_appearance_dt64tz[s]": "AssertionError: Index are different",
Expand Down
30 changes: 29 additions & 1 deletion python/cudf/cudf/tests/general_functions/test_factorize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import cupy as cp
import numpy as np
import pandas as pd
import pytest

import cudf
from cudf.testing import assert_eq
Expand Down Expand Up @@ -86,3 +87,30 @@ def test_factorize_rangeindex_preserves_class():
assert isinstance(cats, cudf.RangeIndex)
assert_eq(cats, p_cats, exact=True)
np.testing.assert_array_equal(labels.get(), p_labels)


@pytest.mark.parametrize("sort", [False, True])
def test_factorize_multiindex(sort):
# The uniques are the distinct rows (without level names), matching
# pandas.
pmi = pd.MultiIndex.from_arrays(
[["bar", "baz", "foo", "bar"], [2, 3, 1, 2]], names=["a", "b"]
)
gmi = cudf.Index(pmi)

expected_codes, expected_uniques = pmi.factorize(sort=sort)
codes, uniques = gmi.factorize(sort=sort)

np.testing.assert_array_equal(codes.get(), expected_codes)
assert_eq(uniques, expected_uniques)


def test_factorize_multiindex_empty():
pmi = pd.MultiIndex.from_arrays(
[pd.Index([], dtype=object), pd.Index([], dtype="f4")]
)
gmi = cudf.Index(pmi)

codes, uniques = gmi.factorize()
assert len(codes) == 0
assert_eq(uniques, gmi[:0])
Loading