Skip to content

Commit 5074510

Browse files
committed
remove calls to obj.dropna()
1 parent 394feb6 commit 5074510

File tree

2 files changed

+16
-5
lines changed

2 files changed

+16
-5
lines changed

pandas/core/groupby/generic.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -538,14 +538,19 @@ def _transform_general(
538538
if isinstance(res, (ABCDataFrame, ABCSeries)):
539539
res = res._values
540540

541-
results.append(klass(res, index=group.index))
541+
indexer = self._get_index(name) if self.dropna else group.index
542+
results.append(klass(res, index=indexer))
542543

543544
# check for empty "results" to avoid concat ValueError
544545
if results:
545546
from pandas.core.reshape.concat import concat
546547

547548
concatenated = concat(results)
548-
result = self._set_result_index_ordered(concatenated)
549+
550+
if not self.dropna:
551+
result = self._set_result_index_ordered(concatenated)
552+
else:
553+
result = concatenated.sort_index()
549554
else:
550555
result = self.obj._constructor(dtype=np.float64)
551556
# we will only try to coerce the result type if

pandas/core/groupby/groupby.py

+9-3
Original file line numberDiff line numberDiff line change
@@ -679,7 +679,7 @@ def _set_group_selection(self):
679679
self._group_selection = ax.difference(Index(groupers), sort=False).tolist()
680680
self._reset_cache("_selected_obj")
681681

682-
def _set_result_index_ordered(self, result):
682+
def _set_result_index_ordered(self, result, dropna=False):
683683
# set the result index on the passed values object and
684684
# return the new object, xref 8046
685685

@@ -690,8 +690,14 @@ def _set_result_index_ordered(self, result):
690690
result.set_axis(index, axis=self.axis, inplace=True)
691691
result = result.sort_index(axis=self.axis)
692692

693-
obj = self.obj.dropna() if self.dropna else self.obj
694-
result.set_axis(obj._get_axis(self.axis), axis=self.axis, inplace=True)
693+
if not dropna:
694+
result.set_axis(self.obj._get_axis(self.axis), axis=self.axis, inplace=True)
695+
else:
696+
result.set_axis(
697+
self._obj_with_exclusions._get_axis(self.axis),
698+
axis=self.axis,
699+
inplace=True,
700+
)
695701
return result
696702

697703
def _dir_additions(self):

0 commit comments

Comments
 (0)