reduce diff

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

Author: varun-edachali-dbx

src/databricks/sql/result_set.py

@@ -36,8 +36,8 @@ class ResultSet(ABC):
 
     def __init__(
         self,
-        connection: "Connection",
-        backend: "DatabricksClient",
+        connection: Connection,
+        backend: DatabricksClient,
         arraysize: int,
         buffer_size_bytes: int,
         command_id: CommandId,
@@ -54,8 +54,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 backend 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
@@ -190,9 +190,9 @@ class ThriftResultSet(ResultSet):
 
     def __init__(
         self,
-        connection: "Connection",
-        execute_response: "ExecuteResponse",
-        thrift_client: "ThriftDatabricksClient",
+        connection: Connection,
+        execute_response: ExecuteResponse,
+        thrift_client: ThriftDatabricksClient,
         session_id_hex: Optional[str],
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
@@ -319,6 +319,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
 
@@ -329,11 +330,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)
 
     def fetchmany_columnar(self, size: int):
         """
@@ -364,7 +365,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()
@@ -373,7 +374,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
@@ -384,7 +385,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)
 
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""

src/databricks/sql/session.py

@@ -150,7 +150,7 @@ def session_id(self) -> SessionId:
         return self._session_id
 
     @property
-    def guid(self):
+    def guid(self) -> Any:
         """Get the raw session ID (backend-specific)"""
         return self._session_id.guid
 
@@ -161,7 +161,7 @@ def guid_hex(self) -> str:
 
     def close(self) -> None:
         """Close the underlying session."""
-        logger.info(f"Closing session {self.guid_hex}")
+        logger.info("Closing session %s", self.guid_hex)
         if not self.is_open:
             logger.debug("Session appears to have been closed already")
             return
@@ -174,13 +174,13 @@ def close(self) -> None:
         except DatabaseError as e:
             if "Invalid SessionHandle" in str(e):
                 logger.warning(
-                    f"Attempted to close session that was already closed: {e}"
+                    "Attempted to close session that was already closed: %s", e
                 )
             else:
                 logger.warning(
-                    f"Attempt to close session raised an exception at the server: {e}"
+                    "Attempt to close session raised an exception at the server: %s", e
                 )
         except Exception as e:
-            logger.error(f"Attempt to close session raised a local exception: {e}")
+            logger.error("Attempt to close session raised a local exception: %s", e)
 
         self.is_open = False