multi frame decompression of lz4

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

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/queue.py

@@ -36,6 +36,25 @@
 logger = logging.getLogger(__name__)
 
 
+def decompress_multi_frame_lz4(attachment: bytes) -> bytes:
+    try:
+        decompressor = lz4.frame.LZ4FrameDecompressor()
+        arrow_file = decompressor.decompress(attachment)
+
+        # the attachment may be a concatenation of multiple LZ4 frames
+        while decompressor.unused_data:
+            remaining_data = decompressor.unused_data
+            arrow_file += decompressor.decompress(remaining_data)
+
+            logger.debug(f"LZ4 decompressed {len(arrow_file)} bytes from attachment")
+
+    except Exception as e:
+        logger.error(f"LZ4 decompression failed: {e}")
+        raise e
+
+    return arrow_file
+
+
 class SeaResultSetQueueFactory(ABC):
     @staticmethod
     def build_queue(
@@ -70,13 +89,14 @@ def build_queue(
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
             if result_data.attachment is not None:
                 arrow_file = (
-                    lz4.frame.decompress(result_data.attachment)
+                    decompress_multi_frame_lz4(result_data.attachment)
                     if lz4_compressed
                     else result_data.attachment
                 )
                 arrow_table = create_arrow_table_from_arrow_file(
                     arrow_file, description
                 )
+                logger.debug(f"Created arrow table with {arrow_table.num_rows} rows")
                 return ArrowQueue(arrow_table, manifest.total_row_count)
 
             # EXTERNAL_LINKS disposition