From 67bc5fbd306be7cdffd216f3791d4024acfa95b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christian=20Br=C3=A6dstrup?=
 <3591721+LinuxChristian@users.noreply.github.com>
Date: Mon, 12 Jul 2021 21:21:24 +0200
Subject: [PATCH] fix: use pandas function to check for NaN (#750)

* fix: use pandas function to check for NaN

Starting with pandas 1.0, an experimental pandas.NA value (singleton) is available to represent scalar missing values as
opposed to numpy.nan. Comparing the variable with itself results in a pandas.NA value that doesn't support type-casting
to boolean. Using the build-in pandas.isna function handles all pandas supported NaN values.

* tests: Skip tests if pandas below required version

* tests: compare expected and actual directly as lists

* Fix pytest.mark.skipif spelling

Co-authored-by: Peter Lamut <plamut@users.noreply.github.com>
---
 google/cloud/bigquery/_pandas_helpers.py |  2 +-
 tests/unit/test__pandas_helpers.py       | 40 ++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
index e93a99eba..285c0e83c 100644
--- a/google/cloud/bigquery/_pandas_helpers.py
+++ b/google/cloud/bigquery/_pandas_helpers.py
@@ -780,7 +780,7 @@ def dataframe_to_json_generator(dataframe):
         output = {}
         for column, value in zip(dataframe.columns, row):
             # Omit NaN values.
-            if value != value:
+            if pandas.isna(value):
                 continue
             output[column] = value
         yield output
diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py
index 39a3d845b..aa87e28f5 100644
--- a/tests/unit/test__pandas_helpers.py
+++ b/tests/unit/test__pandas_helpers.py
@@ -19,6 +19,7 @@
 import operator
 import queue
 import warnings
+import pkg_resources
 
 import mock
 
@@ -47,6 +48,14 @@
 except ImportError:  # pragma: NO COVER
     bigquery_storage = None
 
+PANDAS_MINIUM_VERSION = pkg_resources.parse_version("1.0.0")
+
+if pandas is not None:
+    PANDAS_INSTALLED_VERSION = pkg_resources.get_distribution("pandas").parsed_version
+else:
+    # Set to less than MIN version.
+    PANDAS_INSTALLED_VERSION = pkg_resources.parse_version("0.0.0")
+
 
 skip_if_no_bignumeric = pytest.mark.skipif(
     not _BIGNUMERIC_SUPPORT, reason="BIGNUMERIC support requires pyarrow>=3.0.0",
@@ -734,6 +743,37 @@ def test_list_columns_and_indexes_with_named_index_same_as_column_name(
     assert columns_and_indexes == expected
 
 
+@pytest.mark.skipif(
+    pandas is None or PANDAS_INSTALLED_VERSION < PANDAS_MINIUM_VERSION,
+    reason="Requires `pandas version >= 1.0.0` which introduces pandas.NA",
+)
+def test_dataframe_to_json_generator(module_under_test):
+    utcnow = datetime.datetime.utcnow()
+    df_data = collections.OrderedDict(
+        [
+            ("a_series", [pandas.NA, 2, 3, 4]),
+            ("b_series", [0.1, float("NaN"), 0.3, 0.4]),
+            ("c_series", ["a", "b", pandas.NA, "d"]),
+            ("d_series", [utcnow, utcnow, utcnow, pandas.NaT]),
+            ("e_series", [True, False, True, None]),
+        ]
+    )
+    dataframe = pandas.DataFrame(
+        df_data, index=pandas.Index([4, 5, 6, 7], name="a_index")
+    )
+
+    dataframe = dataframe.astype({"a_series": pandas.Int64Dtype()})
+
+    rows = module_under_test.dataframe_to_json_generator(dataframe)
+    expected = [
+        {"b_series": 0.1, "c_series": "a", "d_series": utcnow, "e_series": True},
+        {"a_series": 2, "c_series": "b", "d_series": utcnow, "e_series": False},
+        {"a_series": 3, "b_series": 0.3, "d_series": utcnow, "e_series": True},
+        {"a_series": 4, "b_series": 0.4, "c_series": "d"},
+    ]
+    assert list(rows) == expected
+
+
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
 def test_list_columns_and_indexes_with_named_index(module_under_test):
     df_data = collections.OrderedDict(