You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/07/27 18:50:02 UTC
(airflow) 01/02: Keep DataFrame XComs working on pandas 3
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch pandas-3-support
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit aa9b69aa9f2b9649054ffc7970f88c6e610e5c97
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Mon Jul 27 18:25:05 2026 +0200
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.
---
.../amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py | 10 +++++++++-
.../salesforce/tests/unit/salesforce/hooks/test_salesforce.py | 9 +++++++--
task-sdk/src/airflow/sdk/serde/serializers/pandas.py | 4 ++++
task-sdk/tests/task_sdk/serde/test_serializers.py | 10 ++++++++--
4 files changed, 28 insertions(+), 5 deletions(-)
diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py b/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
index 5fb50fff9fe..3a960d1fd97 100644
--- a/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_sql_to_s3.py
@@ -24,12 +24,15 @@ from unittest import mock
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 @@ class TestSqlToS3Operator:
)
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")
diff --git a/providers/salesforce/tests/unit/salesforce/hooks/test_salesforce.py b/providers/salesforce/tests/unit/salesforce/hooks/test_salesforce.py
index 39e2e41fde6..923d59a3956 100644
--- a/providers/salesforce/tests/unit/salesforce/hooks/test_salesforce.py
+++ b/providers/salesforce/tests/unit/salesforce/hooks/test_salesforce.py
@@ -24,12 +24,15 @@ from unittest.mock import Mock, patch
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 @@ class TestSalesforceHook:
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,
)
diff --git a/task-sdk/src/airflow/sdk/serde/serializers/pandas.py b/task-sdk/src/airflow/sdk/serde/serializers/pandas.py
index 72a2e818a11..aaa68580abe 100644
--- a/task-sdk/src/airflow/sdk/serde/serializers/pandas.py
+++ b/task-sdk/src/airflow/sdk/serde/serializers/pandas.py
@@ -22,7 +22,11 @@ from typing import TYPE_CHECKING
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
diff --git a/task-sdk/tests/task_sdk/serde/test_serializers.py b/task-sdk/tests/task_sdk/serde/test_serializers.py
index 7e2ade64ebd..04510eeb1af 100644
--- a/task-sdk/tests/task_sdk/serde/test_serializers.py
+++ b/task-sdk/tests/task_sdk/serde/test_serializers.py
@@ -289,12 +289,18 @@ class TestSerializers:
@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
],