Merge pull request #178 from estie-inc/feat/gzip-prealloc

feat: preallocate gzip decode buffer

Author: yuly3

src/bench_support.rs

@@ -48,7 +48,7 @@ pub fn parse_inline_result_table(
 }
 
 pub fn decode_gzip_chunk(body: Bytes) -> Result<Bytes> {
-    rowset::decode_gzip_chunk(body).map_err(crate::Error::from)
+    rowset::decode_gzip_chunk(body, None).map_err(crate::Error::from)
 }
 
 pub fn inline_rows_to_result_table(

src/rowset/gzip.rs

@@ -5,7 +5,12 @@ use flate2::bufread::GzDecoder;
 
 use crate::error::ProtocolError;
 
-pub(crate) fn decode_gzip_chunk(body: Bytes) -> std::result::Result<Bytes, ProtocolError> {
+const MAX_GZIP_PREALLOCATED_BYTES: usize = 128 * 1024 * 1024;
+
+pub(crate) fn decode_gzip_chunk(
+    body: Bytes,
+    expected_uncompressed_size: Option<usize>,
+) -> std::result::Result<Bytes, ProtocolError> {
     if body.is_empty() {
         return Ok(body);
     }
@@ -16,7 +21,8 @@ pub(crate) fn decode_gzip_chunk(body: Bytes) -> std::result::Result<Bytes, Proto
 
     if body[0] == 0x1f && body[1] == 0x8b {
         let mut decoder = GzDecoder::new(&body[..]);
-        let mut decoded = Vec::new();
+        let mut decoded =
+            Vec::with_capacity(gzip_preallocated_capacity(expected_uncompressed_size));
         decoder
             .read_to_end(&mut decoded)
             .map_err(ProtocolError::gzip_decode)?;
@@ -26,6 +32,12 @@ pub(crate) fn decode_gzip_chunk(body: Bytes) -> std::result::Result<Bytes, Proto
     }
 }
 
+fn gzip_preallocated_capacity(expected_uncompressed_size: Option<usize>) -> usize {
+    expected_uncompressed_size
+        .filter(|size| *size <= MAX_GZIP_PREALLOCATED_BYTES)
+        .unwrap_or(0)
+}
+
 #[cfg(test)]
 mod tests {
     use std::error::Error as StdError;
@@ -35,12 +47,26 @@ mod tests {
 
     #[test]
     fn malformed_gzip_is_protocol_error() {
-        let err: Error = decode_gzip_chunk(Bytes::from_static(b"\x1f\x8bgarbage"))
+        let err: Error = decode_gzip_chunk(Bytes::from_static(b"\x1f\x8bgarbage"), None)
             .unwrap_err()
             .into();
 
         assert_eq!(err.kind(), ErrorKind::Protocol);
         assert!(err.to_string().contains("gzip decompression failed"));
         assert!(StdError::source(&err).is_some());
     }
+
+    #[test]
+    fn gzip_capacity_hint_uses_bounded_expected_size() {
+        assert_eq!(gzip_preallocated_capacity(Some(1024)), 1024);
+        assert_eq!(
+            gzip_preallocated_capacity(Some(MAX_GZIP_PREALLOCATED_BYTES)),
+            MAX_GZIP_PREALLOCATED_BYTES
+        );
+        assert_eq!(
+            gzip_preallocated_capacity(Some(MAX_GZIP_PREALLOCATED_BYTES + 1)),
+            0
+        );
+        assert_eq!(gzip_preallocated_capacity(None), 0);
+    }
 }

src/rowset/parser.rs

@@ -501,9 +501,11 @@ pub(crate) async fn parse_remote_chunk_result_table_async(
         }
     };
     let query_id_for_work = Arc::clone(&query_id);
+    let expected_uncompressed_size = workload.uncompressed_bytes;
 
     let parse_work_result = execute_parse_work(workload, blocking_parse_limiter, move || {
-        let bytes = decode_gzip_chunk(body).map_err(QueryScopedRepr::from)?;
+        let bytes =
+            decode_gzip_chunk(body, expected_uncompressed_size).map_err(QueryScopedRepr::from)?;
         parse_table_with_shape(
             schema,
             query_id_for_work,