diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index ea901c3ae..f24a6584a 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 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 index 8c0a4a5a3..cfd561400 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
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies: