From d8bdd4a169c13baecef618484094a6a80ccf4f33 Mon Sep 17 00:00:00 2001
From: Shivam Raj <shivam.raj@databricks.com>
Date: Thu, 24 Apr 2025 10:49:27 +0530
Subject: [PATCH] convert column table to arrow if arrow present

---
 src/databricks/sql/client.py | 15 ++++++++++++++-
 tests/e2e/test_driver.py     |  7 +++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index ea901c3a..48405421 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -1415,9 +1415,22 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
-            results = pyarrow.concat_tables([results, partial_results])
+            if isinstance(results, ColumnTable) and isinstance(
+                partial_results, ColumnTable
+            ):
+                results = self.merge_columnar(results, partial_results)
+            else:
+                results = pyarrow.concat_tables([results, partial_results])
             self._next_row_index += partial_results.num_rows
 
+        # If PyArrow is installed and we have a ColumnTable result,
+        # convert it to a PyArrow Table for consistency with pre-3.5.0 behavior
+        if isinstance(results, ColumnTable) and pyarrow:
+            data = {
+                name: col
+                for name, col in zip(results.column_names, results.column_table)
+            }
+            return pyarrow.Table.from_pydict(data)
         return results
 
     def fetchall_columnar(self):
diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py
index 8c0a4a5a..cfd56140 100644
--- a/tests/e2e/test_driver.py
+++ b/tests/e2e/test_driver.py
@@ -801,6 +801,13 @@ def test_decimal_not_returned_as_strings_arrow(self):
                 decimal_type = arrow_df.field(0).type
                 assert pyarrow.types.is_decimal(decimal_type)
 
+    @skipUnless(pysql_supports_arrow(), "arrow test needs arrow support")
+    def test_catalogs_returns_arrow_table(self):
+        with self.cursor() as cursor:
+            cursor.catalogs()
+            results = cursor.fetchall_arrow()
+            assert isinstance(results, pyarrow.Table)
+
     def test_close_connection_closes_cursors(self):
 
         from databricks.sql.thrift_api.TCLIService import ttypes