From 412e7639209030db74a20a8222505931849cf9e2 Mon Sep 17 00:00:00 2001 From: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:24:41 -0600 Subject: [PATCH 1/3] fix(databricks): limiting pyarrow version to meet requirements Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> --- Makefile | 5 +- pyproject.toml | 11 ++- sqlmesh/core/engine_adapter/databricks.py | 30 +++++++- tests/core/engine_adapter/test_databricks.py | 73 ++++++++++++++++++++ 4 files changed, 116 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 6357d06351..ee5acbb5d1 100644 --- a/Makefile +++ b/Makefile @@ -238,7 +238,10 @@ bigquery-test: guard-BIGQUERY_KEYFILE engine-bigquery-install pytest -n auto -m "bigquery" --reruns 3 --junitxml=test-results/junit-bigquery.xml databricks-test: guard-DATABRICKS_CATALOG guard-DATABRICKS_SERVER_HOSTNAME guard-DATABRICKS_HTTP_PATH guard-DATABRICKS_CONNECT_VERSION engine-databricks-install - $(PIP) install 'databricks-connect==${DATABRICKS_CONNECT_VERSION}' + # Connect 15.3.x requires numpy<2 and pyarrow>=4 (no upper bound). Cap PyArrow + # in the same resolve that downgrades NumPy, or PyArrow 25+ is left installed + # and cannot be imported. + $(PIP) install 'databricks-connect==${DATABRICKS_CONNECT_VERSION}' 'pyarrow<25' pytest -n auto -m "databricks" --reruns 3 --junitxml=test-results/junit-databricks.xml redshift-test: guard-REDSHIFT_HOST guard-REDSHIFT_USER guard-REDSHIFT_PASSWORD guard-REDSHIFT_DATABASE engine-redshift-install diff --git a/pyproject.toml b/pyproject.toml index 9b28dc7a08..dcd46b0e53 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,16 @@ bigquery = [ # pinned an older SQLGlot which is incompatible with SQLMesh bigframes = ["bigframes>=1.32.0"] clickhouse = ["clickhouse-connect"] -databricks = ["databricks-sql-connector[pyarrow]>=4.2.6"] +# PyArrow 25+ requires NumPy >= 2. install-dev currently resolves PyArrow 26 +# with NumPy 2.5. databricks-connect 15.3.6 then requires numpy<2 and +# pyarrow>=4 with no upper bound, so it replaces NumPy with 1.26.4 and leaves +# PyArrow 26 installed. PyArrow fails to import, the SQL connector disables +# Arrow, and fetchall_arrow() returns a ColumnTable with no to_pandas(). +# PyArrow 24.x imports on both NumPy 1.26 and NumPy 2. +databricks = [ + "databricks-sql-connector[pyarrow]>=4.2.6", + "pyarrow<25", +] dev = [ "agate", "beautifulsoup4", diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py index 8d47cdfd61..7f701ff353 100644 --- a/sqlmesh/core/engine_adapter/databricks.py +++ b/sqlmesh/core/engine_adapter/databricks.py @@ -31,6 +31,34 @@ logger = logging.getLogger(__name__) +def _to_pandas(table: t.Any) -> pd.DataFrame: + """Convert a databricks-sql-connector fetchall_arrow() result to pandas. + + Connector 4.x makes PyArrow optional. When PyArrow cannot be imported (PyArrow 25 + requires NumPy >= 2), fetchall_arrow() returns a ColumnTable with ``column_names`` + and ``column_table`` instead of a PyArrow table. + """ + import pandas as pd + + to_pandas = getattr(table, "to_pandas", None) + if callable(to_pandas): + return to_pandas() + + column_names = getattr(table, "column_names", None) + column_table = getattr(table, "column_table", None) + if column_names is None or column_table is None: + raise SQLMeshError( + f"Unexpected Databricks fetchall_arrow() result {type(table).__name__!r} " + "with no to_pandas(). PyArrow may be missing or incompatible with the " + "installed NumPy (PyArrow 25+ requires NumPy >= 2)." + ) + + # Build from rows so duplicate names are kept. A dict would collapse them, + # which is not what PyArrow's to_pandas() does. + columns = [list(values) for values in column_table] + return pd.DataFrame(list(zip(*columns)), columns=list(column_names)) + + def _query_tags( query_tags: t.Optional[t.Union[exp.Expr, str, int, float, bool]], ) -> t.Optional[t.Dict[str, t.Optional[str]]]: @@ -277,7 +305,7 @@ def _fetch_native_df( query, quote_identifiers=quote_identifiers ) self.execute(query) - return self.cursor.fetchall_arrow().to_pandas() + return _to_pandas(self.cursor.fetchall_arrow()) def fetchdf( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py index 65d8e0388f..d591b7278c 100644 --- a/tests/core/engine_adapter/test_databricks.py +++ b/tests/core/engine_adapter/test_databricks.py @@ -837,3 +837,76 @@ def test_databricks_connect_cluster_id_not_overridden_by_falsy_serverless( token="mytoken", cluster_id="0123-456789-mycluster", ) + + +def test_fetch_native_df_uses_arrow_to_pandas(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + arrow_table = pd.DataFrame({"name": ["orders"]}) + fetchall_arrow = adapter.cursor.fetchall_arrow + fetchall_arrow.return_value.to_pandas.return_value = arrow_table + + df = adapter._fetch_native_df("SELECT 1") + + assert df.to_dict(orient="records") == [{"name": "orders"}] + fetchall_arrow.return_value.to_pandas.assert_called_once_with() + + +def test_fetch_native_df_column_table_without_to_pandas(make_mocked_engine_adapter: t.Callable): + """PyArrow 25 fails to import on NumPy < 2, so the connector returns a ColumnTable.""" + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + + class _ColumnTable: + column_names = ["name", "schema", "catalog", "type"] + column_table = [("orders",), ("analytics",), ("main",), ("table",)] + + adapter.cursor.fetchall_arrow.return_value = _ColumnTable() + + df = adapter._fetch_native_df("SELECT 1") + + assert df.to_dict(orient="records") == [ + {"name": "orders", "schema": "analytics", "catalog": "main", "type": "table"} + ] + + +def test_fetch_native_df_empty_column_table(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + + class _ColumnTable: + column_names = ["name", "schema", "catalog", "type"] + column_table = [(), (), (), ()] + + adapter.cursor.fetchall_arrow.return_value = _ColumnTable() + + df = adapter._fetch_native_df("SELECT 1") + + assert list(df.columns) == ["name", "schema", "catalog", "type"] + assert df.empty + + +def test_fetch_native_df_column_table_keeps_duplicate_columns( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + + class _ColumnTable: + column_names = ["a", "a"] + column_table = [("x",), ("y",)] + + adapter.cursor.fetchall_arrow.return_value = _ColumnTable() + + df = adapter._fetch_native_df("SELECT 1") + + assert list(df.columns) == ["a", "a"] + assert df.iloc[0].tolist() == ["x", "y"] + + +def test_fetch_native_df_rejects_unknown_arrow_result(make_mocked_engine_adapter: t.Callable): + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + + class _Weird: + pass + + adapter.cursor.fetchall_arrow.return_value = _Weird() + + with pytest.raises(SQLMeshError, match="Unexpected Databricks fetchall_arrow"): + adapter._fetch_native_df("SELECT 1") From e1fecc49c06bde8587a5b8f778940dbc6c703cbc Mon Sep 17 00:00:00 2001 From: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:35:06 -0600 Subject: [PATCH 2/3] removing unnecessary fix Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> --- sqlmesh/core/engine_adapter/databricks.py | 30 +------- tests/core/engine_adapter/test_databricks.py | 73 -------------------- 2 files changed, 1 insertion(+), 102 deletions(-) diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py index 7f701ff353..8d47cdfd61 100644 --- a/sqlmesh/core/engine_adapter/databricks.py +++ b/sqlmesh/core/engine_adapter/databricks.py @@ -31,34 +31,6 @@ logger = logging.getLogger(__name__) -def _to_pandas(table: t.Any) -> pd.DataFrame: - """Convert a databricks-sql-connector fetchall_arrow() result to pandas. - - Connector 4.x makes PyArrow optional. When PyArrow cannot be imported (PyArrow 25 - requires NumPy >= 2), fetchall_arrow() returns a ColumnTable with ``column_names`` - and ``column_table`` instead of a PyArrow table. - """ - import pandas as pd - - to_pandas = getattr(table, "to_pandas", None) - if callable(to_pandas): - return to_pandas() - - column_names = getattr(table, "column_names", None) - column_table = getattr(table, "column_table", None) - if column_names is None or column_table is None: - raise SQLMeshError( - f"Unexpected Databricks fetchall_arrow() result {type(table).__name__!r} " - "with no to_pandas(). PyArrow may be missing or incompatible with the " - "installed NumPy (PyArrow 25+ requires NumPy >= 2)." - ) - - # Build from rows so duplicate names are kept. A dict would collapse them, - # which is not what PyArrow's to_pandas() does. - columns = [list(values) for values in column_table] - return pd.DataFrame(list(zip(*columns)), columns=list(column_names)) - - def _query_tags( query_tags: t.Optional[t.Union[exp.Expr, str, int, float, bool]], ) -> t.Optional[t.Dict[str, t.Optional[str]]]: @@ -305,7 +277,7 @@ def _fetch_native_df( query, quote_identifiers=quote_identifiers ) self.execute(query) - return _to_pandas(self.cursor.fetchall_arrow()) + return self.cursor.fetchall_arrow().to_pandas() def fetchdf( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py index d591b7278c..65d8e0388f 100644 --- a/tests/core/engine_adapter/test_databricks.py +++ b/tests/core/engine_adapter/test_databricks.py @@ -837,76 +837,3 @@ def test_databricks_connect_cluster_id_not_overridden_by_falsy_serverless( token="mytoken", cluster_id="0123-456789-mycluster", ) - - -def test_fetch_native_df_uses_arrow_to_pandas(make_mocked_engine_adapter: t.Callable): - adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") - arrow_table = pd.DataFrame({"name": ["orders"]}) - fetchall_arrow = adapter.cursor.fetchall_arrow - fetchall_arrow.return_value.to_pandas.return_value = arrow_table - - df = adapter._fetch_native_df("SELECT 1") - - assert df.to_dict(orient="records") == [{"name": "orders"}] - fetchall_arrow.return_value.to_pandas.assert_called_once_with() - - -def test_fetch_native_df_column_table_without_to_pandas(make_mocked_engine_adapter: t.Callable): - """PyArrow 25 fails to import on NumPy < 2, so the connector returns a ColumnTable.""" - adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") - - class _ColumnTable: - column_names = ["name", "schema", "catalog", "type"] - column_table = [("orders",), ("analytics",), ("main",), ("table",)] - - adapter.cursor.fetchall_arrow.return_value = _ColumnTable() - - df = adapter._fetch_native_df("SELECT 1") - - assert df.to_dict(orient="records") == [ - {"name": "orders", "schema": "analytics", "catalog": "main", "type": "table"} - ] - - -def test_fetch_native_df_empty_column_table(make_mocked_engine_adapter: t.Callable): - adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") - - class _ColumnTable: - column_names = ["name", "schema", "catalog", "type"] - column_table = [(), (), (), ()] - - adapter.cursor.fetchall_arrow.return_value = _ColumnTable() - - df = adapter._fetch_native_df("SELECT 1") - - assert list(df.columns) == ["name", "schema", "catalog", "type"] - assert df.empty - - -def test_fetch_native_df_column_table_keeps_duplicate_columns( - make_mocked_engine_adapter: t.Callable, -): - adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") - - class _ColumnTable: - column_names = ["a", "a"] - column_table = [("x",), ("y",)] - - adapter.cursor.fetchall_arrow.return_value = _ColumnTable() - - df = adapter._fetch_native_df("SELECT 1") - - assert list(df.columns) == ["a", "a"] - assert df.iloc[0].tolist() == ["x", "y"] - - -def test_fetch_native_df_rejects_unknown_arrow_result(make_mocked_engine_adapter: t.Callable): - adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") - - class _Weird: - pass - - adapter.cursor.fetchall_arrow.return_value = _Weird() - - with pytest.raises(SQLMeshError, match="Unexpected Databricks fetchall_arrow"): - adapter._fetch_native_df("SELECT 1") From 9be4b247f3c628faf9cb81a17aab851907d02160 Mon Sep 17 00:00:00 2001 From: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:36:21 -0600 Subject: [PATCH 3/3] fix(databricks-ci): moving to ci install only Signed-off-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com> --- pyproject.toml | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dcd46b0e53..9b28dc7a08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,16 +51,7 @@ bigquery = [ # pinned an older SQLGlot which is incompatible with SQLMesh bigframes = ["bigframes>=1.32.0"] clickhouse = ["clickhouse-connect"] -# PyArrow 25+ requires NumPy >= 2. install-dev currently resolves PyArrow 26 -# with NumPy 2.5. databricks-connect 15.3.6 then requires numpy<2 and -# pyarrow>=4 with no upper bound, so it replaces NumPy with 1.26.4 and leaves -# PyArrow 26 installed. PyArrow fails to import, the SQL connector disables -# Arrow, and fetchall_arrow() returns a ColumnTable with no to_pandas(). -# PyArrow 24.x imports on both NumPy 1.26 and NumPy 2. -databricks = [ - "databricks-sql-connector[pyarrow]>=4.2.6", - "pyarrow<25", -] +databricks = ["databricks-sql-connector[pyarrow]>=4.2.6"] dev = [ "agate", "beautifulsoup4",