diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f24a07505..fa3cae951 100644 --- a/src/databricks/sql/result_set.py +++ b/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 diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 319b7de5e..6b111478b 100644 --- a/tests/e2e/test_driver.py +++ b/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 diff --git a/tests/unit/test_pandas_compatibility.py b/tests/unit/test_pandas_compatibility.py index 5434559b2..e266c7ba0 100644 --- a/tests/unit/test_pandas_compatibility.py +++ b/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( {