NaN in FLOAT/DOUBLE columns is returned as None (same as NULL) when pandas is enabled (#978)

src/databricks/sql/result_set.py

@@ -8,8 +8,10 @@
 
 try:
     import pyarrow
+    import pyarrow.compute as pc
 except ImportError:
     pyarrow = None
+    pc = None
 
 if TYPE_CHECKING:
     from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
@@ -130,6 +132,22 @@ def _convert_arrow_table(self, table):
         )
 
         res = df.to_numpy(na_value=None, dtype="object")
+
+        # Pandas' nullable Float dtypes treat IEEE NaN as missing (pd.NA), which
+        # to_numpy(na_value=None) turns into None, making NaN indistinguishable
+        # from SQL NULL. Restore NaN from the Arrow column so results match the
+        # non-pandas path (NULL -> None, NaN -> float('nan')).
+        for col_idx in range(table.num_columns):
+            column = table.column(col_idx)
+            if not pyarrow.types.is_floating(column.type):
+                continue
+            nan_mask = pc.is_nan(column)
+            # Cheap early-out: most float columns contain no NaN. pc.any skips
+            # nulls, so all-null/empty columns yield None and are skipped too.
+            if not pc.any(nan_mask).as_py():
+                continue
+            res[pc.fill_null(nan_mask, False).to_numpy(), col_idx] = float("nan")
+
         return [ResultRow(*v) for v in res]
 
     @property

tests/e2e/test_driver.py

@@ -476,6 +476,72 @@ def test_duplicate_column_names_survive_arrow_fetch(self):
             row = cursor.fetchone()
             assert list(row) == [1, 2]
 
+    @pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}])
+    def test_double_nan_distinct_from_null_with_pandas(self, extra_params):
+        """GH #978: with pandas enabled, IEEE NaN in a DOUBLE column must come
+        back as float('nan'), not None (which is reserved for SQL NULL)."""
+        import math
+
+        query = (
+            "SELECT * FROM VALUES (1, CAST('NaN' AS DOUBLE)), (2, CAST(NULL AS DOUBLE)), "
+            "(3, CAST('Infinity' AS DOUBLE)), (4, CAST(1.5 AS DOUBLE)) AS t(id, d) ORDER BY id"
+        )
+        with self.cursor(extra_params) as cursor:
+            cursor.execute(query)
+            rows = cursor.fetchall()
+
+        assert len(rows) == 4
+        assert isinstance(rows[0].d, float) and math.isnan(rows[0].d)
+        assert rows[1].d is None
+        assert rows[2].d == float("inf")
+        assert rows[3].d == 1.5
+
+    @pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}])
+    def test_float_nan_distinct_from_null_matches_disable_pandas(self, extra_params):
+        """GH #978: FLOAT (float32) NaN must be returned as NaN with pandas
+        enabled, matching the _disable_pandas=True output."""
+        import math
+
+        query = (
+            "SELECT * FROM VALUES (1, CAST('NaN' AS FLOAT)), (2, CAST(NULL AS FLOAT)), "
+            "(3, CAST('Infinity' AS FLOAT)), (4, CAST(1.5 AS FLOAT)) AS t(id, f) ORDER BY id"
+        )
+        with self.cursor(extra_params) as cursor:
+            cursor.execute(query)
+            pandas_rows = cursor.fetchall()
+        with self.cursor(dict(extra_params, _disable_pandas=True)) as cursor:
+            cursor.execute(query)
+            plain_rows = cursor.fetchall()
+
+        assert isinstance(pandas_rows[0].f, float) and math.isnan(pandas_rows[0].f)
+        assert pandas_rows[1].f is None
+
+        def _norm(v):
+            return "NaN" if isinstance(v, float) and math.isnan(v) else v
+
+        assert [tuple(_norm(v) for v in r) for r in pandas_rows] == [
+            tuple(_norm(v) for v in r) for r in plain_rows
+        ]
+
+    @pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}])
+    def test_nan_preserved_in_fetchone_and_fetchmany(self, extra_params):
+        """GH #978: fetchone/fetchmany share _convert_arrow_table and must also
+        keep NaN distinct from NULL."""
+        import math
+
+        query = "SELECT CAST('NaN' AS DOUBLE) AS d, CAST(NULL AS DOUBLE) AS n"
+        with self.cursor(extra_params) as cursor:
+            cursor.execute(query)
+            row = cursor.fetchone()
+            assert isinstance(row.d, float) and math.isnan(row.d)
+            assert row.n is None
+
+            cursor.execute(query)
+            rows = cursor.fetchmany(1)
+            assert len(rows) == 1
+            assert isinstance(rows[0].d, float) and math.isnan(rows[0].d)
+            assert rows[0].n is None
+
     def test_dml_rowcount(self):
         """cursor.rowcount reports the affected-row count for DML
         (INSERT/UPDATE/DELETE) instead of the hardcoded -1, and resets

tests/unit/test_pandas_compatibility.py

@@ -122,6 +122,51 @@ def test_float_types(self):
         self.assertIsNone(rows[1].float32_col)
         self.assertAlmostEqual(rows[2].float64_col, 4.5)
 
+    def test_float_nan_distinct_from_null(self):
+        """GH #978: NaN must stay NaN (not None) on the pandas path, across
+        chunked arrays and duplicate column names, matching disable_pandas."""
+        import math
+
+        nan = float("nan")
+        f32 = pa.chunked_array(
+            [
+                pa.array([nan, None], type=pa.float32()),
+                pa.array([float("inf"), 1.5], type=pa.float32()),
+            ]
+        )
+        f64 = pa.chunked_array(
+            [
+                pa.array([nan], type=pa.float64()),
+                pa.array([None, float("inf"), 1.5], type=pa.float64()),
+            ]
+        )
+        table = pa.Table.from_arrays([f32, f64, f64], names=["f", "d", "d"])
+        description = [
+            ("f", "float", None, None, None, None, None),
+            ("d", "double", None, None, None, None, None),
+            ("d", "double", None, None, None, None, None),
+        ]
+
+        pandas_rows = _make_result_set(description)._convert_arrow_table(table)
+        plain_rows = _make_result_set(description, disable_pandas=True)._convert_arrow_table(table)
+
+        self.assertEqual(len(pandas_rows), 4)
+        for col in range(3):
+            first = pandas_rows[0][col]
+            self.assertIsInstance(first, float)
+            self.assertTrue(math.isnan(first))
+            self.assertIsNone(pandas_rows[1][col])
+            self.assertEqual(pandas_rows[2][col], float("inf"))
+            self.assertAlmostEqual(pandas_rows[3][col], 1.5, places=5)
+
+        def _norm(v):
+            return "NaN" if isinstance(v, float) and math.isnan(v) else v
+
+        self.assertEqual(
+            [tuple(_norm(v) for v in r) for r in pandas_rows],
+            [tuple(_norm(v) for v in r) for r in plain_rows],
+        )
+
     def test_boolean_type(self):
         table = pa.table(
             {