From 7e536f2cf7157afd79796fd4e291490801d31bec Mon Sep 17 00:00:00 2001 From: Dhanush A G Date: Sun, 25 Oct 2026 00:21:30 +0530 Subject: [PATCH 1/2] Support pandas 3+ in the dataframes API (#40709) * Enable pandas 3.x dependency on Python 3.11+ via PEP 508 environment markers * Update tox.ini and Gradle build configs to test pandas 3.0 in py311 * Add defensive attribute checking in frame_base decorators (args_to_kwargs, populate_defaults, with_docs_from) * Support pandas 3.0 Copy-on-Write defaults * Update fillna, groupby, first, last, and applymap to handle removed pandas 3.0 arguments and methods * Add pandas 3.0 string and datetime dtype recognition in pandas_type_compatibility * Update frames_test.py to support pandas 3.0 --- .../apache_beam/dataframe/frame_base.py | 14 ++- sdks/python/apache_beam/dataframe/frames.py | 106 +++++++++++------- .../apache_beam/dataframe/frames_test.py | 19 +++- .../typehints/pandas_type_compatibility.py | 7 +- sdks/python/setup.py | 11 +- .../python/test-suites/tox/py311/build.gradle | 4 + sdks/python/tox.ini | 3 +- 7 files changed, 105 insertions(+), 59 deletions(-) diff --git a/sdks/python/apache_beam/dataframe/frame_base.py b/sdks/python/apache_beam/dataframe/frame_base.py index f5bf024b177b..fe68f28706c6 100644 --- a/sdks/python/apache_beam/dataframe/frame_base.py +++ b/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): diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py index f99df35afe83..f23ed4bd89b6 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/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,9 +4866,11 @@ 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) @@ -4857,7 +4878,8 @@ class _DeferredGroupByCols(frame_base.DeferredFrame): 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(), diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index fbe3ef36d804..cf54106f1956 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/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 diff --git a/sdks/python/apache_beam/typehints/pandas_type_compatibility.py b/sdks/python/apache_beam/typehints/pandas_type_compatibility.py index 45ae27baffe7..a1efb008edc6 100644 --- a/sdks/python/apache_beam/typehints/pandas_type_compatibility.py +++ b/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,12 @@ def dtype_to_fieldtype(dtype): if fieldtype is not None: return fieldtype - elif dtype.kind == 'S': + elif str(dtype) in ('str', 'string') or isinstance(dtype, pd.StringDtype): + return Optional[str] + elif getattr(dtype, 'kind', None) == 'S': return bytes + elif getattr(dtype, 'kind', None) == 'M': + return Optional[datetime.datetime] else: return Any diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 77a5b835a059..d0947936d4e8 100644 --- a/sdks/python/setup.py +++ b/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', diff --git a/sdks/python/test-suites/tox/py311/build.gradle b/sdks/python/test-suites/tox/py311/build.gradle index a05e9972bef4..f02baa164aa6 100644 --- a/sdks/python/test-suites/tox/py311/build.gradle +++ b/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. diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index c69fc2e45ef7..9afd1288fc38 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -365,7 +365,7 @@ 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 = @@ -373,6 +373,7 @@ deps = 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)'" From 36ee9801f86404df86d55fe5e6db251b2c7d7850 Mon Sep 17 00:00:00 2001 From: Dhanush A G Date: Sun, 25 Oct 2026 00:35:19 +0530 Subject: [PATCH 2/2] Address review: support pd.api.types.is_string_dtype and trigger postcommit dependency CI --- .../trigger_files/beam_PostCommit_Python_Dependency.json | 2 +- .../apache_beam/typehints/pandas_type_compatibility.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index 9d28ab31a100..0c0e8ca6a12b 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Dependency.json +++ b/.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 } diff --git a/sdks/python/apache_beam/typehints/pandas_type_compatibility.py b/sdks/python/apache_beam/typehints/pandas_type_compatibility.py index a1efb008edc6..410abcc94f1e 100644 --- a/sdks/python/apache_beam/typehints/pandas_type_compatibility.py +++ b/sdks/python/apache_beam/typehints/pandas_type_compatibility.py @@ -131,7 +131,11 @@ def dtype_to_fieldtype(dtype): if fieldtype is not None: return fieldtype - elif str(dtype) in ('str', 'string') or isinstance(dtype, pd.StringDtype): + 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