More refractor

Author: jprakash-db

src/databricks/sql/client.py

@@ -9,7 +9,6 @@
 import requests
 import json
 import os
-import decimal
 from uuid import UUID
 
 from databricks.sql import __version__
@@ -31,8 +30,6 @@
     transform_paramstyle,
     ColumnTable,
     ColumnQueue,
-    concat_chunked_tables,
-    merge_columnar,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -1402,7 +1399,7 @@ def _convert_columnar_table(self, table: ColumnTable):
             result.append(ResultRow(*curr_row))
 
         return result
-   
+
     def _convert_arrow_table(self, table: "pyarrow.Table"):
 
         column_names = [c[0] for c in self.description]
@@ -1411,7 +1408,7 @@ def _convert_arrow_table(self, table: "pyarrow.Table"):
         if self.connection.disable_pandas is True:
             columns_as_lists = [col.to_pylist() for col in table.itercolumns()]
             return [ResultRow(*row) for row in zip(*columns_as_lists)]
-        
+
         # Need to use nullable types, as otherwise type can change when there are missing values.
         # See https://arrow.apache.org/docs/python/pandas.html#nullable-types
         # NOTE: This api is epxerimental https://pandas.pydata.org/pandas-docs/stable/user_guide/integer_na.html
@@ -1445,7 +1442,7 @@ def _convert_arrow_table(self, table: "pyarrow.Table"):
     @property
     def rownumber(self):
         return self._next_row_index
-    
+
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows of a query result, returning a PyArrow table.
@@ -1494,18 +1491,18 @@ def fetchmany_columnar(self, size: int):
             self._next_row_index += partial_results.num_rows
 
         return results
-    
+
     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
-        
+
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             results.append(partial_results)
             self._next_row_index += partial_results.num_rows
-            
+
         return results.to_arrow_table()
 
     def fetchall_columnar(self):
@@ -1516,7 +1513,7 @@ def fetchall_columnar(self):
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
-            results = merge_columnar(results, partial_results)
+            results.append(partial_results)
             self._next_row_index += partial_results.num_rows
 
         return results

src/databricks/sql/cloudfetch/download_manager.py

@@ -84,8 +84,7 @@ def _schedule_downloads(self):
         """
         While download queue has a capacity, peek pending links and submit them to thread pool.
         """
-        # print("Schedule_downloads")
-        # logger.debug("ResultFileDownloadManager: schedule downloads")
+        logger.debug("ResultFileDownloadManager: schedule downloads")
         while (len(self._download_tasks) < self._max_download_threads) and (
             len(self._pending_links) > 0
         ):

src/databricks/sql/cloudfetch/downloader.py

@@ -91,17 +91,17 @@ def run(self) -> DownloadedFile:
         ResultSetDownloadHandler._validate_link(
             self.link, self.settings.link_expiry_buffer_secs
         )
-        
+
         with self._http_client.execute(
             method=HttpMethod.GET,
             url=self.link.fileLink,
             timeout=self.settings.download_timeout,
             verify=self._ssl_options.tls_verify,
             headers=self.link.httpHeaders
+            # TODO: Pass cert from `self._ssl_options`
         ) as response:
-
             response.raise_for_status()
-            
+
             # Save (and decompress if needed) the downloaded file
             compressed_data = response.content
             decompressed_data = (

src/databricks/sql/common/http.py

@@ -8,6 +8,7 @@
 from typing import Generator
 import logging
 import time
+
 logger = logging.getLogger(__name__)
 
 
@@ -70,10 +71,7 @@ def execute(
         logger.info("Executing HTTP request: %s with url: %s", method.value, url)
         response = None
         try:
-            start_time = time.time()
             response = self.session.request(method.value, url, **kwargs)
-            end_time = time.time()
-            # print(f"Downloaded file in {end_time - start_time} seconds")
             yield response
         except Exception as e:
             logger.error("Error executing HTTP request in DatabricksHttpClient: %s", e)

src/databricks/sql/thrift_backend.py

@@ -36,9 +36,6 @@
     RequestErrorInfo,
     NoRetryReason,
     ResultSetQueueFactory,
-    convert_arrow_based_set_to_arrow_table,
-    convert_decimals_in_arrow_table,
-    convert_column_based_set_to_arrow_table,
 )
 from databricks.sql.types import SSLOptions
 
@@ -633,23 +630,6 @@ def _poll_for_status(self, op_handle):
         )
         return self.make_request(self._client.GetOperationStatus, req)
 
-    def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, description):
-        if t_row_set.columns is not None:
-            (
-                arrow_table,
-                num_rows,
-            ) = convert_column_based_set_to_arrow_table(t_row_set.columns, description)
-        elif t_row_set.arrowBatches is not None:
-            (arrow_table, num_rows,) = convert_arrow_based_set_to_arrow_table(
-                t_row_set.arrowBatches, lz4_compressed, schema_bytes
-            )
-        else:
-            raise OperationalError(
-                "Unsupported TRowSet instance {}".format(t_row_set),
-                session_id_hex=self._session_id_hex,
-            )
-        return convert_decimals_in_arrow_table(arrow_table, description), num_rows
-
     def _get_metadata_resp(self, op_handle):
         req = ttypes.TGetResultSetMetadataReq(operationHandle=op_handle)
         return self.make_request(self._client.GetResultSetMetadata, req)