Keep SQL Server integer columns exact when they contain NULLs

mindsdb/integrations/handlers/mssql_handler/mssql_handler.py

@@ -1,5 +1,6 @@
 from typing import Any, Union, TYPE_CHECKING
 import datetime
+import decimal
 
 import pymssql
 from pymssql import OperationalError
@@ -59,6 +60,38 @@ def _map_type(mssql_type_text: str) -> MYSQL_DATA_TYPE:
     return MYSQL_DATA_TYPE.VARCHAR
 
 
+# pyodbc reports the Python type of each column in cursor.description[1]
+_ODBC_PYTHON_TYPES = {
+    bool: MYSQL_DATA_TYPE.TINYINT,
+    int: MYSQL_DATA_TYPE.INT,
+    float: MYSQL_DATA_TYPE.FLOAT,
+    decimal.Decimal: MYSQL_DATA_TYPE.DECIMAL,
+    datetime.datetime: MYSQL_DATA_TYPE.DATETIME,
+    datetime.date: MYSQL_DATA_TYPE.DATE,
+    datetime.time: MYSQL_DATA_TYPE.TIME,
+    str: MYSQL_DATA_TYPE.TEXT,
+    bytes: MYSQL_DATA_TYPE.BINARY,
+    bytearray: MYSQL_DATA_TYPE.BINARY,
+}
+
+
+def _nullable_dtype(values: list[Any]) -> str | None:
+    """Pick a pandas nullable dtype for integer or bit values, so NULLs don't turn the column into float or object.
+
+    Args:
+        values (list[Any]): the values of one column.
+
+    Returns:
+        str | None: 'Int64' or 'boolean' if all non-NULL values are int or bool, otherwise None.
+    """
+    value_types = {type(value) for value in values if value is not None}
+    if value_types == {int}:
+        return "Int64"
+    if value_types == {bool}:
+        return "boolean"
+    return None
+
+
 def _make_table_response(
     result: list[Union[dict[str, Any], tuple]], cursor: Union[pymssql.Cursor, "pyodbc.Cursor"], use_odbc: bool = False
 ) -> Response:
@@ -85,13 +118,27 @@ def _make_table_response(
         # pymssql with as_dict=True returns list of dicts
         data_frame = pd.DataFrame(result)
 
-    for column in description:
+    for column_index, column in enumerate(description):
         column_name = column[0]
         column_type = column[1]
         column_dtype = data_frame[column_name].dtype
 
-        if use_odbc:
-            # For pyodbc, use type inference based on pandas dtype
+        if (
+            result
+            and column_type in (pymssql.NUMBER, int, bool)
+            and not (pd_types.is_integer_dtype(column_dtype) or pd_types.is_bool_dtype(column_dtype))
+        ):
+            key = column_index if use_odbc else column_name
+            values = [row[key] for row in result]
+            nullable_dtype = _nullable_dtype(values)
+            if nullable_dtype is not None:
+                data_frame[column_name] = pd.array(values, dtype=nullable_dtype)
+                column_dtype = data_frame[column_name].dtype
+
+        if use_odbc and column_type in _ODBC_PYTHON_TYPES:
+            mysql_types.append(_ODBC_PYTHON_TYPES[column_type])
+        elif use_odbc:
+            # Fall back to type inference based on pandas dtype
             if pd_types.is_integer_dtype(column_dtype):
                 mysql_types.append(MYSQL_DATA_TYPE.INT)
             elif pd_types.is_float_dtype(column_dtype):

tests/unit/handlers/test_mssql.py

@@ -932,6 +932,43 @@ def test_types_casting(self):
             self.assertEqual(result_value, input_value)
         # endregion
 
+    def test_types_casting_with_nulls(self):
+        """Test that NULLs in integer and bit columns don't turn them into float or object"""
+        mock_conn = MagicMock()
+        mock_cursor = MagicMock()
+        mock_cursor.__enter__ = MagicMock(return_value=mock_cursor)
+        mock_cursor.__exit__ = MagicMock(return_value=None)
+
+        self.handler.connect = MagicMock(return_value=mock_conn)
+        mock_conn.cursor = MagicMock(return_value=mock_cursor)
+
+        mock_cursor.fetchall.return_value = [
+            {"n_bigint": 9007199254740993, "n_int": 3, "n_bit": True, "n_float": 1.5},
+            {"n_bigint": None, "n_int": None, "n_bit": None, "n_float": None},
+        ]
+        mock_cursor.description = [
+            ("n_bigint", 3, None, None, None, None, None),
+            ("n_int", 3, None, None, None, None, None),
+            ("n_bit", 3, None, None, None, None, None),
+            ("n_float", 3, None, None, None, None, None),
+        ]
+
+        response: TableResponse = self.handler.native_query("SELECT * FROM test_table")
+
+        self.assertEqual(
+            [col.type for col in response.columns],
+            [MYSQL_DATA_TYPE.INT, MYSQL_DATA_TYPE.INT, MYSQL_DATA_TYPE.TINYINT, MYSQL_DATA_TYPE.FLOAT],
+        )
+        df = response.data_frame
+        self.assertEqual(df["n_bigint"].dtype, "Int64")
+        self.assertEqual(df["n_bigint"][0], 9007199254740993)
+        self.assertEqual(df["n_int"].dtype, "Int64")
+        self.assertEqual(df["n_int"][0], 3)
+        self.assertEqual(df["n_bit"].dtype, "boolean")
+        self.assertTrue(df["n_bit"][0])
+        self.assertEqual(df["n_float"].dtype, "float64")
+        self.assertTrue(df[["n_bigint", "n_int", "n_bit", "n_float"]].iloc[1].isna().all())
+
 
 class TestMSSQLHandlerODBC(unittest.TestCase):
     """Tests for MSSQL handler with ODBC connection"""
@@ -1185,6 +1222,57 @@ def __getitem__(self, idx):
             if "pyodbc" in sys.modules:
                 del sys.modules["pyodbc"]
 
+    def test_odbc_types_from_cursor_description(self):
+        """Test that ODBC column types come from cursor.description and NULLs don't turn integers into float"""
+        mock_pyodbc = MagicMock()
+        mock_conn = MagicMock()
+        mock_cursor = MagicMock()
+        mock_cursor.__enter__ = MagicMock(return_value=mock_cursor)
+        mock_cursor.__exit__ = MagicMock(return_value=None)
+
+        mock_cursor.fetchall.return_value = [
+            (None, None, None, None, None, None),
+            (9007199254740993, Decimal("10.50"), datetime.date(2026, 1, 2), datetime.time(12, 30), True, "x"),
+        ]
+        mock_cursor.description = [
+            ("n_bigint", int, None, 19, 19, 0, True),
+            ("n_decimal", Decimal, None, 19, 19, 4, True),
+            ("d_date", datetime.date, None, 10, 10, 0, True),
+            ("d_time", datetime.time, None, 16, 16, 7, True),
+            ("n_bit", bool, None, 1, 1, 0, True),
+            ("t_text", str, None, 10, 10, 0, True),
+        ]
+
+        sys.modules["pyodbc"] = mock_pyodbc
+
+        try:
+            handler = SqlServerHandler("mssql_odbc", connection_data=self.connection_data)
+            handler.connect = MagicMock(return_value=mock_conn)
+            handler.is_connected = True
+            mock_conn.cursor = MagicMock(return_value=mock_cursor)
+
+            response = handler.native_query("SELECT * FROM test")
+
+            self.assertEqual(
+                [col.type for col in response.columns],
+                [
+                    MYSQL_DATA_TYPE.INT,
+                    MYSQL_DATA_TYPE.DECIMAL,
+                    MYSQL_DATA_TYPE.DATE,
+                    MYSQL_DATA_TYPE.TIME,
+                    MYSQL_DATA_TYPE.TINYINT,
+                    MYSQL_DATA_TYPE.TEXT,
+                ],
+            )
+            df = response.data_frame
+            self.assertEqual(df["n_bigint"].dtype, "Int64")
+            self.assertEqual(df["n_bigint"][1], 9007199254740993)
+            self.assertEqual(df["n_bit"].dtype, "boolean")
+            self.assertTrue(df.iloc[0].isna().all())
+        finally:
+            if "pyodbc" in sys.modules:
+                del sys.modules["pyodbc"]
+
 
 if __name__ == "__main__":
     unittest.main()