PERF: assert_frame_equal (#44695)

Author: jbrockmendel

pandas/_testing/asserters.py

@@ -22,7 +22,10 @@
     is_numeric_dtype,
     needs_i8_conversion,
 )
-from pandas.core.dtypes.dtypes import PandasDtype
+from pandas.core.dtypes.dtypes import (
+    CategoricalDtype,
+    PandasDtype,
+)
 from pandas.core.dtypes.missing import array_equivalent
 
 import pandas as pd
@@ -655,7 +658,7 @@ def raise_assert_detail(obj, message, left, right, diff=None, index_values=None)
     if isinstance(left, np.ndarray):
         left = pprint_thing(left)
     elif (
-        is_categorical_dtype(left)
+        isinstance(left, CategoricalDtype)
         or isinstance(left, PandasDtype)
         or isinstance(left, StringDtype)
     ):
@@ -664,7 +667,7 @@ def raise_assert_detail(obj, message, left, right, diff=None, index_values=None)
     if isinstance(right, np.ndarray):
         right = pprint_thing(right)
     elif (
-        is_categorical_dtype(right)
+        isinstance(right, CategoricalDtype)
         or isinstance(right, PandasDtype)
         or isinstance(right, StringDtype)
     ):
@@ -1008,8 +1011,8 @@ def assert_series_equal(
         # is False. We'll still raise if only one is a `Categorical`,
         # regardless of `check_categorical`
         if (
-            is_categorical_dtype(left.dtype)
-            and is_categorical_dtype(right.dtype)
+            isinstance(left.dtype, CategoricalDtype)
+            and isinstance(right.dtype, CategoricalDtype)
             and not check_categorical
         ):
             pass
@@ -1054,7 +1057,9 @@ def assert_series_equal(
             raise AssertionError(msg)
     elif is_interval_dtype(left.dtype) and is_interval_dtype(right.dtype):
         assert_interval_array_equal(left.array, right.array)
-    elif is_categorical_dtype(left.dtype) or is_categorical_dtype(right.dtype):
+    elif isinstance(left.dtype, CategoricalDtype) or isinstance(
+        right.dtype, CategoricalDtype
+    ):
         _testing.assert_almost_equal(
             left._values,
             right._values,
@@ -1106,7 +1111,9 @@ def assert_series_equal(
         assert_attr_equal("name", left, right, obj=obj)
 
     if check_categorical:
-        if is_categorical_dtype(left.dtype) or is_categorical_dtype(right.dtype):
+        if isinstance(left.dtype, CategoricalDtype) or isinstance(
+            right.dtype, CategoricalDtype
+        ):
             assert_categorical_equal(
                 left._values,
                 right._values,
@@ -1315,9 +1322,11 @@ def assert_frame_equal(
     # compare by columns
     else:
         for i, col in enumerate(left.columns):
-            assert col in right
-            lcol = left.iloc[:, i]
-            rcol = right.iloc[:, i]
+            # We have already checked that columns match, so we can do
+            #  fast location-based lookups
+            lcol = left._ixs(i, axis=1)
+            rcol = right._ixs(i, axis=1)
+
             # GH #38183
             # use check_index=False, because we do not want to run
             # assert_index_equal for each column,

pandas/core/frame.py

@@ -3911,7 +3911,8 @@ def _box_col_values(self, values: SingleDataManager, loc: int) -> Series:
         #  we attach the Timestamp object as the name.
         name = self.columns[loc]
         klass = self._constructor_sliced
-        return klass(values, index=self.index, name=name, fastpath=True)
+        # We get index=self.index bc values is a SingleDataManager
+        return klass(values, name=name, fastpath=True)
 
     # ----------------------------------------------------------------------
     # Lookup Caching

pandas/core/series.py

@@ -332,7 +332,11 @@ def __init__(
         ):
             # GH#33357 called with just the SingleBlockManager
             NDFrame.__init__(self, data)
-            self.name = name
+            if fastpath:
+                # e.g. from _box_col_values, skip validation of name
+                object.__setattr__(self, "_name", name)
+            else:
+                self.name = name
             return
 
         # we are called internally, so short-circuit

pandas/tests/frame/methods/test_sort_values.py

@@ -3,8 +3,6 @@
 import numpy as np
 import pytest
 
-from pandas.errors import PerformanceWarning
-
 import pandas as pd
 from pandas import (
     Categorical,
@@ -849,13 +847,7 @@ def test_sort_column_level_and_index_label(
         # Compute result by transposing and sorting on axis=1.
         result = df_idx.T.sort_values(by=sort_names, ascending=ascending, axis=1)
 
-        if len(levels) > 1:
-            # Accessing multi-level columns that are not lexsorted raises a
-            # performance warning
-            with tm.assert_produces_warning(PerformanceWarning):
-                tm.assert_frame_equal(result, expected)
-        else:
-            tm.assert_frame_equal(result, expected)
+        tm.assert_frame_equal(result, expected)
 
     def test_sort_values_pos_args_deprecation(self):
         # https://github.com/pandas-dev/pandas/issues/41485

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -9,8 +9,6 @@
 import numpy as np
 import pytest
 
-from pandas.errors import PerformanceWarning
-
 from pandas.core.dtypes.common import is_integer_dtype
 
 import pandas as pd
@@ -373,9 +371,7 @@ def test_agg_multiple_functions_same_name_with_ohlc_present():
     expected = DataFrame(
         expected_values, columns=expected_columns, index=expected_index
     )
-    # PerformanceWarning is thrown by `assert col in right` in assert_frame_equal
-    with tm.assert_produces_warning(PerformanceWarning):
-        tm.assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_multiple_functions_tuples_and_non_tuples(df):