Keep Databricks integer columns exact when they contain NULLs

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)

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.