feat: support passing index object directly into maybe_set_index

[Riik]

What type of PR is this? (check all applicable)

• 💾 Refactor
• ✨ Feature
• 🐛 Bug Fix
• 🔧 Optimization
• 📝 Documentation
• ✅ Test
• 🐳 Other

Related issues

• Related issue [Enh]: allow passing index object directly to maybe_set_index #1239
• Closes #

Checklist

• Code follows style guide (ruff)
• Tests added
• Documented the changes

If you have comments or can explain your changes, please do so below.

Pandas set_index docs
Pandas set_index test cases
I've renamed the column_names: str | list[str] parameter to keys: str | Series | list[Series | str]. I referenced the Pandas docs where which have keys: label or array-like or list of labels/arrays. Since this is a backwards-compatability function, I figured it is good for it to be close to the Pandas spec, and the original param name column doesn't cover all cases anymore.

[MarcoGorelli]

thanks for your PR!

this would break backwards-compatibility, so I don't think we can just rename to keys unfortunately

[Riik]

That makes sense, I wondered if this was part of the "stable API" or not. Would you advise to stick with the original name (columns) and expand the functionality/meaning, or should we not add this feature and close the issue / PR?

[MarcoGorelli]

one idea could be to add an extra keyword argument index which you can pass an index object directly to (and the disallow passing both columns and index), that would still be backwards-compatible

[Riik]

Right, so we preserve the column_names for backwards comparability, add index, make both optional, and check that they are never both specified. They new signature will then be

def maybe_set_index(df: T, column_names: str | list[str] | None, index: Series | list[Series] | None) -> T

I'll give that a try!

[FBruzzesi]

Hey @Riik thanks for the effort! I left a minor comment and a suggestion to simplify the logic a bit.

Additionally, could we handle the case for which df is actually a series? That's the use case I had in mind when I opened the related issue. I.e. if the native_object is a pandas like series, then I would like to modify its index and clearly that's not possible to do unless another Series/Index/iterable is passed.

Since that's the case, I think we should also raise an error if native_object is a series but column_names is provided

[Riik]

Additionally, could we handle the case for which df is actually a series? That's the use case I had in mind when I opened the related issue. I.e. if the native_object is a pandas like series, then I would like to modify its index and clearly that's not possible to do unless another Series/Index/iterable is passed.

Ah yes, I see the idea now! I think that should work with the current implementation, let me add a test for it

[Riik]

@FBruzzesi I've added a test case and also support for setting the index on a Series. It didn't work out of the box since pd.Series (docs) doesn't have a set_index() method, but you can set the index directly by assignment.

Also, there is the extra case where passing in column_names in combination with a Series is now invalid.

I feel like this function is getting a bit unwieldy, i.e. there a lot of conditionals inside the function. Would it be an option to make a separate function maybe_set_index_series instead? That would clean up the logic of this function and keep it simpler. Or maybe there is other ways to clean up this function a bit more. Let me know if you have any ideas!

[FBruzzesi]

@FBruzzesi I've added a test case and also support for setting the index on a Series. It didn't work out of the box since pd.Series (docs) doesn't have a set_index() method, but you can set the index directly by assignment.

That's great, thanks! It looks good, I just tagged Marco asking if Series.set_axis should be prefered instead.

Also, there is the extra case where passing in column_names in combination with a Series is now invalid.

I feel like this function is getting a bit unwieldy, i.e. there a lot of conditionals inside the function. Would it be an option to make a separate function maybe_set_index_series instead? That would clean up the logic of this function and keep it simpler. Or maybe there is other ways to clean up this function a bit more. Let me know if you have any ideas!

It still seems pretty reasonable in scope and size. I will take a deeper look tomorrow to see if there is a way to cut some corners :)

[FBruzzesi]

@Riik I just merged with main branch and pushed a few changes to parse index and column_names into keys, so that we avoid branching out for dataframe/series in multiple places.

[Riik]

@Riik I just merged with main branch and pushed a few changes to parse index and column_names into keys, so that we avoid branching out for dataframe/series in multiple places.

Thank you! I think the function reads quite nicely now, with the keys internally it is easier to follow the flow.

Anything else that still needs to happen on this PR?

[FBruzzesi]

I am (biased and) happy with the state of this! I will leave the last word to Marco for the set_axis vs s.index=index question before merging!

[MarcoGorelli]

thanks for updating - implementation looks good, just not sure about the argument name df

I did a search on https://github.com/search?q=%22maybe_set_index%28%22+language%3APython&type=code&p=1, and nobody would be affected by such a renaming. I think this is rare enough as a function that we can just change it - the only user I found it scikit-lego, who write

nw.maybe_set_index(result, ["fold", "part"])

(i.e. they pass the first argument positionally, not by name)

---

@FBruzzesi 's suggested obj, as that's what we use in maybe_get_index

[MarcoGorelli]

is it a bit odd for this to be called df when it can be a dataframe or series?

[MarcoGorelli]

thanks @Riik and @FBruzzesi !

[rik-vandervlist-sympower]

🙌 Thanks for the feedback and help, it was a fun PR to work on!