BUG/PERF: use lexsort_indexer in MultiIndex.argsort

doc/source/whatsnew/v1.6.0.rst

@@ -173,6 +173,7 @@ Missing
 
 MultiIndex
 ^^^^^^^^^^
+- Bug in :meth:`MultiIndex.argsort` raising ``TypeError`` when index contains :attr:`NA` (:issue:`48495`)
 - Bug in :meth:`MultiIndex.unique` losing extension array dtype (:issue:`48335`)
 - Bug in :meth:`MultiIndex.union` losing extension array (:issue:`48498`, :issue:`48505`)
 - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`)

pandas/core/indexes/multi.py

@@ -2226,9 +2226,9 @@ def append(self, other):
 
     def argsort(self, *args, **kwargs) -> npt.NDArray[np.intp]:
         if len(args) == 0 and len(kwargs) == 0:
-            # np.lexsort is significantly faster than self._values.argsort()
-            values = [self._get_level_values(i) for i in reversed(range(self.nlevels))]
-            return np.lexsort(values)
+            # lexsort is significantly faster than self._values.argsort()
+            target = self._sort_levels_monotonic()
+            return lexsort_indexer(target._get_codes_for_sorting())
         return self._values.argsort(*args, **kwargs)
 
     @Appender(_index_shared_docs["repeat"] % _index_doc_kwargs)

pandas/tests/indexes/multi/test_setops.py

@@ -379,8 +379,7 @@ def test_union_sort_other_incomparable():
     idx = MultiIndex.from_product([[1, pd.Timestamp("2000")], ["a", "b"]])
 
     # default, sort=None
-    with tm.assert_produces_warning(RuntimeWarning):
-        result = idx.union(idx[:1])
+    result = idx.union(idx[:1])
     tm.assert_index_equal(result, idx)
 
     # sort=False

pandas/tests/indexing/multiindex/test_sorted.py

@@ -2,9 +2,11 @@
 import pytest
 
 from pandas import (
+    NA,
     DataFrame,
     MultiIndex,
     Series,
+    array,
 )
 import pandas._testing as tm
 
@@ -86,6 +88,34 @@ def test_sort_values_key(self):
 
         tm.assert_frame_equal(result, expected)
 
+    def test_argsort_with_na(self):
+        arrays = [
+            array([2, NA, 1], dtype="Int64"),
+            array([1, 2, 3], dtype="Int64"),
+        ]
+        index = MultiIndex.from_arrays(arrays)
+        result = index.argsort()
+        expected = np.array([2, 0, 1])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_sort_values_with_na(self):
+        arrays = [
+            array([2, NA, 1], dtype="Int64"),
+            array([1, 2, 3], dtype="Int64"),
+        ]
+        index = MultiIndex.from_arrays(arrays)
+        index = index.sort_values()
+        result = DataFrame(range(3), index=index)
+
+        arrays = [
+            array([1, 2, NA], dtype="Int64"),
+            array([3, 1, 2], dtype="Int64"),
+        ]
+        index = MultiIndex.from_arrays(arrays)
+        expected = DataFrame(range(3), index=index)
+
+        tm.assert_frame_equal(result, expected)
+
     def test_frame_getitem_not_sorted(self, multiindex_dataframe_random_data):
         frame = multiindex_dataframe_random_data
         df = frame.T