Pandas 1.5 support (#23973)

* [WIP] Pandas 1.5 support

* Partial progress

* Exclude pandas 1.5.0 and 1.5.1 because of pandas-dev/pandas#45725

* Plumb group_keys more places

* Fix bad argument

* Fix bad argument

* Debug

* Debug

* Debug

* Debug

* allow list tests

* Update changes

* fmt

* Lint

* Lint

* Fix pd_version

* Fix pd_version

* Remove CHANGES.md since this needs to go in 2.45 section

Author: damccorm

sdks/python/apache_beam/dataframe/frames.py

@@ -372,24 +372,24 @@ def last(self, offset):
   @frame_base.args_to_kwargs(pd.DataFrame)
   @frame_base.populate_defaults(pd.DataFrame)
   def groupby(self, by, level, axis, as_index, group_keys, **kwargs):
-    """``as_index`` and ``group_keys`` must both be ``True``.
+    """``as_index`` must be ``True``.
 
     Aggregations grouping by a categorical column with ``observed=False`` set
     are not currently parallelizable
     (`Issue 21827 <https://github.com/apache/beam/issues/21827>`_).
     """
     if not as_index:
       raise NotImplementedError('groupby(as_index=False)')
-    if not group_keys:
-      raise NotImplementedError('groupby(group_keys=False)')
 
     if axis in (1, 'columns'):
       return _DeferredGroupByCols(
           expressions.ComputedExpression(
               'groupbycols',
-              lambda df: df.groupby(by, axis=axis, **kwargs), [self._expr],
+              lambda df: df.groupby(
+                  by, axis=axis, group_keys=group_keys, **kwargs), [self._expr],
               requires_partition_by=partitionings.Arbitrary(),
-              preserves_partition_by=partitionings.Arbitrary()))
+              preserves_partition_by=partitionings.Arbitrary()),
+          group_keys=group_keys)
 
     if level is None and by is None:
       raise TypeError("You have to supply one of 'by' and 'level'")
@@ -559,14 +559,17 @@ def prepend_index(df, by):  # type: ignore
         expressions.ComputedExpression(
             'groupbyindex',
             lambda df: df.groupby(
-                level=list(range(df.index.nlevels)), **kwargs), [to_group],
+                level=list(range(df.index.nlevels)),
+                group_keys=group_keys,
+                **kwargs), [to_group],
             requires_partition_by=partitionings.Index(),
             preserves_partition_by=partitionings.Arbitrary()),
         kwargs,
         to_group,
         to_group_with_index,
         grouping_columns=grouping_columns,
-        grouping_indexes=grouping_indexes)
+        grouping_indexes=grouping_indexes,
+        group_keys=group_keys)
 
   @property  # type: ignore
   @frame_base.with_docs_from(pd.DataFrame)
@@ -676,6 +679,7 @@ def replace(self, to_replace, value, limit, method, **kwargs):
     order-sensitive. It cannot be specified.
 
     If ``limit`` is specified this operation is not parallelizable."""
+    # pylint: disable-next=c-extension-no-member
     value_compare = None if PD_VERSION < (1, 4) else lib.no_default
     if method is not None and not isinstance(to_replace,
                                              dict) and value is value_compare:
@@ -4123,6 +4127,7 @@ def __init__(self, expr, kwargs,
                ungrouped_with_index: expressions.Expression[pd.core.generic.NDFrame], # pylint: disable=line-too-long
                grouping_columns,
                grouping_indexes,
+               group_keys,
                projection=None):
     """This object represents the result of::
 
@@ -4149,6 +4154,7 @@ def __init__(self, expr, kwargs,
     self._projection = projection
     self._grouping_columns = grouping_columns
     self._grouping_indexes = grouping_indexes
+    self._group_keys = group_keys
     self._kwargs = kwargs
 
     if (self._kwargs.get('dropna', True) is False and
@@ -4170,6 +4176,7 @@ def __getattr__(self, name):
         self._ungrouped_with_index,
         self._grouping_columns,
         self._grouping_indexes,
+        self._group_keys,
         projection=name)
 
   def __getitem__(self, name):
@@ -4184,6 +4191,7 @@ def __getitem__(self, name):
         self._ungrouped_with_index,
         self._grouping_columns,
         self._grouping_indexes,
+        self._group_keys,
         projection=name)
 
   @frame_base.with_docs_from(DataFrameGroupBy)
@@ -4233,6 +4241,7 @@ def apply(self, func, *args, **kwargs):
     project = _maybe_project_func(self._projection)
     grouping_indexes = self._grouping_indexes
     grouping_columns = self._grouping_columns
+    group_keys = self._group_keys
 
     # Unfortunately pandas does not execute func to determine the right proxy.
     # We run user func on a proxy here to detect the return type and generate
@@ -4321,7 +4330,8 @@ def do_partition_apply(df):
       df = df.reset_index(grouping_columns, drop=True)
 
       gb = df.groupby(level=grouping_indexes or None,
-                      by=grouping_columns or None)
+                      by=grouping_columns or None,
+                      group_keys=group_keys)
 
       gb = project(gb)
 
@@ -4361,6 +4371,7 @@ def fn_wrapper(x, *args, **kwargs):
       fn_wrapper = fn
 
     project = _maybe_project_func(self._projection)
+    group_keys = self._group_keys
 
     # pandas cannot execute fn to determine the right proxy.
     # We run user fn on a proxy here to detect the return type and generate the
@@ -4387,10 +4398,12 @@ def fn_wrapper(x, *args, **kwargs):
     return DeferredDataFrame(
         expressions.ComputedExpression(
             'transform',
-            lambda df: project(df.groupby(level=levels)).transform(
-                fn_wrapper,
-                *args,
-                **kwargs).droplevel(self._grouping_columns),
+            lambda df: project(
+              df.groupby(level=levels, group_keys=group_keys)
+            ).transform(
+              fn_wrapper,
+              *args,
+              **kwargs).droplevel(self._grouping_columns),
             [self._ungrouped_with_index],
             proxy=proxy,
             requires_partition_by=partitionings.Index(levels),
@@ -4551,6 +4564,7 @@ def wrapper(self, *args, **kwargs):
     is_categorical_grouping = any(to_group.get_level_values(i).is_categorical()
                                   for i in self._grouping_indexes)
     groupby_kwargs = self._kwargs
+    group_keys = self._group_keys
 
     # Don't include un-observed categorical values in the preagg
     preagg_groupby_kwargs = groupby_kwargs.copy()
@@ -4562,6 +4576,7 @@ def wrapper(self, *args, **kwargs):
         lambda df: getattr(
             project(
                 df.groupby(level=list(range(df.index.nlevels)),
+                           group_keys=group_keys,
                            **preagg_groupby_kwargs)
             ),
             agg_name)(**kwargs),
@@ -4574,6 +4589,7 @@ def wrapper(self, *args, **kwargs):
         'post_combine_' + post_agg_name,
         lambda df: getattr(
             df.groupby(level=list(range(df.index.nlevels)),
+                       group_keys=group_keys,
                        **groupby_kwargs),
             post_agg_name)(**kwargs),
         [pre_agg],
@@ -4597,6 +4613,7 @@ def wrapper(self, *args, **kwargs):
     assert isinstance(self, DeferredGroupBy)
 
     to_group = self._ungrouped.proxy().index
+    group_keys = self._group_keys
     is_categorical_grouping = any(to_group.get_level_values(i).is_categorical()
                                   for i in self._grouping_indexes)
 
@@ -4606,6 +4623,7 @@ def wrapper(self, *args, **kwargs):
         agg_name,
         lambda df: getattr(project(
             df.groupby(level=list(range(df.index.nlevels)),
+                       group_keys=group_keys,
                        **groupby_kwargs),
         ), agg_name)(**kwargs),
         [self._ungrouped],

sdks/python/apache_beam/dataframe/pandas_doctests_test.py

@@ -127,10 +127,7 @@ def test_ndframe_tests(self):
             'pandas.core.generic.NDFrame.copy': ['*'],
             'pandas.core.generic.NDFrame.droplevel': ['*'],
             'pandas.core.generic.NDFrame.get': ['*'],
-            'pandas.core.generic.NDFrame.rank': [
-                # Modified dataframe
-                'df'
-            ],
+            'pandas.core.generic.NDFrame.rank': ['*'],
             'pandas.core.generic.NDFrame.rename': [
                 # Seems to be an upstream bug. The actual error has a different
                 # message:
@@ -704,14 +701,19 @@ def test_groupby_tests(self):
     result = doctests.testmod(
         pd.core.groupby.groupby,
         use_beam=False,
+        verbose=True,
         wont_implement_ok={
+            'pandas.core.groupby.groupby.GroupBy.first': ['*'],
             'pandas.core.groupby.groupby.GroupBy.head': ['*'],
+            'pandas.core.groupby.groupby.GroupBy.last': ['*'],
             'pandas.core.groupby.groupby.GroupBy.tail': ['*'],
             'pandas.core.groupby.groupby.GroupBy.nth': ['*'],
             'pandas.core.groupby.groupby.GroupBy.cumcount': ['*'],
             'pandas.core.groupby.groupby.GroupBy.resample': ['*'],
         },
         not_implemented_ok={
+            'pandas.core.groupby.groupby.GroupBy.first': ['*'],
+            'pandas.core.groupby.groupby.GroupBy.last': ['*'],
             'pandas.core.groupby.groupby.GroupBy.ngroup': ['*'],
             'pandas.core.groupby.groupby.GroupBy.sample': ['*'],
             'pandas.core.groupby.groupby.GroupBy.rank': ['*'],
@@ -831,6 +833,7 @@ def test_top_level(self):
             'crosstab': ['*'],
             'cut': ['*'],
             'eval': ['*'],
+            'from_dummies': ['*'],
             'get_dummies': ['*'],
             'infer_freq': ['*'],
             'lreshape': ['*'],
@@ -863,7 +866,10 @@ def test_top_level(self):
         },
         skip={
             # error formatting
-            'concat': ['pd.concat([df5, df6], verify_integrity=True)'],
+            'concat': [
+                'pd.concat([df5, df6], verify_integrity=True)',
+                'pd.concat([df7, new_row.to_frame().T], ignore_index=True)'
+            ],
             # doctest DeprecationWarning
             'melt': ['df'],
             # Order-sensitive re-indexing.

sdks/python/setup.py

@@ -349,8 +349,12 @@ def get_portability_package_data():
         # with python 3.10 leading to incorrect stacktrace.
         # This can be removed once dill is updated to version > 0.3.5.1
         # Issue: https://github.com/apache/beam/issues/23566
-          'dataframe': ['pandas>=1.0,<1.5;python_version<"3.10"',
-                        'pandas>=1.4.3,<1.5;python_version>="3.10"'],
+        # Exclude 1.5.0 and 1.5.1 because of
+        # https://github.com/pandas-dev/pandas/issues/45725
+          'dataframe': [
+            'pandas>=1.0,<1.6,!=1.5.0,!=1.5.1;python_version<"3.10"',
+            'pandas>=1.4.3,<1.6,!=1.5.0,!=1.5.1;python_version>="3.10"'
+          ],
           'dask': [
             'dask >= 2022.6',
             'distributed >= 2022.6',

sdks/python/test-suites/tox/py38/build.gradle

@@ -93,6 +93,10 @@ toxTask "testPy38pandas-14", "py38-pandas-14"
 test.dependsOn "testPy38pandas-14"
 preCommitPy38.dependsOn "testPy38pandas-14"
 
+toxTask "testPy38pandas-15", "py38-pandas-15"
+test.dependsOn "testPy38pandas-15"
+preCommitPy38.dependsOn "testPy38pandas-15"
+
 // Create a test task for each minor version of pytorch
 toxTask "testPy38pytorch-19", "py38-pytorch-19"
 test.dependsOn "testPy38pytorch-19"

sdks/python/tox.ini

@@ -299,12 +299,14 @@ commands =
   # selecting tests with -m (BEAM-12985)
   pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pyarrow {posargs}
 
-[testenv:py{37,38,39,310}-pandas-{11,12,13,14}]
+[testenv:py{37,38,39,310}-pandas-{11,12,13,14,15}]
 deps =
   11: pandas>=1.1.0,<1.2.0
   12: pandas>=1.2.0,<1.3.0
   13: pandas>=1.3.0,<1.4.0
   14: pandas>=1.4.0,<1.5.0
+  # Exclude 1.5.0 and 1.5.1 because of https://github.com/pandas-dev/pandas/issues/45725
+  15: pandas>=1.5.2,<1.6.0
 commands =
   # Log pandas and numpy version for debugging
   /bin/sh -c "pip freeze | grep -E '(pandas|numpy)'"