nits: string literalrs around type defs, naming, excess changes

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -602,7 +602,7 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             session_id = SessionId.from_thrift_handle(
                 response.sessionHandle, properties
             )
-            self._session_id_hex = session_id.guid_hex
+            self._session_id_hex = session_id.hex_guid
             return session_id
         except:
             self._transport.close()
@@ -832,7 +832,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
         return execute_response, is_direct_results
 
     def get_execution_result(
-        self, command_id: CommandId, cursor: "Cursor"
+        self, command_id: CommandId, cursor: Cursor
     ) -> "ResultSet":
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
@@ -1044,8 +1044,8 @@ def get_catalogs(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: "Cursor",
-    ) -> "ResultSet":
+        cursor: Cursor,
+    ) -> ResultSet:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1087,7 +1087,7 @@ def get_schemas(
         cursor: Cursor,
         catalog_name=None,
         schema_name=None,
-    ) -> "ResultSet":
+    ) -> ResultSet:
         from databricks.sql.result_set import ThriftResultSet
 
         thrift_handle = session_id.to_thrift_handle()
@@ -1135,7 +1135,7 @@ def get_tables(
         schema_name=None,
         table_name=None,
         table_types=None,
-    ) -> "ResultSet":
+    ) -> ResultSet:
         from databricks.sql.result_set import ThriftResultSet
 
         thrift_handle = session_id.to_thrift_handle()
@@ -1185,7 +1185,7 @@ def get_columns(
         schema_name=None,
         table_name=None,
         column_name=None,
-    ) -> "ResultSet":
+    ) -> ResultSet:
         from databricks.sql.result_set import ThriftResultSet
 
         thrift_handle = session_id.to_thrift_handle()

src/databricks/sql/backend/types.py

@@ -160,7 +160,7 @@ def __str__(self) -> str:
                 if isinstance(self.secret, bytes)
                 else str(self.secret)
             )
-            return f"{self.guid_hex}|{secret_hex}"
+            return f"{self.hex_guid}|{secret_hex}"
         return str(self.guid)
 
     @classmethod
@@ -239,7 +239,7 @@ def to_sea_session_id(self):
         return self.guid
 
     @property
-    def guid_hex(self) -> str:
+    def hex_guid(self) -> str:
         """
         Get a hexadecimal string representation of the session ID.
 

src/databricks/sql/result_set.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import List, Optional, Any, TYPE_CHECKING
 
@@ -33,8 +35,8 @@ class ResultSet(ABC):
 
     def __init__(
         self,
-        connection: "Connection",
-        backend: "DatabricksClient",
+        connection: Connection,
+        backend: DatabricksClient,
         arraysize: int,
         buffer_size_bytes: int,
         command_id: CommandId,
@@ -51,8 +53,8 @@ def __init__(
         A ResultSet manages the results of a single command.
 
         Parameters:
-            :param connection: The parent connection
-            :param backend: The backend client
+            :param connection: The parent connection that was used to execute this command
+            :param backend: The specialised backend client to be invoked in the fetch phase
             :param arraysize: The max number of rows to fetch at a time (PEP-249)
             :param buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch
             :param command_id: The command ID
@@ -156,9 +158,9 @@ class ThriftResultSet(ResultSet):
 
     def __init__(
         self,
-        connection: "Connection",
-        execute_response: "ExecuteResponse",
-        thrift_client: "ThriftDatabricksClient",
+        connection: Connection,
+        execute_response: ExecuteResponse,
+        thrift_client: ThriftDatabricksClient,
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
@@ -314,6 +316,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         if size < 0:
             raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
         results = self.results.next_n_rows(size)
+        partial_result_chunks = [results]
         n_remaining_rows = size - results.num_rows
         self._next_row_index += results.num_rows
 
@@ -324,11 +327,11 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
-            results = pyarrow.concat_tables([results, partial_results])
+            partial_result_chunks.append(partial_results)
             n_remaining_rows -= partial_results.num_rows
             self._next_row_index += partial_results.num_rows
 
-        return results
+        return pyarrow.concat_tables(partial_result_chunks, use_threads=True)
 
     def fetchmany_columnar(self, size: int):
         """
@@ -359,7 +362,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
-
+        partial_result_chunks = [results]
         while not self.has_been_closed_server_side and self.is_direct_results:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
@@ -368,7 +371,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
             ):
                 results = self.merge_columnar(results, partial_results)
             else:
-                results = pyarrow.concat_tables([results, partial_results])
+                partial_result_chunks.append(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
@@ -379,7 +382,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
                 for name, col in zip(results.column_names, results.column_table)
             }
             return pyarrow.Table.from_pydict(data)
-        return results
+        return pyarrow.concat_tables(partial_result_chunks, use_threads=True)
 
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
@@ -452,9 +455,9 @@ class SeaResultSet(ResultSet):
 
     def __init__(
         self,
-        connection: "Connection",
-        execute_response: "ExecuteResponse",
-        sea_client: "SeaDatabricksClient",
+        connection: Connection,
+        execute_response: ExecuteResponse,
+        sea_client: SeaDatabricksClient,
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         result_data=None,

src/databricks/sql/session.py

@@ -156,7 +156,7 @@ def guid(self):
     @property
     def guid_hex(self) -> str:
         """Get the session ID in hex format"""
-        return self._session_id.guid_hex
+        return self._session_id.hex_guid
 
     def close(self) -> None:
         """Close the underlying session."""