From c6c7f57f3da72381ac5fabbbc64c9d456673443c Mon Sep 17 00:00:00 2001 From: Maha Rana Yadavalli <271375718+maharanay22@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:37:04 -0400 Subject: [PATCH] Keep SQL Server integer columns exact when they contain NULLs The SQL Server handler built the result DataFrame without dtypes, so a NULL in an int, bigint or bit column made pandas fall back to float64 or object. Large IDs lost precision (9007199254740993 came back as 9007199254740992.0), small values turned into floats (3 -> 3.0) and the columns were reported as FLOAT or DOUBLE instead of INT. Rebuild integer and bit columns that contain NULLs with the pandas nullable Int64 and boolean dtypes. For ODBC connections, also take the column type from cursor.description, which pyodbc fills with the Python type of the column, so DECIMAL, DATE and TIME columns are no longer reported as TEXT or DATETIME. --- .../handlers/mssql_handler/mssql_handler.py | 53 ++++++++++- tests/unit/handlers/test_mssql.py | 88 +++++++++++++++++++ 2 files changed, 138 insertions(+), 3 deletions(-) diff --git a/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py b/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py index 7b6e42fff3..c537f5a853 100644 --- a/mindsdb/integrations/handlers/mssql_handler/mssql_handler.py +++ b/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): diff --git a/tests/unit/handlers/test_mssql.py b/tests/unit/handlers/test_mssql.py index d7024d5135..38192a7eb8 100644 --- a/tests/unit/handlers/test_mssql.py +++ b/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()