From 220fdd69552543ea866655eb7c8867867b75a82d Mon Sep 17 00:00:00 2001
From: Maha Rana Yadavalli <271375718+maharanay22@users.noreply.github.com>
Date: Fri, 16 Oct 2026 20:39:34 -0400
Subject: [PATCH] Keep Databricks integer columns exact when they contain NULLs

native_query built the result DataFrame straight from the cursor rows, so
a NULL in a BIGINT or INT column made pandas store the column as float64.
Values above 2**53 lost precision (9007199254740993 came back as
9007199254740992.0) and small integers turned into floats (3 -> 3.0).

Rebuild integer and boolean columns that contain NULLs with the pandas
nullable Int64 and boolean dtypes. Columns without NULLs keep their
current dtypes.
---
 .../databricks_handler/databricks_handler.py  | 28 ++++++++++++++++++-
 .../community_handlers/test_databricks.py     | 18 ++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/mindsdb/integrations/handlers/databricks_handler/databricks_handler.py b/mindsdb/integrations/handlers/databricks_handler/databricks_handler.py
index 755308d419..6dcb447f6a 100644
--- a/mindsdb/integrations/handlers/databricks_handler/databricks_handler.py
+++ b/mindsdb/integrations/handlers/databricks_handler/databricks_handler.py
@@ -54,6 +54,32 @@ def _validate_identifier(identifier: str) -> str:
     return identifier
 
 
+def _make_data_frame(result: List[Tuple], columns: List[str]) -> pd.DataFrame:
+    """Build a DataFrame from query rows, keeping integer and boolean columns that contain NULLs exact.
+
+    Without this, pandas turns an integer column with a NULL into float64, which loses precision above 2**53.
+
+    Args:
+        result (List[Tuple]): rows returned by the cursor.
+        columns (List[str]): column names.
+
+    Returns:
+        pd.DataFrame: the query result.
+    """
+    data_frame = pd.DataFrame(result, columns=columns)
+    for index in range(len(columns)):
+        dtype = data_frame.dtypes.iloc[index]
+        if pd.api.types.is_integer_dtype(dtype) or pd.api.types.is_bool_dtype(dtype):
+            continue
+        values = [row[index] for row in result]
+        value_types = {type(value) for value in values if value is not None}
+        if value_types == {int}:
+            data_frame.isetitem(index, pd.array(values, dtype="Int64"))
+        elif value_types == {bool}:
+            data_frame.isetitem(index, pd.array(values, dtype="boolean"))
+    return data_frame
+
+
 def _parse_interval_value(value: Any) -> Optional[int]:
     if isinstance(value, bool):
         return None
@@ -407,7 +433,7 @@ def native_query(self, query: Text) -> Response:
                 if cursor.description:
                     response = Response(
                         RESPONSE_TYPE.TABLE,
-                        data_frame=pd.DataFrame(result, columns=[x[0] for x in cursor.description]),
+                        data_frame=_make_data_frame(result, [x[0] for x in cursor.description]),
                     )
                 else:
                     response = Response(RESPONSE_TYPE.OK)
diff --git a/tests/unit/handlers/community_handlers/test_databricks.py b/tests/unit/handlers/community_handlers/test_databricks.py
index 3b968ada2d..ac8061fb46 100644
--- a/tests/unit/handlers/community_handlers/test_databricks.py
+++ b/tests/unit/handlers/community_handlers/test_databricks.py
@@ -189,6 +189,24 @@ def test_native_query_empty_select_returns_table(self):
         self.assertEqual(list(response.data_frame.columns), ["id", "name"])
         self.assertEqual(len(response.data_frame), 0)
 
+    def test_native_query_keeps_integers_with_nulls_exact(self):
+        self.mock_cursor.set_results(
+            [(9007199254740993, 3, True, 1.5), (None, None, None, None)],
+            ["id", "qty", "flag", "price"],
+        )
+
+        response = self.handler.native_query("SELECT id, qty, flag, price FROM orders")
+
+        df = response.data_frame
+        self.assertEqual(df["id"].dtype, "Int64")
+        self.assertEqual(df["id"][0], 9007199254740993)
+        self.assertEqual(df["qty"].dtype, "Int64")
+        self.assertEqual(df["qty"][0], 3)
+        self.assertEqual(df["flag"].dtype, "boolean")
+        self.assertTrue(df["flag"][0])
+        self.assertEqual(df["price"].dtype, "float64")
+        self.assertTrue(df.iloc[1].isna().all())
+
     def test_get_tables(self):
         """
         Tests if the `get_tables` method to confirm it correctly calls `native_query` with the appropriate SQL commands.