BUG/PERF: algos.union_with_duplicates losing EA dtypes #48900
Conversation
lukemanley
added
Performance
pandas/core/algorithms.py
Outdated
| for i, value in enumerate(unique_array): | ||
| indexer += [i] * int(max(l_count.at[value], r_count.at[value])) | ||
| return unique_array.take(indexer) | ||
| repeats = final_count.reindex(unique_array).values # type: ignore[attr-defined] |
There was a problem hiding this comment.
Looks good,
can you refactor this to return the indexer and operate on the MultiIndex in _union()? This avoids losing the dtype.
There was a problem hiding this comment.
Refactored to preserve EA dtypes. The indexer applies to the combined unique values not the original MI's so I don't think we need to return it.
lukemanley
changed the title
pandas/core/algorithms.py
Outdated
| for i, value in enumerate(unique_array): | ||
| indexer += [i] * int(max(l_count.at[value], r_count.at[value])) | ||
| return unique_array.take(indexer) | ||
| final_count = np.maximum(l_count, r_count).astype("int", copy=False) |
There was a problem hiding this comment.
You've got 2 MultiIndex checks now. I'd rather handle this on the _union level. Is there any performance benefit of doing it like this?
Additionally, I think you can use unique_vals.take(repeats)
That's more in line with what we are doing elsewhere
There was a problem hiding this comment.
I think we need the MultiIndex checks within union_with_duplicates, otherwise we lose dtypes in there. I'm not sure how how this could be done only on the _union level. Let me know if you have any suggestions.
The repeat logic is not quite a simple take:
np.repeat(unique_vals, repeats)
is equivalent to:
indexer = np.arange(len(unique_vals))
indexer = np.repeat(indexer, repeats)
unique_vals.take(indexer)
There was a problem hiding this comment.
You could simply return final_count and handle the rest In the _union level. Alternatively, you could just pass in self and other and handle everything on the lower level. Or add another argument unique_vals to the function
There was a problem hiding this comment.
Thx, makes sense. I pushed the type checking down into union_with_duplicates.
pandas/core/algorithms.py
Outdated
| else: | ||
| unique_vals = unique(concat_compat([lvals, rvals])) | ||
| unique_vals = ensure_wrapped_if_datetimelike(unique_vals) | ||
| repeats = final_count.reindex(unique_vals).values # type: ignore[attr-defined] |
There was a problem hiding this comment.
Could you add the mypy error as a comment?
There was a problem hiding this comment.
I fixed/removed the mypy ignore
| @@ -578,6 +578,30 @@ def test_union_keep_ea_dtype(any_numeric_ea_dtype, val): | |||
| tm.assert_index_equal(result, expected) | |||
|
|
|||
|
|
|||
| def test_union_with_duplicates_keep_ea_dtype(any_numeric_ea_dtype): | |||
| # GH48900 | |||
| mi1 = MultiIndex.from_arrays( | |||
There was a problem hiding this comment.
What happens if theres duplicate NAs?
There was a problem hiding this comment.
Works the same as non-NA. I added the dupe NA case to the test.
|
thx @lukemanley |
doc/source/whatsnew/v1.6.0.rstfile if fixing a bug or adding a new feature.ASV added: