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 pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

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:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy