Honor level and sort_remaining in sort_index(axis=1) - #23367
Conversation
sort_index(axis=1) silently ignored level= and sort_remaining= and always sorted by the full column labels. Sort by the requested levels (stable multi-key, least-significant first), append the remaining levels when sort_remaining=True, resolve integer and named levels with pandas' bounds validation, and place missing labels per na_position independently of the per-key direction. Fixes 1 pandas-test (test_stack_mixed_dtype[True]); its xfail entry is removed. Attribution verified against an isolated build containing only this change.
github-actions
Bot
added
Python
galipremsagar
added
bug
|
/okay to test 068e70c |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
📝 WalkthroughSummary by CodeRabbit
Walkthrough
ChangesAxis-1 sort_index behavior
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with 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.
Inline comments:
In `@python/cudf/cudf/core/indexed_frame.py`:
- Around line 2931-2941: Update the nested _level_number helper in sort_index to
convert a missing named level from the level_names.index(lvl) ValueError into
KeyError with the message “Level {lvl} not found”. Preserve the existing
integer-level validation and successful name lookup behavior.
- Around line 2954-2959: Update the sort_remaining branch in the indexed-frame
sorting logic to append a fixed ascending value for unspecified levels instead
of reusing ascending_per_key[0]. Preserve explicitly provided per-level
directions, and add coverage for list-valued ascending with sort_remaining=True
to verify remaining levels sort ascending.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b5bac8d1-bdd4-451f-b934-fc702af4fbe2
📒 Files selected for processing (3)
python/cudf/cudf/core/indexed_frame.pypython/cudf/cudf/pandas/scripts/pandas-testing-plugin.pypython/cudf/cudf/tests/dataframe/methods/test_sort_index.py
💤 Files with no reviewable changes (1)
- python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py
| def _level_number(lvl): | ||
| if isinstance(lvl, int): | ||
| norm = lvl + nlevels if lvl < 0 else lvl | ||
| if not 0 <= norm < nlevels: | ||
| raise IndexError( | ||
| f"Too many levels: Index has only " | ||
| f"{nlevels} levels, {lvl} is not a valid " | ||
| "level number" | ||
| ) | ||
| return norm | ||
| return level_names.index(lvl) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant implementation and nearby tests
ast-grep outline python/cudf/cudf/core/indexed_frame.py --view expanded | sed -n '1,220p'
printf '\n--- surrounding lines ---\n'
sed -n '2900,2960p' python/cudf/cudf/core/indexed_frame.py
printf '\n--- search for level lookup handling/tests ---\n'
rg -n "Level .* not found|_get_level_number|level_names\.index|sort_index\\(level=|KeyError" python/cudf -g '!**/__pycache__/**'Repository: rapidsai/cudf
Length of output: 23013
Named-level sort_index should raise KeyError
python/cudf/cudf/core/indexed_frame.py:2931-2941 — level_names.index(lvl) currently lets a missing level name bubble up as ValueError. Catch that here and re-raise KeyError(f"Level {lvl} not found") to match pandas and callers that handle missing levels via KeyError.
🤖 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/indexed_frame.py` around lines 2931 - 2941, Update the
nested _level_number helper in sort_index to convert a missing named level from
the level_names.index(lvl) ValueError into KeyError with the message “Level
{lvl} not found”. Preserve the existing integer-level validation and successful
name lookup behavior.
| if sort_remaining: | ||
| seen = set(key_order) | ||
| for i in range(nlevels): | ||
| if i not in seen: | ||
| key_order.append(i) | ||
| ascending_per_key.append(ascending_per_key[0]) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
For pandas DataFrame.sort_index with a MultiIndex, level as a list, ascending as a per-level list, and sort_remaining=True, what ascending direction is applied to the levels not listed in level?
💡 Result:
When using pandas.DataFrame.sort_index with a MultiIndex, a specified level list, and sort_remaining=True, the levels not included in the level argument are sorted in ascending order [1][2]. The sort_remaining=True parameter ensures that after the specified levels have been sorted according to the directions provided in the ascending list, all other levels in the MultiIndex are also sorted [1][2]. The default sort order for these remaining levels is ascending [1][2]. While the ascending parameter allows for individual control over the sort direction of specified levels (when passed as a list of booleans) [1][3], it does not alter the default behavior applied to the remaining levels when sort_remaining=True is active; those levels will be sorted in ascending order regardless of the directions chosen for the specified levels [1][2].
Citations:
- 1: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.sort_index.html
- 2: https://pandas.pydata.org/docs/reference/api/pandas.MultiIndex.sortlevel.html
- 3: https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.sort_index.html?highlight=sort_index
🏁 Script executed:
python3 - <<'PY'
from pathlib import Path
path = Path("python/cudf/cudf/core/indexed_frame.py")
text = path.read_text().splitlines()
for start in [2920, 2940, 2950, 2965, 2990]:
end = min(len(text), start + 40)
print(f"\n--- lines {start+1}-{end} ---")
for i in range(start, end):
print(f"{i+1:4}: {text[i]}")
PYRepository: rapidsai/cudf
Length of output: 10025
🏁 Script executed:
python3 - <<'PY'
from pathlib import Path
path = Path("python/cudf/cudf/core/indexed_frame.py")
lines = path.read_text().splitlines()
for start, end in [(2890, 2975)]:
print(f"--- {start+1}-{end} ---")
for i in range(start, min(end, len(lines))):
print(f"{i+1:4}: {lines[i]}")
PYRepository: rapidsai/cudf
Length of output: 4400
🏁 Script executed:
python3 - <<'PY'
import sys
try:
import pandas as pd
except Exception as e:
print(f"pandas import failed: {e}")
sys.exit(0)
print("pandas", pd.__version__)
mi = pd.MultiIndex.from_tuples([
("a", 2, "x"),
("a", 1, "y"),
("a", 1, "x"),
("b", 2, "y"),
("b", 1, "x"),
("b", 1, "y"),
], names=["L0", "L1", "L2"])
df = pd.DataFrame({"v": range(len(mi))}, index=mi)
for asc in ([False, True], [True, False], [False, False], [True, True]):
try:
out = df.sort_index(level=["L0", "L1"], ascending=asc, sort_remaining=True)
print("\nascending =", asc)
print(list(out.index))
except Exception as e:
print("\nascending =", asc, "raised", type(e).__name__, e)
# Also probe sortlevel on MultiIndex directly if available
if hasattr(mi, "sortlevel"):
print("\nMultiIndex.sortlevel probes:")
for asc in ([False, True], [True, False], [False, False], [True, True]):
try:
idx, _ = mi.sortlevel(level=["L0", "L1"], ascending=asc, sort_remaining=True)
print("ascending =", asc, list(idx))
except Exception as e:
print("ascending =", asc, "raised", type(e).__name__, e)
PYRepository: rapidsai/cudf
Length of output: 198
🏁 Script executed:
python3 - <<'PY'
import sys
try:
import pandas as pd
except Exception as e:
print(f"pandas import failed: {e}")
sys.exit(0)
print("pandas", pd.__version__)
mi = pd.MultiIndex.from_tuples([
("a", 2, "x"),
("a", 1, "y"),
("a", 1, "x"),
("b", 2, "y"),
("b", 1, "x"),
("b", 1, "y"),
], names=["L0", "L1", "L2"])
df = pd.DataFrame({"v": range(len(mi))}, index=mi)
for asc in ([False, True], [True, False], [False, False], [True, True]):
try:
out = df.sort_index(level=["L0", "L1"], ascending=asc, sort_remaining=True)
print("\nascending =", asc)
print(list(out.index))
except Exception as e:
print("\nascending =", asc, "raised", type(e).__name__, e)
if hasattr(mi, "sortlevel"):
print("\nMultiIndex.sortlevel probes:")
for asc in ([False, True], [True, False], [False, False], [True, True]):
try:
idx, _ = mi.sortlevel(level=["L0", "L1"], ascending=asc, sort_remaining=True)
print("ascending =", asc, list(idx))
except Exception as e:
print("ascending =", asc, "raised", type(e).__name__, e)
PYRepository: rapidsai/cudf
Length of output: 198
Keep sort_remaining levels ascending
When ascending is a list, this appends ascending_per_key[0], so ascending=[False, True] makes the remaining levels sort descending. sort_remaining=True should leave unspecified levels ascending; add a test for list ascending + sort_remaining=True and use a fixed ascending default for the appended levels.
🤖 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/indexed_frame.py` around lines 2954 - 2959, Update the
sort_remaining branch in the indexed-frame sorting logic to append a fixed
ascending value for unspecified levels instead of reusing ascending_per_key[0].
Preserve explicitly provided per-level directions, and add coverage for
list-valued ascending with sort_remaining=True to verify remaining levels sort
ascending.
vyasr
left a comment
vyasr
left a comment
There was a problem hiding this comment.
Some of the behavior here is a bit confusing.
Related, this is a lot of edge case handling to pass a single pandas test :(
|
|
||
| key_order = [_level_number(lvl) for lvl in level] | ||
| ascending_per_key: list[bool] = ( | ||
| [bool(flag) for flag in cast("Iterable[bool]", ascending)] |
There was a problem hiding this comment.
Why are we casting an iterable of bools to a single bool here? What is the intended semantic? I'm confused by this behavior.
There was a problem hiding this comment.
Fair — all three of these confusion points were symptoms of hand-rolling pandas' label-sort semantics. I've deleted the whole block: since column labels are host-side pandas metadata anyway, the new code has pandas itself compute the order (a pd.Series(range(n), index=columns).sort_index(...) yields the positional indexer), inheriting pandas' exact level resolution/validation, per-level ascending, sort_remaining, and na_position semantics with no edge-case code of our own. Net diff is now smaller than the original one-liner path it replaced was hiding: the old sorted(self._column_names) fallback also couldn't honor na_position for NaN labels, which this fixes for free. Also expanded the classic test coverage well beyond the one pandas test: per-level ascending lists, ascending-length mismatch, unknown level names, flat columns with level=0/name, and NaN column labels with na_position on the no-level path (~1090 params total, all compared against pandas), plus the full pandas-tests test_sort_index.py file runs clean.
| for i in range(nlevels): | ||
| if i not in seen: | ||
| key_order.append(i) | ||
| ascending_per_key.append(ascending_per_key[0]) |
There was a problem hiding this comment.
Why do we append ascending_per_key[0] and not ascending_per_key[i]?
There was a problem hiding this comment.
Gone with the rewrite — pandas' own sort_remaining handling applies now. (For the record, the [0] mirrored what pandas does internally: lexsort_indexer extends a scalar-ish ascending with the first spec's direction for the remaining levels — but encoding that by hand was exactly the kind of trivia this code shouldn't own.)
| def _is_na_label(value) -> bool: | ||
| return value is None or ( | ||
| isinstance(value, float) and value != value | ||
| ) |
There was a problem hiding this comment.
This code is removed in the rewrite.
Replace the hand-rolled multi-key label sort with a host-side pandas computation: a Series holding the original positions, indexed by the column labels, sorted by its index yields the positional indexer while inheriting pandas' exact level resolution and validation, per-level ascending, sort_remaining, and na_position semantics. This also fixes the plain (no level) axis=1 path, which used python sorted() and could not honor na_position for NaN labels. Expands test coverage: per-level ascending lists, ascending-length mismatch, unknown level names, flat columns with level 0/name, and NaN column labels with na_position on the no-level path.
|
Reworked in 98123f4 per the review: the hand-rolled multi-key label sort (all ~65 lines of edge-case handling) is deleted. The column labels are host-side pandas metadata, so the axis=1 path now asks pandas for the sorted order directly — Test coverage is expanded well beyond the single pandas test: per-level ascending lists, ascending-length-mismatch |
|
/okay to test 097a25a |
galipremsagar
added
5 - Ready to Merge
|
/okay to test f1286b0 |
|
/merge |
4f082f3
into
NVIDIA:release/26.08
4 participants
Description
Split out of #23255 (6/6).
sort_index(axis=1)silently ignoredlevel=andsort_remaining=and always sorted by the full column labels. Sort by the requested levels (stable multi-key, least significant first), append the remaining levels whensort_remaining=True, resolve integer and named levels with pandas' bounds validation, and place missing labels perna_positionindependently of the per-key sort direction.Fixes 1 pandas-test (
test_stack_mixed_dtype[True], which sorts the stacked frame's columns by level); its xfail entry is removed. Attribution verified by running the node id against an isolated build containing only this change (passes) and a clean build (fails).Independent of the other #23255 split PRs; can merge in any order.
Checklist