Skip to content

Commit 3e2fad4

Browse files
author
peco-engineer-bot[bot]
committed
NaN in FLOAT/DOUBLE columns is returned as None (same as NULL) when pandas is enabled (#978)
Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 85b69ea commit 3e2fad4

3 files changed

Lines changed: 128 additions & 0 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,23 @@ def _convert_arrow_table(self, table):
130130
)
131131

132132
res = df.to_numpy(na_value=None, dtype="object")
133+
134+
import pyarrow.compute as pc
135+
136+
# Pandas' nullable Float dtypes treat IEEE NaN as missing (pd.NA), which
137+
# to_numpy(na_value=None) turns into None, making NaN indistinguishable
138+
# from SQL NULL. Restore NaN from the Arrow column so results match the
139+
# non-pandas path (NULL -> None, NaN -> float('nan')).
140+
for col_idx in range(table.num_columns):
141+
column = table.column(col_idx)
142+
if not pyarrow.types.is_floating(column.type):
143+
continue
144+
nan_mask = pc.fill_null(pc.is_nan(column), False)
145+
if not pc.any(nan_mask).as_py():
146+
continue
147+
for row_idx in nan_mask.to_numpy().nonzero()[0]:
148+
res[row_idx, col_idx] = float("nan")
149+
133150
return [ResultRow(*v) for v in res]
134151

135152
@property

‎tests/e2e/test_driver.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,6 +476,72 @@ def test_duplicate_column_names_survive_arrow_fetch(self):
476476
row = cursor.fetchone()
477477
assert list(row) == [1, 2]
478478

479+
@pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}])
480+
def test_double_nan_distinct_from_null_with_pandas(self, extra_params):
481+
"""GH #978: with pandas enabled, IEEE NaN in a DOUBLE column must come
482+
back as float('nan'), not None (which is reserved for SQL NULL)."""
483+
import math
484+
485+
query = (
486+
"SELECT * FROM VALUES (1, CAST('NaN' AS DOUBLE)), (2, CAST(NULL AS DOUBLE)), "
487+
"(3, CAST('Infinity' AS DOUBLE)), (4, CAST(1.5 AS DOUBLE)) AS t(id, d) ORDER BY id"
488+
)
489+
with self.cursor(extra_params) as cursor:
490+
cursor.execute(query)
491+
rows = cursor.fetchall()
492+
493+
assert len(rows) == 4
494+
assert isinstance(rows[0].d, float) and math.isnan(rows[0].d)
495+
assert rows[1].d is None
496+
assert rows[2].d == float("inf")
497+
assert rows[3].d == 1.5
498+
499+
@pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}])
500+
def test_float_nan_distinct_from_null_matches_disable_pandas(self, extra_params):
501+
"""GH #978: FLOAT (float32) NaN must be returned as NaN with pandas
502+
enabled, matching the _disable_pandas=True output."""
503+
import math
504+
505+
query = (
506+
"SELECT * FROM VALUES (1, CAST('NaN' AS FLOAT)), (2, CAST(NULL AS FLOAT)), "
507+
"(3, CAST('Infinity' AS FLOAT)), (4, CAST(1.5 AS FLOAT)) AS t(id, f) ORDER BY id"
508+
)
509+
with self.cursor(extra_params) as cursor:
510+
cursor.execute(query)
511+
pandas_rows = cursor.fetchall()
512+
with self.cursor(dict(extra_params, _disable_pandas=True)) as cursor:
513+
cursor.execute(query)
514+
plain_rows = cursor.fetchall()
515+
516+
assert isinstance(pandas_rows[0].f, float) and math.isnan(pandas_rows[0].f)
517+
assert pandas_rows[1].f is None
518+
519+
def _norm(v):
520+
return "NaN" if isinstance(v, float) and math.isnan(v) else v
521+
522+
assert [tuple(_norm(v) for v in r) for r in pandas_rows] == [
523+
tuple(_norm(v) for v in r) for r in plain_rows
524+
]
525+
526+
@pytest.mark.parametrize("extra_params", [{}, {"use_sea": True}])
527+
def test_nan_preserved_in_fetchone_and_fetchmany(self, extra_params):
528+
"""GH #978: fetchone/fetchmany share _convert_arrow_table and must also
529+
keep NaN distinct from NULL."""
530+
import math
531+
532+
query = "SELECT CAST('NaN' AS DOUBLE) AS d, CAST(NULL AS DOUBLE) AS n"
533+
with self.cursor(extra_params) as cursor:
534+
cursor.execute(query)
535+
row = cursor.fetchone()
536+
assert isinstance(row.d, float) and math.isnan(row.d)
537+
assert row.n is None
538+
539+
cursor.execute(query)
540+
rows = cursor.fetchmany(1)
541+
assert len(rows) == 1
542+
assert isinstance(rows[0].d, float) and math.isnan(rows[0].d)
543+
assert rows[0].n is None
544+
479545
def test_dml_rowcount(self):
480546
"""cursor.rowcount reports the affected-row count for DML
481547
(INSERT/UPDATE/DELETE) instead of the hardcoded -1, and resets

‎tests/unit/test_pandas_compatibility.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,51 @@ def test_float_types(self):
122122
self.assertIsNone(rows[1].float32_col)
123123
self.assertAlmostEqual(rows[2].float64_col, 4.5)
124124

125+
def test_float_nan_distinct_from_null(self):
126+
"""GH #978: NaN must stay NaN (not None) on the pandas path, across
127+
chunked arrays and duplicate column names, matching disable_pandas."""
128+
import math
129+
130+
nan = float("nan")
131+
f32 = pa.chunked_array(
132+
[
133+
pa.array([nan, None], type=pa.float32()),
134+
pa.array([float("inf"), 1.5], type=pa.float32()),
135+
]
136+
)
137+
f64 = pa.chunked_array(
138+
[
139+
pa.array([nan], type=pa.float64()),
140+
pa.array([None, float("inf"), 1.5], type=pa.float64()),
141+
]
142+
)
143+
table = pa.Table.from_arrays([f32, f64, f64], names=["f", "d", "d"])
144+
description = [
145+
("f", "float", None, None, None, None, None),
146+
("d", "double", None, None, None, None, None),
147+
("d", "double", None, None, None, None, None),
148+
]
149+
150+
pandas_rows = _make_result_set(description)._convert_arrow_table(table)
151+
plain_rows = _make_result_set(description, disable_pandas=True)._convert_arrow_table(table)
152+
153+
self.assertEqual(len(pandas_rows), 4)
154+
for col in range(3):
155+
first = pandas_rows[0][col]
156+
self.assertIsInstance(first, float)
157+
self.assertTrue(math.isnan(first))
158+
self.assertIsNone(pandas_rows[1][col])
159+
self.assertEqual(pandas_rows[2][col], float("inf"))
160+
self.assertAlmostEqual(pandas_rows[3][col], 1.5, places=5)
161+
162+
def _norm(v):
163+
return "NaN" if isinstance(v, float) and math.isnan(v) else v
164+
165+
self.assertEqual(
166+
[tuple(_norm(v) for v in r) for r in pandas_rows],
167+
[tuple(_norm(v) for v in r) for r in plain_rows],
168+
)
169+
125170
def test_boolean_type(self):
126171
table = pa.table(
127172
{

0 commit comments

Comments
 (0)