[PECOBLR-361] convert column table to arrow if arrow present (#551)

Signed-off-by: Sai Shree Pradhan <[email protected]>

Author: shivam2680

src/databricks/sql/client.py

@@ -1433,9 +1433,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):

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