From d0eea008acc3709742ca07c226bc76f19ba50da5 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Tue, 22 Nov 2022 23:03:33 +0000 Subject: [PATCH 1/8] BUG: MultiIndex.putmask losing ea dtype --- doc/source/whatsnew/v2.0.0.rst | 1 + pandas/core/indexes/base.py | 1 - pandas/core/indexes/multi.py | 29 +++++++++++++++++++++ pandas/tests/indexes/multi/test_indexing.py | 28 ++++++++++++++++++++ 4 files changed, 58 insertions(+), 1 deletion(-) diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst index 02ea290995c8d..9814d684c7f77 100644 --- a/doc/source/whatsnew/v2.0.0.rst +++ b/doc/source/whatsnew/v2.0.0.rst @@ -700,6 +700,7 @@ MultiIndex - Bug in :meth:`MultiIndex.union` not sorting when sort=None and index contains missing values (:issue:`49010`) - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`) - Bug in :meth:`MultiIndex.symmetric_difference` losing extension array (:issue:`48607`) +- Bug in :meth:`MultiIndex.putmask` losing extension array (:issue:`49830`) - Bug in :meth:`MultiIndex.value_counts` returning a :class:`Series` indexed by flat index of tuples instead of a :class:`MultiIndex` (:issue:`49558`) - diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 0bc568fb122ed..9c1f5fea7603d 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -5159,7 +5159,6 @@ def _concat(self, to_concat: list[Index], name: Hashable) -> Index: return Index._with_infer(result, name=name) - @final def putmask(self, mask, value) -> Index: """ Return a new Index of the values set with the mask. diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 04a57c1709382..115b4c31008b6 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -79,6 +79,7 @@ ) import pandas.core.algorithms as algos +from pandas.core.array_algos.putmask import validate_putmask from pandas.core.arrays import Categorical from pandas.core.arrays.categorical import factorize_from_iterables import pandas.core.common as com @@ -3659,6 +3660,34 @@ def _validate_fill_value(self, item): raise ValueError("Item must have length equal to number of levels.") return item + def putmask(self, mask, value: MultiIndex) -> Index: + mask, noop = validate_putmask(self, mask) + if noop: + return self.copy() + + if len(mask) == len(value): + subset = value[mask].remove_unused_levels() + else: + subset = value.remove_unused_levels() + + new_levels = [] + new_codes = [] + + for i, (value_level, level, level_codes) in enumerate( + zip(subset.levels, self.levels, self.codes) + ): + new_elements = value_level.difference(level) + new_level = level.append(new_elements) + value_codes = new_level.get_indexer_for(subset.get_level_values(i)) + new_code = ensure_int64(level_codes) + new_code[mask] = value_codes + new_levels.append(new_level) + new_codes.append(new_code) + + return MultiIndex( + levels=new_levels, codes=new_codes, names=self.names, verify_integrity=False + ) + def insert(self, loc: int, item) -> MultiIndex: """ Make new MultiIndex inserting new item at location diff --git a/pandas/tests/indexes/multi/test_indexing.py b/pandas/tests/indexes/multi/test_indexing.py index 552b3753083fe..4c879c8ff5736 100644 --- a/pandas/tests/indexes/multi/test_indexing.py +++ b/pandas/tests/indexes/multi/test_indexing.py @@ -162,6 +162,34 @@ def test_putmask_multiindex_other(self): expected = MultiIndex.from_tuples([right[0], right[1], left[2]]) tm.assert_index_equal(result, expected) + def test_putmask_keep_dtype(self, any_numeric_ea_dtype): + # GH#49830 + midx = MultiIndex.from_arrays( + [pd.Series([1, 2, 3], dtype=any_numeric_ea_dtype), [10, 11, 12]] + ) + midx2 = MultiIndex.from_arrays( + [pd.Series([5, 6, 7], dtype=any_numeric_ea_dtype), [-1, -2, -3]] + ) + result = midx.putmask([True, False, False], midx2) + expected = MultiIndex.from_arrays( + [pd.Series([5, 2, 3], dtype=any_numeric_ea_dtype), [-1, 11, 12]] + ) + tm.assert_index_equal(result, expected) + + def test_putmask_keep_dtype_shorter_value(self, any_numeric_ea_dtype): + # GH#49830 + midx = MultiIndex.from_arrays( + [pd.Series([1, 2, 3], dtype=any_numeric_ea_dtype), [10, 11, 12]] + ) + midx2 = MultiIndex.from_arrays( + [pd.Series([5], dtype=any_numeric_ea_dtype), [-1]] + ) + result = midx.putmask([True, False, False], midx2) + expected = MultiIndex.from_arrays( + [pd.Series([5, 2, 3], dtype=any_numeric_ea_dtype), [-1, 11, 12]] + ) + tm.assert_index_equal(result, expected) + class TestGetIndexer: def test_get_indexer(self): From b73beee00381de87ace7bd9ce5d5df3113881d16 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Wed, 23 Nov 2022 00:33:50 +0000 Subject: [PATCH 2/8] Fix typing --- pandas/core/array_algos/putmask.py | 10 ++++++++-- pandas/core/indexes/multi.py | 15 ++++++++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/pandas/core/array_algos/putmask.py b/pandas/core/array_algos/putmask.py index 17622e78d1b12..3e2c711d12f26 100644 --- a/pandas/core/array_algos/putmask.py +++ b/pandas/core/array_algos/putmask.py @@ -3,7 +3,10 @@ """ from __future__ import annotations -from typing import Any +from typing import ( + TYPE_CHECKING, + Any, +) import numpy as np @@ -19,6 +22,9 @@ from pandas.core.arrays import ExtensionArray +if TYPE_CHECKING: + from pandas import MultiIndex + def putmask_inplace(values: ArrayLike, mask: npt.NDArray[np.bool_], value: Any) -> None: """ @@ -96,7 +102,7 @@ def putmask_without_repeat( def validate_putmask( - values: ArrayLike, mask: np.ndarray + values: ArrayLike | MultiIndex, mask: np.ndarray ) -> tuple[npt.NDArray[np.bool_], bool]: """ Validate mask and check if this putmask operation is a no-op. diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 115b4c31008b6..3841fe08a444c 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -3660,7 +3660,20 @@ def _validate_fill_value(self, item): raise ValueError("Item must have length equal to number of levels.") return item - def putmask(self, mask, value: MultiIndex) -> Index: + def putmask(self, mask, value: MultiIndex) -> MultiIndex: + """ + Return a new MultiIndex of the values set with the mask. + + Parameters + ---------- + mask : array like + value : MultiIndex + Must either be the same length as self or length one + + Returns + ------- + MultiIndex + """ mask, noop = validate_putmask(self, mask) if noop: return self.copy() From 55f6ad266422c9b0bab310d607d3fd3946f8faac Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Wed, 23 Nov 2022 15:20:23 +0000 Subject: [PATCH 3/8] Add asv --- asv_bench/benchmarks/multiindex_object.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/asv_bench/benchmarks/multiindex_object.py b/asv_bench/benchmarks/multiindex_object.py index 97c710be6d5a1..de55268e0407b 100644 --- a/asv_bench/benchmarks/multiindex_object.py +++ b/asv_bench/benchmarks/multiindex_object.py @@ -379,4 +379,26 @@ def time_isin_large(self, dtype): self.midx.isin(self.values_large) +class Putmask: + def setup(self): + N = 10**5 + level1 = range(1_000) + + level2 = date_range(start="1/1/2000", periods=N // 1000) + self.midx = MultiIndex.from_product([level1, level2]) + + level1 = range(1_000, 2_000) + self.midx_values = MultiIndex.from_product([level1, level2]) + + level2 = date_range(start="1/1/2010", periods=N // 1000) + self.midx_values_different = MultiIndex.from_product([level1, level2]) + self.mask = np.array([True, False] * (N // 2)) + + def time_putmask(self): + self.midx.putmask(self.mask, self.midx_values) + + def time_putmask_all_different(self): + self.midx.putmask(self.mask, self.midx_values_different) + + from .pandas_vb_common import setup # noqa: F401 isort:skip From 6bf213d682cbb7cd7a44f55cdcd0776d6fdd2974 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Wed, 23 Nov 2022 22:06:50 +0000 Subject: [PATCH 4/8] Simplify and add whatsnew --- doc/source/whatsnew/v2.0.0.rst | 1 + pandas/core/indexes/multi.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst index 9814d684c7f77..14cd13fb362d6 100644 --- a/doc/source/whatsnew/v2.0.0.rst +++ b/doc/source/whatsnew/v2.0.0.rst @@ -586,6 +586,7 @@ Performance improvements - Performance improvement in :class:`MultiIndex` set operations with sort=None (:issue:`49010`) - Performance improvement in :meth:`.DataFrameGroupBy.mean`, :meth:`.SeriesGroupBy.mean`, :meth:`.DataFrameGroupBy.var`, and :meth:`.SeriesGroupBy.var` for extension array dtypes (:issue:`37493`) - Performance improvement in :meth:`MultiIndex.isin` when ``level=None`` (:issue:`48622`, :issue:`49577`) +- Performance improvement in :meth:`MultiIndex.putmask` (:issue:`49830`) - Performance improvement in :meth:`Index.union` and :meth:`MultiIndex.union` when index contains duplicates (:issue:`48900`) - Performance improvement in :meth:`Series.fillna` for pyarrow-backed dtypes (:issue:`49722`) - Performance improvement for :meth:`Series.value_counts` with nullable dtype (:issue:`48338`) diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 3841fe08a444c..eba1993afa26a 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -3689,8 +3689,7 @@ def putmask(self, mask, value: MultiIndex) -> MultiIndex: for i, (value_level, level, level_codes) in enumerate( zip(subset.levels, self.levels, self.codes) ): - new_elements = value_level.difference(level) - new_level = level.append(new_elements) + new_level = level.union(value_level, sort=False) value_codes = new_level.get_indexer_for(subset.get_level_values(i)) new_code = ensure_int64(level_codes) new_code[mask] = value_codes From f6d0a799f418eb84a7d2817ba541b96b9852bd76 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Wed, 23 Nov 2022 23:11:39 +0000 Subject: [PATCH 5/8] BUG: MultiIndex.join losing dtype --- doc/source/whatsnew/v2.0.0.rst | 1 + pandas/core/indexes/base.py | 19 +++------------ pandas/tests/indexes/multi/test_join.py | 31 +++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 16 deletions(-) diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst index 14cd13fb362d6..03c8f0638acba 100644 --- a/doc/source/whatsnew/v2.0.0.rst +++ b/doc/source/whatsnew/v2.0.0.rst @@ -701,6 +701,7 @@ MultiIndex - Bug in :meth:`MultiIndex.union` not sorting when sort=None and index contains missing values (:issue:`49010`) - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`) - Bug in :meth:`MultiIndex.symmetric_difference` losing extension array (:issue:`48607`) +- Bug in :meth:`MultiIndex.join` losing dtypes when :class:`MultiIndex` has duplicates and is non-monotonic (:issue:`49830`) - Bug in :meth:`MultiIndex.putmask` losing extension array (:issue:`49830`) - Bug in :meth:`MultiIndex.value_counts` returning a :class:`Series` indexed by flat index of tuples instead of a :class:`MultiIndex` (:issue:`49558`) - diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 9c1f5fea7603d..54d7284abaa42 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4587,22 +4587,9 @@ def _join_non_unique( ) mask = left_idx == -1 - join_array = self._values.take(left_idx) - right = other._values.take(right_idx) - - if isinstance(join_array, np.ndarray): - # error: Argument 3 to "putmask" has incompatible type - # "Union[ExtensionArray, ndarray[Any, Any]]"; expected - # "Union[_SupportsArray[dtype[Any]], _NestedSequence[ - # _SupportsArray[dtype[Any]]], bool, int, float, complex, - # str, bytes, _NestedSequence[Union[bool, int, float, - # complex, str, bytes]]]" - np.putmask(join_array, mask, right) # type: ignore[arg-type] - else: - join_array._putmask(mask, right) - - join_index = self._wrap_joined_index(join_array, other) - + join_idx = self.take(left_idx) + right = other.take(right_idx) + join_index = join_idx.putmask(mask, right) return join_index, left_idx, right_idx @final diff --git a/pandas/tests/indexes/multi/test_join.py b/pandas/tests/indexes/multi/test_join.py index aa2f2ca5af7bd..3a1cb846868da 100644 --- a/pandas/tests/indexes/multi/test_join.py +++ b/pandas/tests/indexes/multi/test_join.py @@ -225,3 +225,34 @@ def test_join_multi_with_nan(): index=MultiIndex.from_product([["A"], [1.0, 2.0]], names=["id1", "id2"]), ) tm.assert_frame_equal(result, expected) + + +def test_join_dtypes(any_numeric_ea_dtype): + # GH#49830 + midx = MultiIndex.from_arrays([Series([1, 2], dtype=any_numeric_ea_dtype), [3, 4]]) + midx2 = MultiIndex.from_arrays( + [Series([1, 0, 0], dtype=any_numeric_ea_dtype), [3, 4, 4]] + ) + result = midx.join(midx2, how="outer") + expected = MultiIndex.from_arrays( + [Series([0, 0, 1, 2], dtype=any_numeric_ea_dtype), [4, 4, 3, 4]] + ) + tm.assert_index_equal(result, expected) + + +def test_join_dtypes_all_nan(any_numeric_ea_dtype): + # GH#49830 + midx = MultiIndex.from_arrays( + [Series([1, 2], dtype=any_numeric_ea_dtype), [np.nan, np.nan]] + ) + midx2 = MultiIndex.from_arrays( + [Series([1, 0, 0], dtype=any_numeric_ea_dtype), [np.nan, np.nan, np.nan]] + ) + result = midx.join(midx2, how="outer") + expected = MultiIndex.from_arrays( + [ + Series([0, 0, 1, 2], dtype=any_numeric_ea_dtype), + [np.nan, np.nan, np.nan, np.nan], + ] + ) + tm.assert_index_equal(result, expected) From d5259733c9020c941a4b8ed34ac1e509cfdb9177 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Wed, 23 Nov 2022 23:32:19 +0000 Subject: [PATCH 6/8] BUG: MultiIndex.join losing dtype --- doc/source/whatsnew/v2.0.0.rst | 2 +- pandas/core/indexes/base.py | 16 +++++++++++----- pandas/core/indexes/datetimelike.py | 4 ++-- pandas/tests/indexes/multi/test_join.py | 9 +++++---- 4 files changed, 19 insertions(+), 12 deletions(-) diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst index 03c8f0638acba..fa71af405a077 100644 --- a/doc/source/whatsnew/v2.0.0.rst +++ b/doc/source/whatsnew/v2.0.0.rst @@ -701,7 +701,7 @@ MultiIndex - Bug in :meth:`MultiIndex.union` not sorting when sort=None and index contains missing values (:issue:`49010`) - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`) - Bug in :meth:`MultiIndex.symmetric_difference` losing extension array (:issue:`48607`) -- Bug in :meth:`MultiIndex.join` losing dtypes when :class:`MultiIndex` has duplicates and is non-monotonic (:issue:`49830`) +- Bug in :meth:`MultiIndex.join` losing dtypes when :class:`MultiIndex` has duplicates (:issue:`49830`) - Bug in :meth:`MultiIndex.putmask` losing extension array (:issue:`49830`) - Bug in :meth:`MultiIndex.value_counts` returning a :class:`Series` indexed by flat index of tuples instead of a :class:`MultiIndex` (:issue:`49558`) - diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 54d7284abaa42..8f59a12c5e128 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4766,10 +4766,10 @@ def _join_monotonic( ridx = None elif how == "inner": join_array, lidx, ridx = self._inner_indexer(other) - join_index = self._wrap_joined_index(join_array, other) + join_index = self._wrap_joined_index(join_array, other, lidx, ridx) elif how == "outer": join_array, lidx, ridx = self._outer_indexer(other) - join_index = self._wrap_joined_index(join_array, other) + join_index = self._wrap_joined_index(join_array, other, lidx, ridx) else: if how == "left": join_array, lidx, ridx = self._left_indexer(other) @@ -4780,20 +4780,26 @@ def _join_monotonic( elif how == "outer": join_array, lidx, ridx = self._outer_indexer(other) - join_index = self._wrap_joined_index(join_array, other) + join_index = self._wrap_joined_index(join_array, other, lidx, ridx) lidx = None if lidx is None else ensure_platform_int(lidx) ridx = None if ridx is None else ensure_platform_int(ridx) return join_index, lidx, ridx - def _wrap_joined_index(self: _IndexT, joined: ArrayLike, other: _IndexT) -> _IndexT: + def _wrap_joined_index( + self: _IndexT, joined: ArrayLike, other: _IndexT, lidx, ridx + ) -> _IndexT: assert other.dtype == self.dtype if isinstance(self, ABCMultiIndex): name = self.names if self.names == other.names else None # error: Incompatible return value type (got "MultiIndex", # expected "_IndexT") - return self._constructor(joined, name=name) # type: ignore[return-value] + mask = lidx == -1 + join_idx = self.take(lidx) + right = other.take(ridx) + join_index = join_idx.putmask(mask, right) + return join_index.set_names(name) # type: ignore[return-value] else: name = get_op_result_name(self, other) return self._constructor._with_infer(joined, name=name, dtype=self.dtype) diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index c6c8695ab01da..20db210f1fae9 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -602,9 +602,9 @@ def _get_join_freq(self, other): freq = self.freq return freq - def _wrap_joined_index(self, joined, other): + def _wrap_joined_index(self, joined, other, lidx, ridx): assert other.dtype == self.dtype, (other.dtype, self.dtype) - result = super()._wrap_joined_index(joined, other) + result = super()._wrap_joined_index(joined, other, lidx, ridx) result._data._freq = self._get_join_freq(other) return result diff --git a/pandas/tests/indexes/multi/test_join.py b/pandas/tests/indexes/multi/test_join.py index 3a1cb846868da..4fff862961920 100644 --- a/pandas/tests/indexes/multi/test_join.py +++ b/pandas/tests/indexes/multi/test_join.py @@ -227,16 +227,17 @@ def test_join_multi_with_nan(): tm.assert_frame_equal(result, expected) -def test_join_dtypes(any_numeric_ea_dtype): +@pytest.mark.parametrize("val", [0, 5]) +def test_join_dtypes(any_numeric_ea_dtype, val): # GH#49830 midx = MultiIndex.from_arrays([Series([1, 2], dtype=any_numeric_ea_dtype), [3, 4]]) midx2 = MultiIndex.from_arrays( - [Series([1, 0, 0], dtype=any_numeric_ea_dtype), [3, 4, 4]] + [Series([1, val, val], dtype=any_numeric_ea_dtype), [3, 4, 4]] ) result = midx.join(midx2, how="outer") expected = MultiIndex.from_arrays( - [Series([0, 0, 1, 2], dtype=any_numeric_ea_dtype), [4, 4, 3, 4]] - ) + [Series([val, val, 1, 2], dtype=any_numeric_ea_dtype), [4, 4, 3, 4]] + ).sort_values() tm.assert_index_equal(result, expected) From 869283b04f751548f631f75134703930227848ae Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Tue, 29 Nov 2022 00:59:23 +0100 Subject: [PATCH 7/8] Add test --- pandas/tests/frame/methods/test_combine_first.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pandas/tests/frame/methods/test_combine_first.py b/pandas/tests/frame/methods/test_combine_first.py index e838c8fabf456..ad6122501bc19 100644 --- a/pandas/tests/frame/methods/test_combine_first.py +++ b/pandas/tests/frame/methods/test_combine_first.py @@ -543,3 +543,17 @@ def test_combine_first_int64_not_cast_to_float64(): result = df_1.combine_first(df_2) expected = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [12, 34, 65]}) tm.assert_frame_equal(result, expected) + + +def test_midx_losing_dtype(): + # GH#49830 + midx = MultiIndex.from_arrays([[0, 0], [np.nan, np.nan]]) + midx2 = MultiIndex.from_arrays([[1, 1], [np.nan, np.nan]]) + df1 = DataFrame({"a": [None, 4]}, index=midx) + df2 = DataFrame({"a": [3, 3]}, index=midx2) + result = df1.combine_first(df2) + expected_midx = MultiIndex.from_arrays( + [[0, 0, 1, 1], [np.nan, np.nan, np.nan, np.nan]] + ) + expected = DataFrame({"a": [np.nan, 4, 3, 3]}, index=expected_midx) + tm.assert_frame_equal(result, expected) From 1871916d8fb7dcade0bb95386d0be3bb1f62cb88 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Tue, 29 Nov 2022 21:34:20 +0100 Subject: [PATCH 8/8] Add types --- pandas/core/indexes/base.py | 13 ++++++++++--- pandas/core/indexes/datetimelike.py | 4 +++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 86a452625dc3f..29cda0fc72d0b 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -4736,8 +4736,8 @@ def _join_monotonic( ret_index = other if how == "right" else self return ret_index, None, None - ridx: np.ndarray | None - lidx: np.ndarray | None + ridx: npt.NDArray[np.intp] | None + lidx: npt.NDArray[np.intp] | None if self.is_unique and other.is_unique: # We can perform much better than the general case @@ -4765,6 +4765,9 @@ def _join_monotonic( elif how == "outer": join_array, lidx, ridx = self._outer_indexer(other) + assert lidx is not None + assert ridx is not None + join_index = self._wrap_joined_index(join_array, other, lidx, ridx) lidx = None if lidx is None else ensure_platform_int(lidx) @@ -4772,7 +4775,11 @@ def _join_monotonic( return join_index, lidx, ridx def _wrap_joined_index( - self: _IndexT, joined: ArrayLike, other: _IndexT, lidx, ridx + self: _IndexT, + joined: ArrayLike, + other: _IndexT, + lidx: npt.NDArray[np.intp], + ridx: npt.NDArray[np.intp], ) -> _IndexT: assert other.dtype == self.dtype diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index bc73fbb64f02b..afe84befb6b31 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -625,7 +625,9 @@ def _get_join_freq(self, other): freq = self.freq return freq - def _wrap_joined_index(self, joined, other, lidx, ridx): + def _wrap_joined_index( + self, joined, other, lidx: npt.NDArray[np.intp], ridx: npt.NDArray[np.intp] + ): assert other.dtype == self.dtype, (other.dtype, self.dtype) result = super()._wrap_joined_index(joined, other, lidx, ridx) result._data._freq = self._get_join_freq(other)