Add support for pandas 3 based xcoms in airflow (#71103) (#71169)

* Keep DataFrame XComs working on pandas 3

pandas 3 exposes its public classes from the `pandas` namespace, so a
DataFrame is now qualified as `pandas.DataFrame` rather than
`pandas.core.frame.DataFrame`. The serializer was registered only under the
old name, so pushing a DataFrame through XCom raised "cannot serialize
object of type <class 'pandas.DataFrame'>". Both names are registered so
values written by either version stay readable.

pandas 3 also infers a str column where it used to infer object, and keeps
its missing values as NA instead of stringifying them, which the amazon and
salesforce tests asserted on.

* Document pandas 3 impact on DataFrame XComs

Deployments need to know that every component has to carry the pandas 3
support before pandas 3 reaches any worker, that a rollback strands the
XComs written in the meantime, and that a pulled DataFrame now takes its
dtypes from the reader's pandas version.

* Name the pandas 3 newsfragment after its own pull request

* Add regression test for the pandas 2/3 cross-version registry lookup

* rename newsfragment file

---------


(cherry picked from commit 4091ccf)

Co-authored-by: Jarek Potiuk <jarek@potiuk.com>

Author: 2 people

airflow-core/newsfragments/71103.significant.rst

@@ -0,0 +1,33 @@
+pandas 3 changes how DataFrame XComs are stored and read back
+
+pandas 3 exposes its public classes from the ``pandas`` namespace, so a DataFrame is qualified as
+``pandas.DataFrame`` instead of ``pandas.core.frame.DataFrame``. XComs record that name alongside the
+serialized value, so the name written into the metadata database depends on the pandas version of the
+component that pushed the value. Airflow registers both names, and a DataFrame written by either
+pandas version can be read by either — no configuration change is needed, and existing XComs stay
+readable.
+
+What you should do:
+
+* **Roll this Airflow version out to every component before pandas 3 reaches any of them** — workers
+  in particular. A component that predates this change cannot read a DataFrame XCom written under
+  pandas 3, and fails the pull with:
+
+  .. code-block:: text
+
+      ImportError: pandas.DataFrame was not found in allow list for deserialization imports.
+      To allow it, add it to allowed_deserialization_classes in the configuration
+
+  The message points at configuration, but the allow list is not the cause and changing it does not
+  help. The rows are not corrupt: they become readable again as soon as the reader is upgraded.
+
+* **Treat a downgrade as a one-way door for those XComs.** Rolling back to an Airflow version without
+  this change strands any DataFrame XCom written while on pandas 3, with the same error, until you
+  roll forward again.
+
+* **Review Dags that inspect the dtypes of a pulled DataFrame.** The pandas version of the *reader*
+  determines what a pulled DataFrame looks like, not the version that wrote it. Under pandas 3, a
+  column of strings comes back with the ``str`` dtype rather than ``object``, and its missing values
+  come back as ``nan`` rather than ``None``. Values are unchanged, but downstream code that branches
+  on ``dtype == "object"``, checks cells with ``is None``, or compares against a reference frame with
+  ``DataFrame.equals()`` can behave differently after the upgrade.

providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py

@@ -24,12 +24,15 @@
 import pandas as pd
 import polars as pl
 import pytest
+from packaging.version import Version
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.models import Connection
 from airflow.providers.amazon.aws.transfers.sql_to_s3 import SqlToS3Operator
 from airflow.providers.common.compat.sdk import AirflowException
 
+PANDAS_3_PLUS = Version(pd.__version__).major >= 3
+
 
 class TestSqlToS3Operator:
     @pytest.mark.parametrize(
@@ -156,7 +159,12 @@ def test_fix_dtypes(self, params):
         )
         dirty_df = pd.DataFrame({"strings": ["a", "b", None], "ints": [1, 2, None]})
         op._fix_dtypes(df=dirty_df, file_format=op.file_format)
-        assert dirty_df["strings"].values[2] == params["null_string_result"]
+        if PANDAS_3_PLUS:
+            # pandas 3 infers a str column rather than object, and keeps its missing values as NA
+            # instead of the object None (csv) or the "None" it used to be stringified to (parquet)
+            assert pd.isna(dirty_df["strings"].values[2])
+        else:
+            assert dirty_df["strings"].values[2] == params["null_string_result"]
         assert dirty_df["ints"].dtype.kind == "i"
 
     @mock.patch("airflow.providers.amazon.aws.transfers.sql_to_s3.S3Hook")

providers/salesforce/tests/unit/salesforce/hooks/test_salesforce.py

@@ -24,12 +24,15 @@
 import numpy as np
 import pandas as pd
 import pytest
+from packaging.version import Version
 from requests import Session as request_session
 from simple_salesforce import Salesforce, api
 
 from airflow.models.connection import Connection
 from airflow.providers.salesforce.hooks.salesforce import SalesforceHook
 
+PANDAS_3_PLUS = Version(pd.__version__).major >= 3
+
 
 class TestSalesforceHook:
     def setup_method(self):
@@ -349,10 +352,12 @@ def test_write_object_to_file_csv(self, mock_data_frame):
         data_frame = self.salesforce_hook.write_object_to_file(query_results=[], filename=filename, fmt="csv")
 
         mock_data_frame.return_value.to_csv.assert_called_once_with(filename, index=False)
-        # Note that the latest version of pandas dataframes (1.1.2) returns "nan" rather than "None" here
+        # Note that the latest version of pandas dataframes (1.1.2) returns "nan" rather than "None" here,
+        # and pandas 3 keeps missing values as NA instead of stringifying them to "nan" at all
+        missing = np.nan if PANDAS_3_PLUS else "nan"
         pd.testing.assert_frame_equal(
             data_frame,
-            pd.DataFrame({"test": [1, 2, 3], "dict": ["nan", "nan", str({"foo": "bar"})]}),
+            pd.DataFrame({"test": [1, 2, 3], "dict": [missing, missing, str({"foo": "bar"})]}),
             check_index_type=False,
         )
 

task-sdk/src/airflow/sdk/serde/serializers/pandas.py

@@ -22,7 +22,11 @@
 from airflow.sdk.module_loading import qualname
 
 # lazy loading for performance reasons
+# pandas 3 moved the public classes to the `pandas` namespace, so a DataFrame is qualified as
+# `pandas.DataFrame` there and as `pandas.core.frame.DataFrame` on pandas 2. Both are registered so
+# that XComs serialized by either version stay readable.
 serializers = [
+    "pandas.DataFrame",
     "pandas.core.frame.DataFrame",
 ]
 deserializers = serializers

task-sdk/tests/task_sdk/serde/test_serializers.py

@@ -281,6 +281,25 @@ def test_pandas(self):
         d = deserialize(e)
         assert i.equals(d)
 
+    @pytest.mark.parametrize(
+        "classname",
+        ["pandas.DataFrame", "pandas.core.frame.DataFrame"],
+    )
+    def test_pandas_deserializes_regardless_of_writer_qualname(self, classname):
+        """
+        A DataFrame XCom must deserialize under either pandas major's registry qualname.
+
+        The installed pandas only ever produces its own qualname, so this forges the tag to
+        force the lookup through the other registry entry.
+        """
+        import pandas as pd
+
+        i = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
+        e = serialize(i)
+        e[CLASSNAME] = classname
+        d = deserialize(e)
+        assert i.equals(d)
+
     def test_pandas_serializers(self):
         from airflow.sdk.serde.serializers.pandas import serialize
 
@@ -289,12 +308,18 @@ def test_pandas_serializers(self):
     @pytest.mark.parametrize(
         ("klass", "version", "data", "msg"),
         [
-            (pd.DataFrame, 999, "", r"serialized 999 of pandas.core.frame.DataFrame > 1"),  # version too new
+            # pandas 3 qualifies the class as pandas.DataFrame, pandas 2 as pandas.core.frame.DataFrame
+            (
+                pd.DataFrame,
+                999,
+                "",
+                r"serialized 999 of pandas(\.core\.frame)?\.DataFrame > 1",
+            ),  # version too new
             (
                 pd.DataFrame,
                 1,
                 123,
-                r"serialized pandas.core.frame.DataFrame has wrong data type .*<class 'int'>",
+                r"serialized pandas(\.core\.frame)?\.DataFrame has wrong data type .*<class 'int'>",
             ),  # bad payload type
             (str, 1, "", r"do not know how to deserialize builtins.str"),  # bad class
         ],