[peco-engineer-bot]

Summary

Automated fix for #978 — NaN in FLOAT/DOUBLE columns is returned as None (same as NULL) when pandas is enabled.

With pandas enabled, a NaN in a FLOAT or DOUBLE column now comes back as float('nan') instead of None, so it can be told apart from SQL NULL; NULL still comes back as None. In ResultSet._convert_arrow_table (src/databricks/sql/result_set.py), after the pandas to_numpy conversion, the code now uses pyarrow.compute.is_nan to find the NaN cells in each floating-point Arrow column and puts NaN back in those cells. The output now matches the _disable_pandas=True path. All 4 new tests failed before the change and pass after it: the 3 E2E tests ran against the live warehouse on both Thrift and SEA (6 cases), plus the unit test. The full unit suite also passes (1267 passed, run with -m "not realkernel").

Root cause & plan

Root cause: ResultSet._convert_arrow_table (src/databricks/sql/result_set.py, ~L109-133) maps pyarrow float32/float64 to pandas' nullable Float32Dtype/Float64Dtype in the to_pandas types_mapper. That conversion treats IEEE NaN as missing (pd.NA), and df.to_numpy(na_value=None, dtype="object") then turns it into None. The result is that NaN can't be told apart from SQL NULL. The _disable_pandas=True path uses as_py() and correctly returns float('nan'). All three backends (Thrift, SEA, kernel) call this shared method from fetchone/fetchmany/fetchall, so all of them are affected.
Files: src/databricks/sql/result_set.py, tests/e2e/test_driver.py, tests/unit/test_pandas_compatibility.py
Planned coverage:

• In tests/e2e/test_driver.py add a test parametrized with extra_params [{}, {"use_sea": True}]. It runs e.g. 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 the default (pandas enabled) connection and calls fetchall(). Assert: row0.d is a float and math.isnan(row0.d); row1.d is None; row2.d == float('inf'); row3.d == 1.5. Today row0.d is None, so the test fails. (DOUBLE NaN vs NULL, pandas enabled (E2E, Thrift + SEA))
• Same as above but with CAST(... AS FLOAT) values. Assert NaN comes back as a float where math.isnan is true, and NULL comes back as None. Also check that the result matches the _disable_pandas=True output for the same query, comparing NaN with isnan. (FLOAT (float32) NaN vs NULL, pandas enabled (E2E))
• Run SELECT CAST('NaN' AS DOUBLE) AS d, CAST(NULL AS DOUBLE) AS n and call fetchone(). Assert math.isnan(row.d) and row.n is None. These paths share _convert_arrow_table, so this guards all fetch variants. (fetchone/fetchmany paths also preserve NaN (E2E))
• In tests/unit/test_pandas_compatibility.py add a test that builds a pa.table with float32/float64 arrays containing [nan, None, inf, 1.5], split across multiple chunks via pa.chunked_array, plus duplicate column names. Assert the pandas path gives nan, None, inf, 1.5 and equals the disable_pandas path, comparing NaN with isnan. This is in addition to the E2E tests. _(Unit edge cases for convert_arrow_table)

Files changed

• tests/e2e/test_driver.py
• tests/unit/test_pandas_compatibility.py
• src/databricks/sql/result_set.py

Test plan

• tests/e2e/test_driver.py::TestPySQLCoreSuite::test_double_nan_distinct_from_null_with_pandas — fails (red) against the original code, passes (green) after the fix
• tests/e2e/test_driver.py::TestPySQLCoreSuite::test_float_nan_distinct_from_null_matches_disable_pandas — fails (red) against the original code, passes (green) after the fix
• tests/e2e/test_driver.py::TestPySQLCoreSuite::test_nan_preserved_in_fetchone_and_fetchmany — fails (red) against the original code, passes (green) after the fix
• tests/unit/test_pandas_compatibility.py::TestConvertArrowTablePandasCompat::test_float_nan_distinct_from_null — fails (red) against the original code, passes (green) after the fix

🤖 Generated by engineer-bot (bug-fix flow) — review before merge.

[peco-review-bot]

Verdict: 1 Low

Looks good: the fix restores NaN from the Arrow column after the pandas conversion, so NaN and NULL are now distinct. It covers chunked arrays and duplicate column names and is tested at both the unit and E2E levels. One low-severity suggestion: the per-cell restore loop could be a single vectorized assignment.

[peco-review-bot]

Verdict: 1 Low

Looks good. The fix is correct and well-scoped, with one low-severity performance note. The Arrow-side is_nan mask (with nulls filled as False) is applied after pandas conversion, so it handles chunked columns, duplicate column names and empty tables. If a pandas version already keeps NaN, the fix just writes NaN again, so nothing changes.

[peco-review-bot]

✅ No issues identified by the review bot.