diff --git a/CHANGELOG.md b/CHANGELOG.md index a78c3baf9..fcccde626 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Fix fetching results that contain duplicate column names, which failed with `ArrowInvalid: Can't unify schema with duplicate field names` since 4.2.6. # 4.7.0 (2026-10-20) - Upgrade Databricks SQL Kernel to 1.2.0. diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index dd3525547..1869ad903 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -927,6 +927,11 @@ def _concat_arrow_tables(table_chunks: List["pyarrow.Table"]) -> "pyarrow.Table" raises ``TypeError: unexpected keyword argument``. Fall back to the legacy ``promote=True`` there (equivalent to ``promote_options="default"``). """ + first_schema = table_chunks[0].schema + if all(chunk.schema.equals(first_schema) for chunk in table_chunks[1:]): + # Nothing to promote. Schema unification rejects duplicate column names + # (e.g. SELECT 1 AS a, 2 AS a), so skip it when the chunks already agree. + return pyarrow.concat_tables(table_chunks) try: return pyarrow.concat_tables(table_chunks, promote_options="default") except TypeError: diff --git a/tests/unit/test_util.py b/tests/unit/test_util.py index 94a96b54a..33b2605aa 100644 --- a/tests/unit/test_util.py +++ b/tests/unit/test_util.py @@ -152,6 +152,38 @@ def test_concat_table_chunks_arrow_table(self): assert result_table.column("col1").to_pylist() == [1, 2, 3, 4] assert result_table.column("col2").to_pylist() == [5, 6, 7, 8] + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_concat_table_chunks_arrow_table_duplicate_column_names(self): + arrow_table1 = pyarrow.Table.from_arrays( + [pyarrow.array([1]), pyarrow.array([2])], names=["a", "a"] + ) + arrow_table2 = pyarrow.Table.from_arrays( + [pyarrow.array([3]), pyarrow.array([4])], names=["a", "a"] + ) + + result_table = concat_table_chunks([arrow_table1, arrow_table2]) + + assert result_table.column_names == ["a", "a"] + assert result_table.column(0).to_pylist() == [1, 3] + assert result_table.column(1).to_pylist() == [2, 4] + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_concat_table_chunks_arrow_table_promotes_nullability(self): + nullable = pyarrow.Table.from_arrays( + [pyarrow.array(["x"])], + schema=pyarrow.schema([pyarrow.field("col1", pyarrow.string())]), + ) + not_null = pyarrow.Table.from_arrays( + [pyarrow.array(["y"])], + schema=pyarrow.schema( + [pyarrow.field("col1", pyarrow.string(), nullable=False)] + ), + ) + + result_table = concat_table_chunks([nullable, not_null]) + + assert result_table.column("col1").to_pylist() == ["x", "y"] + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") def test_concat_table_chunks_arrow_old_pyarrow_fallback(self, monkeypatch): # pyarrow < 14 has no `promote_options` kwarg; a base connector install @@ -159,7 +191,12 @@ def test_concat_table_chunks_arrow_old_pyarrow_fallback(self, monkeypatch): # concat_table_chunks must fall back to the legacy `promote=True` API # rather than raise TypeError (regression: DBR LTS install check). arrow_table1 = pyarrow.Table.from_pydict({"col1": [1, 2]}) - arrow_table2 = pyarrow.Table.from_pydict({"col1": [3, 4]}) + arrow_table2 = pyarrow.Table.from_arrays( + [pyarrow.array([3, 4])], + schema=pyarrow.schema( + [pyarrow.field("col1", pyarrow.int64(), nullable=False)] + ), + ) real_concat = pyarrow.concat_tables