databricks · shivam2680 · May 16, 2025 · Apr 24, 2025 · May 8, 2025
diff --git 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 PyArrow Table
+        # Valid only for metadata commands result set
+        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
@@ -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