Support pandas 3+ in the dataframes API (#40709)

.github/trigger_files/beam_PostCommit_Python_Dependency.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run",
-    "modification": 7
+    "modification": 8
   }

sdks/python/apache_beam/dataframe/frame_base.py

@@ -32,6 +32,8 @@
 from apache_beam.dataframe import expressions
 from apache_beam.dataframe import partitionings
 
+PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2]))
+
 
 class DeferredBase(object):
 
@@ -491,7 +493,7 @@ def args_to_kwargs(base_type, removed_method=False, removed_args=None):
            running Pandas version.
   """
   def wrap(func):
-    if removed_method:
+    if removed_method or not hasattr(base_type, func.__name__):
       # Do no processing, let Beam function itself raise the error if called.
       return func
 
@@ -566,12 +568,12 @@ def with_docs_from(base_type, name=None, removed_method=False):
   version of Pandas.
   """
   def wrap(func):
-    if removed_method:
+    fn_name = name or func.__name__
+    if removed_method or not hasattr(base_type, fn_name):
       func.__doc__ = (
           "This method has been removed in the current version of Pandas.")
       return func
 
-    fn_name = name or func.__name__
     orig_doc = getattr(base_type, fn_name).__doc__
     if orig_doc is None:
       return func
@@ -659,7 +661,7 @@ def populate_defaults(base_type, removed_method=False, removed_args=None):
            running Pandas version.
   """
   def wrap(func):
-    if removed_method:
+    if removed_method or not hasattr(base_type, func.__name__):
       return func
 
     base_argspec = getfullargspec(unwrap(getattr(base_type, func.__name__)))
@@ -690,8 +692,10 @@ def wrap(func):
     # to mean that copy is the value of copy_on_write. Since
     # copy_on_write will always be true for Beam, just fill it
     # in here. In pandas 1, the default was True anyway.
+    # In pandas 3, Copy-on-Write is permanently enabled and methods
+    # either omit copy or set it to None.
     if 'copy' in arg_to_default and arg_to_default['copy'] is None:
-      arg_to_default['copy'] = True
+      arg_to_default['copy'] = None if PD_VERSION >= (3, 0) else True
 
     @functools.wraps(func)
     def wrapper(**kwargs):

sdks/python/apache_beam/dataframe/frames.py

@@ -264,10 +264,12 @@ def swaplevel(self, **kwargs):
             preserves_partition_by=partitionings.Arbitrary()))
 
   @frame_base.with_docs_from(pd.DataFrame)
-  @frame_base.args_to_kwargs(pd.DataFrame)
-  @frame_base.populate_defaults(pd.DataFrame)
+  @frame_base.args_to_kwargs(
+      pd.DataFrame, removed_args=['method'] if PD_VERSION >= (3, 0) else None)
+  @frame_base.populate_defaults(
+      pd.DataFrame, removed_args=['method'] if PD_VERSION >= (3, 0) else None)
   @frame_base.maybe_inplace
-  def fillna(self, value, method, axis, limit, **kwargs):
+  def fillna(self, value=None, method=None, axis=None, limit=None, **kwargs):
     """When ``axis="index"``, both ``method`` and ``limit`` must be ``None``.
     otherwise this operation is order-sensitive."""
     # Default value is None, but is overriden with index.
@@ -317,12 +319,21 @@ def __init__(self, value):
       get_value = lambda x: x
       requires = partitionings.Arbitrary()
 
+    def _do_fillna(df, value):
+      val = get_value(value)
+      if method is not None:
+        if PD_VERSION >= (3, 0):
+          fill_fn = getattr(df, method, None)
+          if fill_fn is not None:
+            return fill_fn(axis=axis, limit=limit)
+        return df.fillna(val, method=method, axis=axis, limit=limit, **kwargs)
+      else:
+        return df.fillna(val, axis=axis, limit=limit, **kwargs)
+
     return frame_base.DeferredFrame.wrap(
         # yapf: disable
         expressions.ComputedExpression(
-            'fillna', lambda df, value: df.fillna(
-                get_value(value), method=method, axis=axis, limit=limit, **
-                kwargs), [self._expr, value_expr],
+            'fillna', _do_fillna, [self._expr, value_expr],
             preserves_partition_by=partitionings.Arbitrary(),
             requires_partition_by=requires))
 
@@ -335,40 +346,44 @@ def __init__(self, value):
   if hasattr(pd.DataFrame, 'pad'):
     pad = _fillna_alias('pad')
 
-  @frame_base.with_docs_from(pd.DataFrame)
-  def first(self, offset):
-    per_partition = expressions.ComputedExpression(
-        'first-per-partition', lambda df: df.sort_index().first(offset=offset),
-        [self._expr],
-        preserves_partition_by=partitionings.Arbitrary(),
-        requires_partition_by=partitionings.Arbitrary())
-    with expressions.allow_non_parallel_operations(True):
-      return frame_base.DeferredFrame.wrap(
-          expressions.ComputedExpression(
-              'first', lambda df: df.sort_index().first(offset=offset),
-              [per_partition],
-              preserves_partition_by=partitionings.Arbitrary(),
-              requires_partition_by=partitionings.Singleton()))
+  if hasattr(pd.DataFrame, 'first'):
+    @frame_base.with_docs_from(pd.DataFrame)
+    def first(self, offset):
+      per_partition = expressions.ComputedExpression(
+          'first-per-partition', lambda df: df.sort_index().first(offset=offset),
+          [self._expr],
+          preserves_partition_by=partitionings.Arbitrary(),
+          requires_partition_by=partitionings.Arbitrary())
+      with expressions.allow_non_parallel_operations(True):
+        return frame_base.DeferredFrame.wrap(
+            expressions.ComputedExpression(
+                'first', lambda df: df.sort_index().first(offset=offset),
+                [per_partition],
+                preserves_partition_by=partitionings.Arbitrary(),
+                requires_partition_by=partitionings.Singleton()))
 
-  @frame_base.with_docs_from(pd.DataFrame)
-  def last(self, offset):
-    per_partition = expressions.ComputedExpression(
-        'last-per-partition', lambda df: df.sort_index().last(offset=offset),
-        [self._expr],
-        preserves_partition_by=partitionings.Arbitrary(),
-        requires_partition_by=partitionings.Arbitrary())
-    with expressions.allow_non_parallel_operations(True):
-      return frame_base.DeferredFrame.wrap(
-          expressions.ComputedExpression(
-              'last', lambda df: df.sort_index().last(offset=offset),
-              [per_partition],
-              preserves_partition_by=partitionings.Arbitrary(),
-              requires_partition_by=partitionings.Singleton()))
+  if hasattr(pd.DataFrame, 'last'):
+    @frame_base.with_docs_from(pd.DataFrame)
+    def last(self, offset):
+      per_partition = expressions.ComputedExpression(
+          'last-per-partition', lambda df: df.sort_index().last(offset=offset),
+          [self._expr],
+          preserves_partition_by=partitionings.Arbitrary(),
+          requires_partition_by=partitionings.Arbitrary())
+      with expressions.allow_non_parallel_operations(True):
+        return frame_base.DeferredFrame.wrap(
+            expressions.ComputedExpression(
+                'last', lambda df: df.sort_index().last(offset=offset),
+                [per_partition],
+                preserves_partition_by=partitionings.Arbitrary(),
+                requires_partition_by=partitionings.Singleton()))
 
   @frame_base.with_docs_from(pd.DataFrame)
-  @frame_base.args_to_kwargs(pd.DataFrame)
-  @frame_base.populate_defaults(pd.DataFrame)
-  def groupby(self, by, level, axis, as_index, group_keys, **kwargs):
+  @frame_base.args_to_kwargs(
+      pd.DataFrame, removed_args=['axis'] if PD_VERSION >= (3, 0) else None)
+  @frame_base.populate_defaults(
+      pd.DataFrame, removed_args=['axis'] if PD_VERSION >= (3, 0) else None)
+  def groupby(self, by=None, level=None, axis=0, as_index=True, group_keys=True, **kwargs):
     """``as_index`` must be ``True``.
 
     Aggregations grouping by a categorical column with ``observed=False`` set
@@ -379,6 +394,9 @@ def groupby(self, by, level, axis, as_index, group_keys, **kwargs):
       raise NotImplementedError('groupby(as_index=False)')
 
     if axis in (1, 'columns'):
+      if PD_VERSION >= (3, 0):
+        raise NotImplementedError(
+            "groupby(axis=1) is removed in Pandas 3.0+")
       return _DeferredGroupByCols(
           expressions.ComputedExpression(
               'groupbycols', lambda df: df.groupby(
@@ -3034,7 +3052,8 @@ def aggregate(self, func, axis, *args, **kwargs):
 
   agg = aggregate
 
-  applymap = frame_base._elementwise_method('applymap', base=pd.DataFrame)
+  if hasattr(pd.DataFrame, 'applymap'):
+    applymap = frame_base._elementwise_method('applymap', base=pd.DataFrame)
   if PD_VERSION >= (2, 1):
     map = frame_base._elementwise_method('map', base=pd.DataFrame)
   add_prefix = frame_base._elementwise_method('add_prefix', base=pd.DataFrame)
@@ -4847,17 +4866,20 @@ class _DeferredGroupByCols(frame_base.DeferredFrame):
   describe = frame_base.not_implemented_method('describe',
                                                base_type=DataFrameGroupBy)
   diff = frame_base._elementwise_method('diff', base=DataFrameGroupBy)
-  fillna = frame_base._elementwise_method('fillna', base=DataFrameGroupBy)
+  if hasattr(DataFrameGroupBy, 'fillna'):
+    fillna = frame_base._elementwise_method('fillna', base=DataFrameGroupBy)
   filter = frame_base._elementwise_method('filter', base=DataFrameGroupBy)
-  first = frame_base._elementwise_method('first', base=DataFrameGroupBy)
+  if hasattr(DataFrameGroupBy, 'first'):
+    first = frame_base._elementwise_method('first', base=DataFrameGroupBy)
   get_group = frame_base._elementwise_method('get_group', base=DataFrameGroupBy)
   head = frame_base.wont_implement_method(
       DataFrameGroupBy, 'head', explanation=_PEEK_METHOD_EXPLANATION)
   hist = frame_base.wont_implement_method(
       DataFrameGroupBy, 'hist', reason="plotting-tools")
   idxmax = frame_base._elementwise_method('idxmax', base=DataFrameGroupBy)
   idxmin = frame_base._elementwise_method('idxmin', base=DataFrameGroupBy)
-  last = frame_base._elementwise_method('last', base=DataFrameGroupBy)
+  if hasattr(DataFrameGroupBy, 'last'):
+    last = frame_base._elementwise_method('last', base=DataFrameGroupBy)
   max = frame_base._elementwise_method('max', base=DataFrameGroupBy)
   mean = frame_base._elementwise_method('mean', base=DataFrameGroupBy)
   median = frame_base._elementwise_method('median', base=DataFrameGroupBy)
@@ -5170,7 +5192,7 @@ def get_dummies(self, **kwargs):
             'get_dummies',
             lambda series: pd.concat(
               [proxy, series.str.get_dummies(**kwargs)]
-              ).fillna(value=0, method=None).astype('int64'),
+              ).fillna(value=0).astype('int64'),
             [self._expr],
             proxy=proxy,
             requires_partition_by=partitionings.Arbitrary(),

sdks/python/apache_beam/dataframe/frames_test.py

@@ -946,7 +946,9 @@ def test_corrwith_bad_axis(self):
     self._run_error_test(lambda df: df.corrwith(df.a, axis=2), df)
     self._run_error_test(lambda df: df.corrwith(df, axis=5), df)
 
-  @unittest.skipIf(PD_VERSION < (1, 2), "na_action added in pandas 1.2.0")
+  @unittest.skipIf(
+      PD_VERSION < (1, 2) or PD_VERSION >= (3, 0),
+      "na_action added in pandas 1.2.0; applymap removed in pandas 3.0")
   @pytest.mark.filterwarnings(
       "ignore:The default of observed=False is deprecated:FutureWarning")
   def test_applymap_na_action(self):
@@ -1054,11 +1056,16 @@ def test_fillna_columns(self):
          [np.nan, 3, np.nan, 4], [3, np.nan, np.nan, 4]],
         columns=list('ABCD'))
 
-    self._run_test(lambda df: df.fillna(method='ffill', axis='columns'), df)
-    self._run_test(
-        lambda df: df.fillna(method='ffill', axis='columns', limit=1), df)
-    self._run_test(
-        lambda df: df.fillna(method='bfill', axis='columns', limit=1), df)
+    if PD_VERSION < (3, 0):
+      self._run_test(lambda df: df.fillna(method='ffill', axis='columns'), df)
+      self._run_test(
+          lambda df: df.fillna(method='ffill', axis='columns', limit=1), df)
+      self._run_test(
+          lambda df: df.fillna(method='bfill', axis='columns', limit=1), df)
+    else:
+      self._run_test(lambda df: df.ffill(axis='columns'), df)
+      self._run_test(lambda df: df.ffill(axis='columns', limit=1), df)
+      self._run_test(lambda df: df.bfill(axis='columns', limit=1), df)
 
     # Intended behavior is unclear here. See
     # https://github.com/pandas-dev/pandas/issues/40989

sdks/python/apache_beam/typehints/pandas_type_compatibility.py

@@ -52,6 +52,7 @@
 """
 
 # ruff: noqa: UP006
+import datetime
 from typing import Any
 from typing import List
 from typing import Optional
@@ -130,8 +131,16 @@ def dtype_to_fieldtype(dtype):
 
   if fieldtype is not None:
     return fieldtype
-  elif dtype.kind == 'S':
+  elif getattr(dtype, 'kind', None) == 'O' or dtype is object:
+    return Any
+  elif (str(dtype) in ('str', 'string') or
+        isinstance(dtype, pd.StringDtype) or
+        pd.api.types.is_string_dtype(dtype)):
+    return Optional[str]
+  elif getattr(dtype, 'kind', None) == 'S':
     return bytes
+  elif getattr(dtype, 'kind', None) == 'M':
+    return Optional[datetime.datetime]
   else:
     return Any
 

sdks/python/setup.py

@@ -159,8 +159,9 @@ def cythonize(*args, **kwargs):
 # must update the below "docs" and "test" for extras_require
 dataframe_dependency = [
     # pandas 2.2.2 is the first version that works with numpy 2.0.0+
-    # TODO(https://github.com/apache/beam/issues/40709): support pandas 3+
-    'pandas>=2.2.2,<2.4',
+    # Python 3.10 is not supported by pandas 3.0+
+    'pandas>=2.2.2,<2.4; python_version < "3.11"',
+    'pandas>=2.2.2,<3.2; python_version >= "3.11"',
 ]
 
 milvus_dependency = ['pymilvus>=2.5.10,<3.0.0']
@@ -491,7 +492,8 @@ def get_portability_package_data():
               'docstring-parser>=0.15,<1.0',
               'docutils>=0.18.1',
               'markdown',
-              'pandas<2.6.0',
+              'pandas<2.4; python_version < "3.11"',
+              'pandas<3.2; python_version >= "3.11"',
               'openai',
               'virtualenv-clone>=0.5,<1.0',
           ],
@@ -502,7 +504,8 @@ def get_portability_package_data():
               'jinja2>=3.0,<3.2',
               'joblib>=1.0.1',
               'mock>=1.0.1,<6.0.0',
-              'pandas<2.6.0',
+              'pandas<2.4; python_version < "3.11"',
+              'pandas<3.2; python_version >= "3.11"',
               'parameterized>=0.7.1,<0.10.0',
               'pydot>=1.2.0,<2',
               'pyhamcrest>=1.9,!=1.10.0,<3.0.0',

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

@@ -94,6 +94,10 @@ toxTask "testPy311pandas-23", "py311-pandas-23", "${posargs}"
 test.dependsOn "testPy311pandas-23"
 postCommitPyDep.dependsOn "testPy311pandas-23"
 
+toxTask "testPy311pandas-30", "py311-pandas-30", "${posargs}"
+test.dependsOn "testPy311pandas-30"
+postCommitPyDep.dependsOn "testPy311pandas-30"
+
 // TODO(https://github.com/apache/beam/issues/30908): Revise what are we testing
 
 // Create a test task for the lower and upper bounds of supported pytorch.

sdks/python/tox.ini

@@ -365,14 +365,15 @@ commands =
   /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pyarrow {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
 
-[testenv:py311-pandas-{22,23}]
+[testenv:py311-pandas-{22,23,30}]
 # Keep in line with dataframe_dependency in setup.py. pandas 2.2.2 is the first
 # release compatible with numpy 2.
 deps =
   # Lower bound of the supported range; also keeps numpy 1.x covered.
   22: pandas>=2.2.2,<2.3.0
   22: numpy>=1.23.2,<2.0.0
   23: pandas>=2.3.0,<2.4.0
+  30: pandas>=3.0.0,<3.1.0
 commands =
   # Log pandas and numpy version for debugging
   /bin/sh -c "pip freeze | grep -E '(pandas|numpy)'"