Skip to content

BUG: Fix getitem dtype preservation with multiindexes #51895

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
merged 14 commits into from
Apr 25, 2023
Merged

BUG: Fix getitem dtype preservation with multiindexes #51895

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
a459276
BUG/TST fix dtype preservation with multindex
m-richards Mar 11, 2023
eab4c6a
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
m-richards Mar 11, 2023
a3fc02f
lint
m-richards Mar 11, 2023
1be4d25
Update pandas/tests/indexing/multiindex/test_multiindex.py
m-richards Mar 11, 2023
e865548
cleanups
m-richards Mar 11, 2023
fda3df8
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
m-richards Mar 11, 2023
533e0d8
switch to iloc, reindex fails in some cases
m-richards Mar 12, 2023
1d221b1
suggestions from code review
m-richards Mar 14, 2023
968228a
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
m-richards Mar 14, 2023
f56e76c
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
jorisvandenbossche Mar 16, 2023
e3c17af
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
m-richards Apr 23, 2023
6f260fe
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
m-richards Apr 23, 2023
079d543
Merge remote-tracking branch 'upstream/main' into fix_multiindex_dtype
m-richards Apr 24, 2023
6badc9e
address code review comments
m-richards Apr 25, 2023
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
11 changes: 2 additions & 9 deletions pandas/core/frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3827,15 +3827,8 @@ def _getitem_multilevel(self, key):
if isinstance(loc, (slice, np.ndarray)):
new_columns = self.columns[loc]
result_columns = maybe_droplevels(new_columns, key)
if self._is_mixed_type:
result = self.reindex(columns=new_columns)
result.columns = result_columns
else:
new_values = self.values[:, loc]
Copy link
Member

Choose a reason for hiding this comment

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

I guess this is faster than iloc? Did not profile but that's the only reason I can think of that it is done this way. To preserve the performance, I guess we can check something like _can_fast_transpose here as well?

Copy link
Member

Choose a reason for hiding this comment

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

But in the case where this matters, I expect that locis a slice? And then I would expect iloc to be fast as well? (it has separate paths for single blocks AFAIR)
But yes, something to verify (indexing the numpy array will always be faster since you don't have the overhead of our indexing machinery. But if we also end up doing the same, the difference should be acceptable).

Copy link
Member

Choose a reason for hiding this comment

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

(and I think a bit of extra overhead is also acceptable to avoid needing the custom add_references code here (as in #51944)

Copy link
Member

Choose a reason for hiding this comment

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

Sounds good as long as is not to much overhead. We should check at least. That was my reasoning for adding the reference code instead of switching to iloc

Copy link
Member

Choose a reason for hiding this comment

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

Small test:

columns = pd.MultiIndex.from_product([range(10), range(20)])
df = pd.DataFrame(np.random.randn(10000, len(columns)), columns=columns).copy()
slice = df.columns.get_loc(0)
# gives slice(0, 20, None)

In [47]: %timeit df.iloc[:, slice]
96.2 µs ± 4.58 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [48]: %timeit pd.DataFrame(df.values[:, slice], index=df.index)
33.9 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

So iloc is a bit slower, having a bit more overhead (mostly in creating the resulting MultiIndex columns, which we then override ..). But this mostly fixed overhead regardless of the size of the data, and we are speaking about microseconds.
The actual subsetting still happens with a slice. And just to be sure to compare to a non-slice case (where we do an actual "take") selecting the same subset:

In [49]: idx = np.arange(20)

In [50]: %timeit df.iloc[:, idx]
576 µs ± 19.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

result = self._constructor(
new_values, index=self.index, columns=result_columns
)
result = result.__finalize__(self)
result = self.iloc[:, loc]
result.columns = result_columns

# If there is only one column being returned, and its name is
# either an empty string, or a tuple with an empty string as its
Expand Down
17 changes: 17 additions & 0 deletions pandas/tests/indexing/multiindex/test_multiindex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,3 +206,20 @@ def test_multiindex_with_na_missing_key(self):
)
with pytest.raises(KeyError, match="missing_key"):
df[[("missing_key",)]]

def test_multiindex_dtype_preservation(self):
# GH51261
columns = MultiIndex.from_tuples([("A", "B")], names=["lvl1", "lvl2"])
df = DataFrame(["value"], columns=columns).astype("category")
assert (df["A"].dtypes == "category").all()
Copy link
Member

Choose a reason for hiding this comment

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

Could you use pandas.core.dtypes.common.is_categorical_dtype(df["A"]) here?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

sure


# geopandas 1763 analogue
df = DataFrame(
[[1, 0], [0, 1]],
columns=[
["foo", "foo"],
["location", "location"],
["x", "y"],
],
).assign(bools=Series([True, False], dtype="boolean"))
assert df["bools"].dtype == "boolean"
Copy link
Member

Choose a reason for hiding this comment

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

Could you use pandas.core.dtypes.common.is_bool_dtype(df["bools"]) here?

Copy link
Member

Choose a reason for hiding this comment

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

For this case we actually want to be more explicit to check the extension dtype (is_bool_dtype would also pass with numpy bool dtype). Possible alternative is isinstance(.., BooleanDtype)