result compression

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -7,12 +7,13 @@
 logger = logging.getLogger(__name__)
 
 
-def test_sea_query_execution():
+def test_sea_query_execution_with_compression():
     """
-    Test executing a query using the SEA backend.
+    Test executing a query using the SEA backend with result compression.
 
     This function connects to a Databricks SQL endpoint using the SEA backend,
-    executes a simple query, and verifies that execution completes successfully.
+    executes a simple query with result compression enabled and disabled,
+    and verifies that execution completes successfully.
     """
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
@@ -27,7 +28,8 @@ def test_sea_query_execution():
         sys.exit(1)
 
     try:
-        # Create connection with SEA backend
+        # Test with compression enabled
+        logger.info("Creating connection with LZ4 compression enabled")
         connection = Connection(
             server_hostname=server_hostname,
             http_path=http_path,
@@ -36,27 +38,50 @@ def test_sea_query_execution():
             schema="default",
             use_sea=True,
             user_agent_entry="SEA-Test-Client",
-            use_cloud_fetch=False,
+            use_cloud_fetch=True,  # Enable cloud fetch to use compression
+            enable_query_result_lz4_compression=True,  # Enable LZ4 compression
         )
 
         logger.info(
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
         logger.info(f"backend type: {type(connection.session.backend)}")
 
-        # Create a cursor and execute a simple query
+        # Execute a simple query with compression enabled
         cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
-
-        logger.info("Executing query: SELECT 1 as test_value")
+        logger.info("Executing query with LZ4 compression: SELECT 1 as test_value")
         cursor.execute("SELECT 1 as test_value")
+        logger.info("Query with compression executed successfully")
+        cursor.close()
+        connection.close()
+        logger.info("Successfully closed SEA session with compression enabled")
 
-        # We don't fetch results yet since we haven't implemented the fetch functionality
-        logger.info("Query executed successfully")
+        # Test with compression disabled
+        logger.info("Creating connection with LZ4 compression disabled")
+        connection = Connection(
+            server_hostname=server_hostname,
+            http_path=http_path,
+            access_token=access_token,
+            catalog=catalog,
+            schema="default",
+            use_sea=True,
+            user_agent_entry="SEA-Test-Client",
+            use_cloud_fetch=False,  # Enable cloud fetch
+            enable_query_result_lz4_compression=False,  # Disable LZ4 compression
+        )
 
-        # Close cursor and connection
+        logger.info(
+            f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
+        )
+        
+        # Execute a simple query with compression disabled
+        cursor = connection.cursor(arraysize=0, buffer_size_bytes=0)
+        logger.info("Executing query without compression: SELECT 1 as test_value")
+        cursor.execute("SELECT 1 as test_value")
+        logger.info("Query without compression executed successfully")
         cursor.close()
         connection.close()
-        logger.info("Successfully closed SEA session")
+        logger.info("Successfully closed SEA session with compression disabled")
 
     except Exception as e:
         logger.error(f"Error during SEA query execution test: {str(e)}")
@@ -65,7 +90,7 @@ def test_sea_query_execution():
         logger.error(traceback.format_exc())
         sys.exit(1)
 
-    logger.info("SEA query execution test completed successfully")
+    logger.info("SEA query execution test with compression completed successfully")
 
 
 def test_sea_session():
@@ -133,5 +158,5 @@ def test_sea_session():
     # Test session management
     test_sea_session()
 
-    # Test query execution
-    test_sea_query_execution()
+    # Test query execution with compression
+    test_sea_query_execution_with_compression()

src/databricks/sql/backend/models/requests.py

@@ -7,7 +7,6 @@
 from typing import Dict, List, Any, Optional, Union
 from dataclasses import dataclass, field
 
-
 @dataclass
 class StatementParameter:
     """Parameter for a SQL statement."""
@@ -21,7 +20,6 @@ class StatementParameter:
 class ExecuteStatementRequest:
     """Request to execute a SQL statement."""
 
-    # TODO: result_compression key
     warehouse_id: str
     statement: str
     session_id: str
@@ -34,6 +32,7 @@ class ExecuteStatementRequest:
     parameters: Optional[List[StatementParameter]] = None
     catalog: Optional[str] = None
     schema: Optional[str] = None
+    result_compression: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
@@ -58,6 +57,9 @@ def to_dict(self) -> Dict[str, Any]:
 
         if self.schema:
             result["schema"] = self.schema
+            
+        if self.result_compression:
+            result["result_compression"] = self.result_compression
 
         if self.parameters:
             result["parameters"] = [

src/databricks/sql/backend/sea_backend.py

@@ -277,6 +277,8 @@ def execute_command(
 
         format = "ARROW_STREAM" if use_cloud_fetch else "JSON_ARRAY"
         disposition = "EXTERNAL_LINKS" if use_cloud_fetch else "INLINE"
+        result_compression = "LZ4_FRAME" if lz4_compression else "NONE"
+        
         request = ExecuteStatementRequest(
             warehouse_id=self.warehouse_id,
             session_id=sea_session_id,
@@ -288,6 +290,7 @@ def execute_command(
             row_limit=max_rows if max_rows > 0 else None,
             byte_limit=max_bytes if max_bytes > 0 else None,
             parameters=sea_parameters if sea_parameters else None,
+            result_compression=result_compression,
         )
 
         response_data = self.http_client._make_request(