From ec82a2a59ea5c0025264354b3ed8c917832fbafb Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Thu, 19 Mar 2020 16:32:06 -0700 Subject: [PATCH] REF: pass align_keys to apply --- pandas/core/internals/managers.py | 38 +++++++++++++++---------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 66e96af05eb71..e83b71bd966cc 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -27,7 +27,7 @@ ) from pandas.core.dtypes.concat import concat_compat from pandas.core.dtypes.dtypes import ExtensionDtype -from pandas.core.dtypes.generic import ABCExtensionArray, ABCSeries +from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries from pandas.core.dtypes.missing import isna import pandas.core.algorithms as algos @@ -375,7 +375,7 @@ def reduce(self, func, *args, **kwargs): return res - def apply(self: T, f, filter=None, **kwargs) -> T: + def apply(self: T, f, filter=None, align_keys=None, **kwargs) -> T: """ Iterate over the blocks, collect and create a new BlockManager. @@ -390,6 +390,7 @@ def apply(self: T, f, filter=None, **kwargs) -> T: ------- BlockManager """ + align_keys = align_keys or [] result_blocks = [] # fillna: Series/DataFrame is responsible for making sure value is aligned @@ -404,28 +405,14 @@ def apply(self: T, f, filter=None, **kwargs) -> T: self._consolidate_inplace() + align_copy = False if f == "where": align_copy = True - if kwargs.get("align", True): - align_keys = ["other", "cond"] - else: - align_keys = ["cond"] - elif f == "putmask": - align_copy = False - if kwargs.get("align", True): - align_keys = ["new", "mask"] - else: - align_keys = ["mask"] - else: - align_keys = [] - # TODO(EA): may interfere with ExtensionBlock.setitem for blocks - # with a .values attribute. aligned_args = { k: kwargs[k] for k in align_keys - if not isinstance(kwargs[k], ABCExtensionArray) - and hasattr(kwargs[k], "values") + if isinstance(kwargs[k], (ABCSeries, ABCDataFrame)) } for b in self.blocks: @@ -561,13 +548,24 @@ def isna(self, func) -> "BlockManager": return self.apply("apply", func=func) def where(self, **kwargs) -> "BlockManager": - return self.apply("where", **kwargs) + if kwargs.pop("align", True): + align_keys = ["other", "cond"] + else: + align_keys = ["cond"] + + return self.apply("where", align_keys=align_keys, **kwargs) def setitem(self, indexer, value) -> "BlockManager": return self.apply("setitem", indexer=indexer, value=value) def putmask(self, **kwargs): - return self.apply("putmask", **kwargs) + + if kwargs.pop("align", True): + align_keys = ["new", "mask"] + else: + align_keys = ["mask"] + + return self.apply("putmask", align_keys=align_keys, **kwargs) def diff(self, n: int, axis: int) -> "BlockManager": return self.apply("diff", n=n, axis=axis)