Skip to content

Commit 0794565

Browse files
peco-engineer-bot[bot]peco-engineer-bot[bot]mehmetcankurtIsaac
authored
Queries returning duplicate column names fail with ArrowInvalid: Can't unify schema with duplicate field names since 4.2.6 (#989) (#995)
* Queries returning duplicate column names fail with `ArrowInvalid: Can't unify schema with duplicate field names` since 4.2.6 (#989) Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com> * ai: apply changes for #995 (1 review thread) Addresses: - #4341406782 at src/databricks/sql/utils.py:940 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> * ai: apply changes for #995 (1 review thread) Addresses: - #4341457400 at tests/e2e/test_driver.py:456 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> * Skip duplicate-column concat test when pyarrow is not installed The non-pyarrow unit test jobs import pyarrow as None, so the new test failed with AttributeError instead of being skipped like its neighbors. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Mehmet Can Kurt <mehmet.cankurt@databricks.com> --------- Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> Signed-off-by: Mehmet Can Kurt <mehmet.cankurt@databricks.com> Co-authored-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com> Co-authored-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> Co-authored-by: Mehmet Can Kurt <mehmet.cankurt@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 92e4e4b commit 0794565

3 files changed

Lines changed: 58 additions & 0 deletions

File tree

‎src/databricks/sql/utils.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -926,7 +926,25 @@ def _concat_arrow_tables(table_chunks: List["pyarrow.Table"]) -> "pyarrow.Table"
926926
bundled pyarrow -- e.g. DBR 13.3 (12.x) / 14.3 -- where ``promote_options``
927927
raises ``TypeError: unexpected keyword argument``. Fall back to the legacy
928928
``promote=True`` there (equivalent to ``promote_options="default"``).
929+
930+
Schema promotion triggers Arrow's schema unification, which rejects
931+
duplicate field names (``ArrowInvalid: Can't unify schema with duplicate
932+
field names``) even when no promotion is actually needed -- breaking valid
933+
SQL like ``SELECT 1 AS a, 2 AS a`` or ``SELECT * FROM a JOIN b`` (GH #989).
934+
Promotion is only needed when chunk schemas actually differ (e.g. the
935+
nullability case from GH #418), so when every chunk shares the first
936+
chunk's schema we concatenate without promotion.
937+
938+
Known limitation: a result set that has *both* duplicate column names *and*
939+
chunks whose schemas genuinely differ falls through to the promotion path
940+
and still re-raises ``ArrowInvalid: Can't unify schema with duplicate field
941+
names``. Arrow cannot unify a schema with duplicate field names, so there is
942+
no promotion-based fix for that combination without renaming fields; it is a
943+
narrow, pre-existing edge left unhandled here.
929944
"""
945+
if all(table.schema.equals(table_chunks[0].schema) for table in table_chunks[1:]):
946+
return pyarrow.concat_tables(table_chunks)
947+
930948
try:
931949
return pyarrow.concat_tables(table_chunks, promote_options="default")
932950
except TypeError:

‎tests/e2e/test_driver.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -453,6 +453,29 @@ def test_create_table_will_return_empty_result_set(self, extra_params):
453453
finally:
454454
cursor.execute("DROP TABLE IF EXISTS {}".format(table_name))
455455

456+
@skipUnless(pysql_supports_arrow(), "arrow test needs arrow support")
457+
def test_duplicate_column_names_survive_arrow_fetch(self):
458+
"""Queries whose result has two columns with the same name must survive
459+
the Arrow fetch path (concat_table_chunks). Regressed in 4.2.6 by the
460+
promote_options change (GH #989); 4.2.5 returned Row(a=1, a=2)."""
461+
with self.cursor({}) as cursor:
462+
cursor.execute("SELECT 1 AS a, 2 AS a")
463+
arrow_table = cursor.fetchall_arrow()
464+
assert arrow_table.column_names == ["a", "a"]
465+
assert arrow_table.column(0).to_pylist() == [1]
466+
assert arrow_table.column(1).to_pylist() == [2]
467+
468+
with self.cursor({}) as cursor:
469+
cursor.execute("SELECT 1 AS a, 2 AS a")
470+
rows = cursor.fetchall()
471+
assert len(rows) == 1
472+
assert list(rows[0]) == [1, 2]
473+
474+
with self.cursor({}) as cursor:
475+
cursor.execute("SELECT 1 AS a, 2 AS a")
476+
row = cursor.fetchone()
477+
assert list(row) == [1, 2]
478+
456479
def test_dml_rowcount(self):
457480
"""cursor.rowcount reports the affected-row count for DML
458481
(INSERT/UPDATE/DELETE) instead of the hardcoded -1, and resets

‎tests/unit/test_util.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,23 @@ def old_concat(tables, promote_options=None, **kwargs):
174174
result_table = concat_table_chunks([arrow_table1, arrow_table2])
175175
assert result_table.column("col1").to_pylist() == [1, 2, 3, 4]
176176

177+
@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
178+
def test_concat_table_chunks_duplicate_column_names(self):
179+
# Duplicate column names are valid SQL (e.g. SELECT 1 AS a, 2 AS a).
180+
# Schema promotion rejects them with ArrowInvalid, so when chunk
181+
# schemas match we must concatenate without promotion (GH #989).
182+
arrow_table1 = pyarrow.Table.from_arrays(
183+
[pyarrow.array([1]), pyarrow.array([2])], names=["a", "a"]
184+
)
185+
arrow_table2 = pyarrow.Table.from_arrays(
186+
[pyarrow.array([3]), pyarrow.array([4])], names=["a", "a"]
187+
)
188+
189+
result_table = concat_table_chunks([arrow_table1, arrow_table2])
190+
assert result_table.column_names == ["a", "a"]
191+
assert result_table.column(0).to_pylist() == [1, 3]
192+
assert result_table.column(1).to_pylist() == [2, 4]
193+
177194
def test_concat_table_chunks_empty(self):
178195
result_table = concat_table_chunks([])
179196
assert result_table == []

0 commit comments

Comments
 (0)