ENH: Only apply first group once in fast GroupBy.apply

Author: fjetter

doc/source/user_guide/groupby.rst

@@ -946,23 +946,6 @@ that is itself a series, and possibly upcast the result to a DataFrame:
  So depending on the path taken, and exactly what you are grouping. Thus the grouped columns(s) may be included in
  the output as well as set the indices.
 
-.. warning::
-
- In the current implementation apply calls func twice on the
- first group to decide whether it can take a fast or slow code
- path. This can lead to unexpected behavior if func has
- side-effects, as they will take effect twice for the first
- group.
-
- .. ipython:: python
-
- d = pd.DataFrame({"a": ["x", "y"], "b": [1, 2]})
- def identity(df):
- print(df)
- return df
-
- d.groupby("a").apply(identity)
-
 
 Other useful features
 ---------------------

doc/source/whatsnew/v0.25.0.rst

@@ -26,6 +26,51 @@ Other Enhancements
 Backwards incompatible API changes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
+GroupBy.apply on ``DataFrame`` evaluates first group only once
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+(:issue:`2936`, :issue:`2656`, :issue:`7739`, :issue:`10519`, :issue:`12155`,
+:issue:`20084`, :issue:`21417`)
+
+The implementation of ``DataFrame.groupby.apply`` previously evaluated func
+consistently twice on the first group to infer if it is safe to use a fast
+code path. Particularly for functions with side effects, this was an undesired
+behavior and may have led to surprises.
+
+Now every group is evaluated only a single time.
+
+Previous behavior:
+
+.. code-block:: ipython
+
+ In [2]: df = pd.DataFrame({"a": ["x", "y"], "b": [1, 2]})
+
+ In [3]: side_effects = []
+
+ In [4]: def func_fast_apply(group):
+ ...: side_effects.append(group.name)
+ ...: return len(group)
+ ...:
+
+ In [5]: df.groupby("a").apply(func_fast_apply)
+
+ In [6]: assert side_effects == ["x", "x", "y"]
+
+New behavior:
+
+.. ipython:: python
+
+ df = pd.DataFrame({"a": ["x", "y"], "b": [1, 2]})
+
+ side_effects = []
+ def func(group):
+ side_effects.append(group.name)
+ return group
+
+ df.groupby("a").apply(func)
+ assert side_effects == ["x", "y"]
+
+
 .. _whatsnew_0250.api.other:
 
 Other API Changes

pandas/_libs/reduction.pyx

@@ -507,44 +507,41 @@ def apply_frame_axis0(object frame, object f, object names,
 
  results = []
 
- # Need to infer if our low-level mucking is going to cause a segfault
- if n > 0:
- chunk = frame.iloc[starts[0]:ends[0]]
- object.__setattr__(chunk, 'name', names[0])
- try:
- result = f(chunk)
- if result is chunk:
- raise InvalidApply('Function unsafe for fast apply')
- except:
- raise InvalidApply('Let this error raise above us')
-
  slider = BlockSlider(frame)
 
  mutated = False
+ status = 0
  item_cache = slider.dummy._item_cache
  try:
  for i in range(n):
  slider.move(starts[i], ends[i])
 
  item_cache.clear() # ugh
+ chunk = slider.dummy
+ object.__setattr__(chunk, 'name', names[i])
 
- object.__setattr__(slider.dummy, 'name', names[i])
- piece = f(slider.dummy)
-
- # I'm paying the price for index-sharing, ugh
  try:
- if piece.index is slider.dummy.index:
+ piece = f(chunk)
+ except:
+ raise InvalidApply('Let this error raise above us')
+ # Need to infer if low level index slider will cause segfaults
+ if i == 0 and piece is chunk:
+ status = 1
+ try:
+ if piece.index is chunk.index:
  piece = piece.copy(deep='all')
  else:
  mutated = True
  except AttributeError:
  pass
 
  results.append(piece)
+ if status > 0:
+ break
  finally:
  slider.reset()
 
- return results, mutated
+ return results, mutated, status
 
 
 cdef class BlockSlider:

pandas/core/groupby/ops.py

@@ -165,14 +165,17 @@ def apply(self, f, data, axis=0):
  mutated = self.mutated
  splitter = self._get_splitter(data, axis=axis)
  group_keys = self._get_group_keys()
-
+ status = 0
+ result_values = []
  # oh boy
  f_name = com.get_callable_name(f)
  if (f_name not in base.plotting_methods and
  hasattr(splitter, 'fast_apply') and axis == 0):
  try:
- values, mutated = splitter.fast_apply(f, group_keys)
- return group_keys, values, mutated
+ result = splitter.fast_apply(f, group_keys)
+ result_values, mutated, status = result
+ if status == 0:
+ return group_keys, result_values, mutated
  except reduction.InvalidApply:
  # we detect a mutation of some kind
  # so take slow path
@@ -181,9 +184,10 @@ def apply(self, f, data, axis=0):
  # raise this error to the caller
  pass
 
- result_values = []
  for key, (i, group) in zip(group_keys, splitter):
  object.__setattr__(group, 'name', key)
+ if status > 0 and i == 0:
+ continue
 
  # group might be modified
  group_axes = _get_axes(group)
@@ -854,10 +858,7 @@ def fast_apply(self, f, names):
  return [], True
 
  sdata = self._get_sorted_data()
- results, mutated = reduction.apply_frame_axis0(sdata, f, names,
- starts, ends)
-
- return results, mutated
+ return reduction.apply_frame_axis0(sdata, f, names, starts, ends)
 
  def _chop(self, sdata, slice_obj):
  if self.axis == 0:

pandas/tests/groupby/test_apply.py

@@ -101,10 +101,32 @@ def f(g):
  splitter = grouper._get_splitter(g._selected_obj, axis=g.axis)
  group_keys = grouper._get_group_keys()
 
- values, mutated = splitter.fast_apply(f, group_keys)
+ values, mutated, status = splitter.fast_apply(f, group_keys)
+ assert status == 0
  assert not mutated
 
 
+def test_group_apply_once_per_group():
+ # GH24748 ,GH2936, GH2656, GH7739, GH10519, GH12155, GH20084, GH21417
+ df = pd.DataFrame({'a': [0, 0, 1, 1, 2, 2], 'b': np.arange(6)})
+
+ names = []
+
+ def f_copy(group):
+ names.append(group.name)
+ return group.copy()
+ df.groupby("a").apply(f_copy)
+ assert names == [0, 1, 2]
+
+ def f_nocopy(group):
+ names.append(group.name)
+ return group
+ names = []
+ # this takes the slow apply path
+ df.groupby("a").apply(f_nocopy)
+ assert names == [0, 1, 2]
+
+
 def test_apply_with_mixed_dtype():
  # GH3480, apply with mixed dtype on axis=1 breaks in 0.11
  df = DataFrame({'foo1': np.random.randn(6),

pandas/tests/groupby/test_groupby.py

@@ -1420,20 +1420,30 @@ def foo(x):
 
 def test_group_name_available_in_inference_pass():
  # gh-15062
+ # GH24748 ,GH2936, GH2656, GH7739, GH10519, GH12155, GH20084, GH21417
  df = pd.DataFrame({'a': [0, 0, 1, 1, 2, 2], 'b': np.arange(6)})
 
  names = []
 
- def f(group):
+ def f_fast(group):
  names.append(group.name)
  return group.copy()
 
- df.groupby('a', sort=False, group_keys=False).apply(f)
- # we expect 2 zeros because we call ``f`` once to see if a faster route
- # can be used.
- expected_names = [0, 0, 1, 2]
+ df.groupby('a', sort=False, group_keys=False).apply(f_fast)
+
+ # every group should appear once, i.e. apply is called once per group
+ expected_names = [0, 1, 2]
  assert names == expected_names
 
+ names_slow = []
+
+ def f_slow(group):
+ names_slow.append(group.name)
+ return group
+
+ df.groupby('a', sort=False, group_keys=False).apply(f_slow)
+ assert names_slow == [0, 1, 2]
+
 
 def test_no_dummy_key_names(df):
  # see gh-1291