Skip to content

Keep Databricks integer columns exact when they contain NULLs - #12

Open
maharanay22 wants to merge 1 commit into
mindsdb:mainfrom
maharanay22:fix-databricks-nullable-int-types
Open

maharanay22 wants to merge 1 commit into
mindsdb:mainfrom
maharanay22:fix-databricks-nullable-int-types

Conversation

@maharanay22

Copy link
Copy Markdown

Description

DatabricksHandler.native_query() builds the result with pd.DataFrame(result, columns=...) straight from the cursor rows. When a BIGINT, INT, SMALLINT or TINYINT column contains a NULL, pandas stores the whole column as float64. That changes the data MindsDB returns:

  1. Values above 2**53 lose precision. A BIGINT ID of 9007199254740993 comes back as 9007199254740992.0, which is a different ID
  2. Small integers turn into floats, so a quantity of 3 comes back as 3.0

BOOLEAN columns with a NULL have a similar problem and come back as object

This change builds the DataFrame in a small helper, _make_data_frame(). After the normal construction, any column that pandas couldn't keep as int or bool is rebuilt from the raw values with the pandas nullable Int64 or boolean dtype, but only when every non-NULL value in it is an int (or every one is a bool). Columns without NULLs keep their current dtypes, and float, decimal, string and date columns are left alone. The PostgreSQL handler already returns nullable Int64 / boolean columns, so this matches what MindsDB gets from it. Columns are handled by position, so queries that return two columns with the same name still work

Fixes: no linked issue, because issues are turned off for this repository. The bug is described in full above

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Verification Process

To ensure the changes are working as expected:

  • Test Location: tests/unit/handlers/community_handlers/test_databricks.py, new test TestTableOperations::test_native_query_keeps_integers_with_nulls_exact
  • Verification Steps: run native_query() with a mocked connection whose cursor returns databricks.sql.types.Row rows for id BIGINT, amount DECIMAL(10,2) and qty INT, where the second row is all NULLs

Before (current main):

[{'id': 9007199254740992.0, 'amount': Decimal('10.50'), 'qty': 3.0}, {'id': nan, 'amount': None, 'qty': nan}] {'id': dtype('float64'), 'amount': dtype('O'), 'qty': dtype('float64')}

After:

[{'id': 9007199254740993, 'amount': Decimal('10.50'), 'qty': 3}, {'id': None, 'amount': None, 'qty': None}] {'id': Int64Dtype(), 'amount': dtype('O'), 'qty': Int64Dtype()}

Unit tests. The new test fails on main and passes with the fix, and the existing Databricks tests still pass:

# main
FAILED tests/unit/handlers/community_handlers/test_databricks.py::TestTableOperations::test_native_query_keeps_integers_with_nulls_exact
E       AssertionError: dtype('float64') != 'Int64'
1 failed, 39 passed

# this branch
40 passed

I didn't have a Databricks workspace to test against, so the check uses the connector's own Row type with a mocked connection. ruff check and ruff format --check (0.11.11) pass on the changed files. No dependencies changed. databricks_handler.py is stored with CRLF line endings in the repository, so the new lines use CRLF too and the diff only touches the lines that changed

Additional Media:

  • I have attached a brief loom video or screenshots showcasing the new functionality or change.

I didn't record a video or take screenshots. The before and after terminal output above shows the change instead

Checklist:

  • My code follows the style guidelines(PEP 8) of MindsDB.
  • I have appropriately commented on my code, especially in complex areas.
  • Necessary documentation updates are either made or tracked in issues.
  • Relevant unit and integration tests are updated or added.

No documentation change is needed

I used an AI assistant to help investigate this and draft the change and tests, and I reviewed and tested it myself

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant