From d8bdd4a169c13baecef618484094a6a80ccf4f33 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 24 Apr 2025 10:49:27 +0530 Subject: [PATCH 1/2] 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 ea901c3ae..484054213 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 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 From 9503b7099352307ec6817002f25237568069cbd0 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 8 May 2025 15:19:35 +0530 Subject: [PATCH 2/2] modified comment --- src/databricks/sql/client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 484054213..f24a6584a 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1423,8 +1423,8 @@ def fetchall_arrow(self) -> "pyarrow.Table": 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 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 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