[SC-110400] Enabling compression in Python SQL Connector (#49)

Signed-off-by: Mohit Singla <mohit.singla@databricks.com>
Co-authored-by: Moe Derakhshani <moe.derakhshani@databricks.com>

Author: mohitsingla-db

poetry.lock

@@ -13,6 +13,7 @@ python = "^3.7.1"
 thrift = "^0.13.0"
 pandas = "^1.3.0"
 pyarrow = "^9.0.0"
+lz4 = "^4.0.2"
 requests=">2.18.1"
 oauthlib=">=3.1.0"
 

pyproject.toml

@@ -152,6 +152,7 @@ def read(self) -> Optional[OAuthToken]:
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)
         self.disable_pandas = kwargs.get("_disable_pandas", False)
+        self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True)
 
         auth_provider = get_python_sql_connector_auth_provider(
             server_hostname, **kwargs
@@ -318,6 +319,7 @@ def execute(
             session_handle=self.connection._session_handle,
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
+            lz4_compression=self.connection.lz4_compression,
             cursor=self,
         )
         self.active_result_set = ResultSet(
@@ -614,6 +616,7 @@ def __init__(
         self.has_been_closed_server_side = execute_response.has_been_closed_server_side
         self.has_more_rows = execute_response.has_more_rows
         self.buffer_size_bytes = result_buffer_size_bytes
+        self.lz4_compressed = execute_response.lz4_compressed
         self.arraysize = arraysize
         self.thrift_backend = thrift_backend
         self.description = execute_response.description
@@ -642,6 +645,7 @@ def _fill_results_buffer(self):
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             expected_row_start_offset=self._next_row_index,
+            lz4_compressed=self.lz4_compressed,
             arrow_schema_bytes=self._arrow_schema_bytes,
             description=self.description,
         )

src/databricks/sql/client.py

@@ -4,6 +4,7 @@
 import math
 import time
 import threading
+import lz4.frame
 from ssl import CERT_NONE, CERT_REQUIRED, create_default_context
 
 import pyarrow
@@ -451,7 +452,7 @@ def open_session(self, session_configuration, catalog, schema):
                 initial_namespace = None
 
             open_session_req = ttypes.TOpenSessionReq(
-                client_protocol_i64=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V5,
+                client_protocol_i64=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V6,
                 client_protocol=None,
                 initialNamespace=initial_namespace,
                 canUseMultipleCatalogs=True,
@@ -507,7 +508,7 @@ def _poll_for_status(self, op_handle):
         )
         return self.make_request(self._client.GetOperationStatus, req)
 
-    def _create_arrow_table(self, t_row_set, schema_bytes, description):
+    def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, description):
         if t_row_set.columns is not None:
             (
                 arrow_table,
@@ -520,7 +521,7 @@ def _create_arrow_table(self, t_row_set, schema_bytes, description):
                 arrow_table,
                 num_rows,
             ) = ThriftBackend._convert_arrow_based_set_to_arrow_table(
-                t_row_set.arrowBatches, schema_bytes
+                t_row_set.arrowBatches, lz4_compressed, schema_bytes
             )
         else:
             raise OperationalError("Unsupported TRowSet instance {}".format(t_row_set))
@@ -545,13 +546,20 @@ def _convert_decimals_in_arrow_table(table, description):
         return table
 
     @staticmethod
-    def _convert_arrow_based_set_to_arrow_table(arrow_batches, schema_bytes):
+    def _convert_arrow_based_set_to_arrow_table(
+        arrow_batches, lz4_compressed, schema_bytes
+    ):
         ba = bytearray()
         ba += schema_bytes
         n_rows = 0
-        for arrow_batch in arrow_batches:
-            n_rows += arrow_batch.rowCount
-            ba += arrow_batch.batch
+        if lz4_compressed:
+            for arrow_batch in arrow_batches:
+                n_rows += arrow_batch.rowCount
+                ba += lz4.frame.decompress(arrow_batch.batch)
+        else:
+            for arrow_batch in arrow_batches:
+                n_rows += arrow_batch.rowCount
+                ba += arrow_batch.batch
         arrow_table = pyarrow.ipc.open_stream(ba).read_all()
         return arrow_table, n_rows
 
@@ -708,7 +716,6 @@ def _results_message_to_execute_response(self, resp, operation_state):
                     ]
                 )
             )
-
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
         has_more_rows = (
@@ -725,12 +732,16 @@ def _results_message_to_execute_response(self, resp, operation_state):
             .serialize()
             .to_pybytes()
         )
-
+        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         if direct_results and direct_results.resultSet:
             assert direct_results.resultSet.results.startRowOffset == 0
             assert direct_results.resultSetMetadata
+
             arrow_results, n_rows = self._create_arrow_table(
-                direct_results.resultSet.results, schema_bytes, description
+                direct_results.resultSet.results,
+                lz4_compressed,
+                schema_bytes,
+                description,
             )
             arrow_queue_opt = ArrowQueue(arrow_results, n_rows, 0)
         else:
@@ -740,6 +751,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             status=operation_state,
             has_been_closed_server_side=has_been_closed_server_side,
             has_more_rows=has_more_rows,
+            lz4_compressed=lz4_compressed,
             command_handle=resp.operationHandle,
             description=description,
             arrow_schema_bytes=schema_bytes,
@@ -783,7 +795,9 @@ def _check_direct_results_for_error(t_spark_direct_results):
                     t_spark_direct_results.closeOperation
                 )
 
-    def execute_command(self, operation, session_handle, max_rows, max_bytes, cursor):
+    def execute_command(
+        self, operation, session_handle, max_rows, max_bytes, lz4_compression, cursor
+    ):
         assert session_handle is not None
 
         spark_arrow_types = ttypes.TSparkArrowTypes(
@@ -802,7 +816,7 @@ def execute_command(self, operation, session_handle, max_rows, max_bytes, cursor
                 maxRows=max_rows, maxBytes=max_bytes
             ),
             canReadArrowResult=True,
-            canDecompressLZ4Result=False,
+            canDecompressLZ4Result=lz4_compression,
             canDownloadResult=False,
             confOverlay={
                 # We want to receive proper Timestamp arrow types.
@@ -916,6 +930,7 @@ def fetch_results(
         max_rows,
         max_bytes,
         expected_row_start_offset,
+        lz4_compressed,
         arrow_schema_bytes,
         description,
     ):
@@ -941,7 +956,7 @@ def fetch_results(
                 )
             )
         arrow_results, n_rows = self._create_arrow_table(
-            resp.results, arrow_schema_bytes, description
+            resp.results, lz4_compressed, arrow_schema_bytes, description
         )
         arrow_queue = ArrowQueue(arrow_results, n_rows)
 

src/databricks/sql/thrift_backend.py

@@ -40,7 +40,7 @@ def remaining_rows(self) -> pyarrow.Table:
 
 ExecuteResponse = namedtuple(
     "ExecuteResponse",
-    "status has_been_closed_server_side has_more_rows description "
+    "status has_been_closed_server_side has_more_rows description lz4_compressed "
     "command_handle arrow_queue arrow_schema_bytes",
 )
 

src/databricks/sql/utils.py

@@ -49,11 +49,15 @@ def test_query_with_large_wide_result_set(self):
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 1000
         with self.cursor() as cursor:
-            uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)])
-            cursor.execute("SELECT id, {uuids} FROM RANGE({rows})".format(uuids=uuids, rows=rows))
-            for row_id, row in enumerate(self.fetch_rows(cursor, rows, fetchmany_size)):
-                self.assertEqual(row[0], row_id)  # Verify no rows are dropped in the middle.
-                self.assertEqual(len(row[1]), 36)
+            for lz4_compression in [False, True]:
+                cursor.connection.lz4_compression=lz4_compression
+                uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)])
+                cursor.execute("SELECT id, {uuids} FROM RANGE({rows})".format(uuids=uuids, rows=rows))
+                self.assertEqual(lz4_compression, cursor.active_result_set.lz4_compressed)
+                for row_id, row in enumerate(self.fetch_rows(cursor, rows, fetchmany_size)):
+                    self.assertEqual(row[0], row_id)  # Verify no rows are dropped in the middle.
+                    self.assertEqual(len(row[1]), 36)
+
 
     def test_query_with_large_narrow_result_set(self):
         resultSize = 300 * 1000 * 1000  # 300 MB
@@ -85,10 +89,10 @@ def test_long_running_query(self):
                 start = time.time()
 
                 cursor.execute("""SELECT count(*)
-                         FROM RANGE({scale}) x
-                         JOIN RANGE({scale0}) y
-                         ON from_unixtime(x.id * y.id, "yyyy-MM-dd") LIKE "%not%a%date%" 
-                         """.format(scale=scale_factor * scale0, scale0=scale0))
+                        FROM RANGE({scale}) x
+                        JOIN RANGE({scale0}) y
+                        ON from_unixtime(x.id * y.id, "yyyy-MM-dd") LIKE "%not%a%date%" 
+                        """.format(scale=scale_factor * scale0, scale0=scale0))
 
                 n, = cursor.fetchone()
                 self.assertEqual(n, 0)

tests/e2e/common/large_queries_mixin.py

@@ -510,6 +510,20 @@ def test_timezone_with_timestamp(self):
                 self.assertEqual(arrow_result_table.field(0).type, ts_type)
                 self.assertEqual(arrow_result_value, expected.timestamp() * 1000000)
 
+    @skipUnless(pysql_supports_arrow(), 'arrow test needs arrow support')
+    def test_can_flip_compression(self):
+        with self.cursor() as cursor:
+            cursor.execute("SELECT array(1,2,3,4)")
+            cursor.fetchall()
+            lz4_compressed = cursor.active_result_set.lz4_compressed
+            #The endpoint should support compression
+            self.assertEqual(lz4_compressed, True)
+            cursor.connection.lz4_compression=False
+            cursor.execute("SELECT array(1,2,3,4)")
+            cursor.fetchall()
+            lz4_compressed = cursor.active_result_set.lz4_compressed
+            self.assertEqual(lz4_compressed, False)
+
     def _should_have_native_complex_types(self):
         return pysql_has_version(">=", 2) and is_thrift_v5_plus(self.arguments)
 

tests/e2e/driver_tests.py

@@ -38,6 +38,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
                 has_been_closed_server_side=True,
                 has_more_rows=False,
                 description=Mock(),
+                lz4_compressed=Mock(),
                 command_handle=None,
                 arrow_queue=arrow_queue,
                 arrow_schema_bytes=schema.serialize().to_pybytes()))
@@ -50,7 +51,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
     def make_dummy_result_set_from_batch_list(batch_list):
         batch_index = 0
 
-        def fetch_results(op_handle, max_rows, max_bytes, expected_row_start_offset,
+        def fetch_results(op_handle, max_rows, max_bytes, expected_row_start_offset, lz4_compressed,
                           arrow_schema_bytes, description):
             nonlocal batch_index
             results = FetchTests.make_arrow_queue(batch_list[batch_index])
@@ -71,6 +72,7 @@ def fetch_results(op_handle, max_rows, max_bytes, expected_row_start_offset,
                 has_more_rows=True,
                 description=[(f'col{col_id}', 'integer', None, None, None, None, None)
                              for col_id in range(num_cols)],
+                lz4_compressed=Mock(),
                 command_handle=None,
                 arrow_queue=None,
                 arrow_schema_bytes=None))