Merge branch 'master' of https://github.com/pandas-dev/pandas into cln-arith

Author: jbrockmendel

doc/source/whatsnew/v1.2.0.rst

@@ -285,6 +285,8 @@ Performance improvements
 - ``Styler`` uuid method altered to compress data transmission over web whilst maintaining reasonably low table collision probability (:issue:`36345`)
 - Performance improvement in :meth:`pd.to_datetime` with non-ns time unit for ``float`` ``dtype`` columns (:issue:`20445`)
 - Performance improvement in setting values on a :class:`IntervalArray` (:issue:`36310`)
+- The internal index method :meth:`~Index._shallow_copy` now makes the new index and original index share cached attributes,
+  avoiding creating these again, if created on either. This can speed up operations that depend on creating copies of existing indexes (:issue:`36840`)
 - Performance improvement in :meth:`RollingGroupby.count` (:issue:`35625`)
 
 .. ---------------------------------------------------------------------------
@@ -435,7 +437,7 @@ ExtensionArray
 
 - Fixed Bug where :class:`DataFrame` column set to scalar extension type via a dict instantion was considered an object type rather than the extension type (:issue:`35965`)
 - Fixed bug where ``astype()`` with equal dtype and ``copy=False`` would return a new object (:issue:`284881`)
--
+- Fixed bug when applying a NumPy ufunc with multiple outputs to a :class:`pandas.arrays.IntegerArray` returning None (:issue:`36913`)
 
 
 Other

pandas/core/arrays/integer.py

@@ -412,7 +412,7 @@ def reconstruct(x):
 
         result = getattr(ufunc, method)(*inputs2, **kwargs)
         if isinstance(result, tuple):
-            tuple(reconstruct(x) for x in result)
+            return tuple(reconstruct(x) for x in result)
         else:
             return reconstruct(result)
 

pandas/core/computation/expressions.py

@@ -248,6 +248,7 @@ def where(cond, a, b, use_numexpr=True):
     use_numexpr : bool, default True
         Whether to try to use numexpr.
     """
+    assert _where is not None
     return _where(cond, a, b) if use_numexpr else _where_standard(cond, a, b)
 
 

pandas/core/computation/pytables.py

@@ -42,7 +42,10 @@ class Term(ops.Term):
     env: PyTablesScope
 
     def __new__(cls, name, env, side=None, encoding=None):
-        klass = Constant if not isinstance(name, str) else cls
+        if isinstance(name, str):
+            klass = cls
+        else:
+            klass = Constant
         return object.__new__(klass)
 
     def __init__(self, name, env: PyTablesScope, side=None, encoding=None):
@@ -83,6 +86,7 @@ class BinOp(ops.BinOp):
 
     op: str
     queryables: Dict[str, Any]
+    condition: Optional[str]
 
     def __init__(self, op: str, lhs, rhs, queryables: Dict[str, Any], encoding):
         super().__init__(op, lhs, rhs)
@@ -184,10 +188,8 @@ def convert_value(self, v) -> "TermValue":
 
         def stringify(value):
             if self.encoding is not None:
-                encoder = partial(pprint_thing_encoded, encoding=self.encoding)
-            else:
-                encoder = pprint_thing
-            return encoder(value)
+                return pprint_thing_encoded(value, encoding=self.encoding)
+            return pprint_thing(value)
 
         kind = ensure_decoded(self.kind)
         meta = ensure_decoded(self.meta)
@@ -257,9 +259,11 @@ def __repr__(self) -> str:
     def invert(self):
         """ invert the filter """
         if self.filter is not None:
-            f = list(self.filter)
-            f[1] = self.generate_filter_op(invert=True)
-            self.filter = tuple(f)
+            self.filter = (
+                self.filter[0],
+                self.generate_filter_op(invert=True),
+                self.filter[2],
+            )
         return self
 
     def format(self):

pandas/core/indexes/base.py

@@ -561,12 +561,12 @@ def _shallow_copy(self, values=None, name: Label = no_default):
         name : Label, defaults to self.name
         """
         name = self.name if name is no_default else name
-        cache = self._cache.copy() if values is None else {}
-        if values is None:
-            values = self._values
 
-        result = self._simple_new(values, name=name)
-        result._cache = cache
+        if values is not None:
+            return self._simple_new(values, name=name)
+
+        result = self._simple_new(self._values, name=name)
+        result._cache = self._cache
         return result
 
     def is_(self, other) -> bool:

pandas/core/indexes/datetimelike.py

@@ -671,17 +671,15 @@ def _with_freq(self, freq):
 
     def _shallow_copy(self, values=None, name: Label = lib.no_default):
         name = self.name if name is lib.no_default else name
-        cache = self._cache.copy() if values is None else {}
 
-        if values is None:
-            values = self._data
-
-        if isinstance(values, np.ndarray):
+        if values is not None:
             # TODO: We would rather not get here
-            values = type(self._data)(values, dtype=self.dtype)
+            if isinstance(values, np.ndarray):
+                values = type(self._data)(values, dtype=self.dtype)
+            return self._simple_new(values, name=name)
 
-        result = type(self)._simple_new(values, name=name)
-        result._cache = cache
+        result = self._simple_new(self._data, name=name)
+        result._cache = self._cache
         return result
 
     # --------------------------------------------------------------------

pandas/core/indexes/interval.py

@@ -335,12 +335,12 @@ def _shallow_copy(
         self, values: Optional[IntervalArray] = None, name: Label = lib.no_default
     ):
         name = self.name if name is lib.no_default else name
-        cache = self._cache.copy() if values is None else {}
-        if values is None:
-            values = self._data
 
-        result = self._simple_new(values, name=name)
-        result._cache = cache
+        if values is not None:
+            return self._simple_new(values, name=name)
+
+        result = self._simple_new(self._data, name=name)
+        result._cache = self._cache
         return result
 
     @cache_readonly

pandas/core/indexes/period.py

@@ -260,12 +260,12 @@ def _has_complex_internals(self) -> bool:
 
     def _shallow_copy(self, values=None, name: Label = no_default):
         name = name if name is not no_default else self.name
-        cache = self._cache.copy() if values is None else {}
-        if values is None:
-            values = self._data
 
-        result = self._simple_new(values, name=name)
-        result._cache = cache
+        if values is not None:
+            return self._simple_new(values, name=name)
+
+        result = self._simple_new(self._data, name=name)
+        result._cache = self._cache
         return result
 
     def _maybe_convert_timedelta(self, other):

pandas/core/indexes/range.py

@@ -397,13 +397,13 @@ def __iter__(self):
     def _shallow_copy(self, values=None, name: Label = no_default):
         name = self.name if name is no_default else name
 
-        if values is None:
-            result = self._simple_new(self._range, name=name)
-            result._cache = self._cache.copy()
-            return result
-        else:
+        if values is not None:
             return Int64Index._simple_new(values, name=name)
 
+        result = self._simple_new(self._range, name=name)
+        result._cache = self._cache
+        return result
+
     @doc(Int64Index.copy)
     def copy(self, name=None, deep=False, dtype=None, names=None):
         name = self._validate_names(name=name, names=names, deep=deep)[0]

pandas/core/window/rolling.py

@@ -398,7 +398,7 @@ def _insert_on_column(self, result: "DataFrame", obj: "DataFrame"):
 
         if self.on is not None and not self._on.equals(obj.index):
             name = self._on.name
-            extra_col = Series(self._on, index=obj.index, name=name)
+            extra_col = Series(self._on, index=self.obj.index, name=name)
             if name in result.columns:
                 # TODO: sure we want to overwrite results?
                 result[name] = extra_col
@@ -2263,7 +2263,7 @@ def _get_window_indexer(self, window: int) -> GroupbyRollingIndexer:
         """
         rolling_indexer: Type[BaseIndexer]
         indexer_kwargs: Optional[Dict] = None
-        index_array = self.obj.index.asi8
+        index_array = self._on.asi8
         if isinstance(self.window, BaseIndexer):
             rolling_indexer = type(self.window)
             indexer_kwargs = self.window.__dict__

pandas/tests/arrays/integer/test_function.py

@@ -64,6 +64,20 @@ def test_ufuncs_binary_int(ufunc):
     tm.assert_extension_array_equal(result, expected)
 
 
+def test_ufunc_binary_output():
+    a = integer_array([1, 2, np.nan])
+    result = np.modf(a)
+    expected = np.modf(a.to_numpy(na_value=np.nan, dtype="float"))
+
+    assert isinstance(result, tuple)
+    assert len(result) == 2
+
+    for x, y in zip(result, expected):
+        # TODO(FloatArray): This will return an extension array.
+        # y = integer_array(y)
+        tm.assert_numpy_array_equal(x, y)
+
+
 @pytest.mark.parametrize("values", [[0, 1], [0, None]])
 def test_ufunc_reduce_raises(values):
     a = integer_array(values)

pandas/tests/base/test_misc.py

@@ -128,7 +128,8 @@ def test_memory_usage(index_or_series_obj):
     )
 
     if len(obj) == 0:
-        assert res_deep == res == 0
+        expected = 0 if isinstance(obj, Index) else 80
+        assert res_deep == res == expected
     elif is_object or is_categorical:
         # only deep will pick them up
         assert res_deep > res

pandas/tests/indexes/common.py

@@ -935,28 +935,22 @@ def test_contains_requires_hashable_raises(self):
         with pytest.raises(TypeError, match=msg):
             {} in idx._engine
 
-    def test_copy_copies_cache(self):
-        # GH32898
+    def test_copy_shares_cache(self):
+        # GH32898, GH36840
         idx = self.create_index()
         idx.get_loc(idx[0])  # populates the _cache.
         copy = idx.copy()
 
-        # check that the copied cache is a copy of the original
-        assert idx._cache == copy._cache
-        assert idx._cache is not copy._cache
-        # cache values should reference the same object
-        for key, val in idx._cache.items():
-            assert copy._cache[key] is val, key
+        assert copy._cache is idx._cache
 
-    def test_shallow_copy_copies_cache(self):
-        # GH32669
+    def test_shallow_copy_shares_cache(self):
+        # GH32669, GH36840
         idx = self.create_index()
         idx.get_loc(idx[0])  # populates the _cache.
         shallow_copy = idx._shallow_copy()
 
-        # check that the shallow_copied cache is a copy of the original
-        assert idx._cache == shallow_copy._cache
-        assert idx._cache is not shallow_copy._cache
-        # cache values should reference the same object
-        for key, val in idx._cache.items():
-            assert shallow_copy._cache[key] is val, key
+        assert shallow_copy._cache is idx._cache
+
+        shallow_copy = idx._shallow_copy(idx._data)
+        assert shallow_copy._cache is not idx._cache
+        assert shallow_copy._cache == {}

pandas/tests/window/test_grouper.py

@@ -428,3 +428,32 @@ def test_groupby_rolling_empty_frame(self):
         result = expected.groupby(["s1", "s2"]).rolling(window=1).sum()
         expected.index = pd.MultiIndex.from_tuples([], names=["s1", "s2", None])
         tm.assert_frame_equal(result, expected)
+
+    def test_groupby_rolling_string_index(self):
+        # GH: 36727
+        df = pd.DataFrame(
+            [
+                ["A", "group_1", pd.Timestamp(2019, 1, 1, 9)],
+                ["B", "group_1", pd.Timestamp(2019, 1, 2, 9)],
+                ["Z", "group_2", pd.Timestamp(2019, 1, 3, 9)],
+                ["H", "group_1", pd.Timestamp(2019, 1, 6, 9)],
+                ["E", "group_2", pd.Timestamp(2019, 1, 20, 9)],
+            ],
+            columns=["index", "group", "eventTime"],
+        ).set_index("index")
+
+        groups = df.groupby("group")
+        df["count_to_date"] = groups.cumcount()
+        rolling_groups = groups.rolling("10d", on="eventTime")
+        result = rolling_groups.apply(lambda df: df.shape[0])
+        expected = pd.DataFrame(
+            [
+                ["A", "group_1", pd.Timestamp(2019, 1, 1, 9), 1.0],
+                ["B", "group_1", pd.Timestamp(2019, 1, 2, 9), 2.0],
+                ["H", "group_1", pd.Timestamp(2019, 1, 6, 9), 3.0],
+                ["Z", "group_2", pd.Timestamp(2019, 1, 3, 9), 1.0],
+                ["E", "group_2", pd.Timestamp(2019, 1, 20, 9), 1.0],
+            ],
+            columns=["index", "group", "eventTime", "count_to_date"],
+        ).set_index(["group", "index"])
+        tm.assert_frame_equal(result, expected)

setup.cfg

@@ -154,15 +154,9 @@ check_untyped_defs=False
 [mypy-pandas.core.computation.expr]
 check_untyped_defs=False
 
-[mypy-pandas.core.computation.expressions]
-check_untyped_defs=False
-
 [mypy-pandas.core.computation.ops]
 check_untyped_defs=False
 
-[mypy-pandas.core.computation.pytables]
-check_untyped_defs=False
-
 [mypy-pandas.core.computation.scope]
 check_untyped_defs=False