feat: add ZSTD network compression support and dynamic compression detection

klickhouse/Cargo.toml

@@ -37,6 +37,7 @@ chrono-tz = "0.10"
 futures-util = { version = "0.3", default-features = false, features = ["std"] }
 tokio-stream = "0.1"
 lz4 = { version = "1.28", optional = true }
+zstd = { version = "0.14", optional = true }
 klickhouse_derive = { version = "=0.15.0", optional = true, path = "../klickhouse_derive" }
 cityhash-rs = "1.0"
 bb8 = { version = "0.9", optional = true }
@@ -63,7 +64,7 @@ env_logger = "0.11"
 [features]
 default = ["derive", "compression", "serde", "bfloat16"]
 derive = ["klickhouse_derive"]
-compression = ["lz4"]
+compression = ["dep:lz4", "dep:zstd"]
 geo-types = ["dep:geo-types"]
 refinery = ["refinery-core", "time", "dep:async-trait"]
 serde = ["dep:serde", "serde_json", "uuid/serde", "chrono/serde", "half?/serde"]

klickhouse/README.md

@@ -28,7 +28,7 @@ $ cargo nextest run
 ## Feature flags
 
 - `derive`: Enable [klickhouse_derive], providing a derive macro for the [Row] trait. Default.
-- `compression`: `lz4` compression for client/server communication. Default.
+- `compression`: `lz4`/`zstd` compression for client/server communication. Default.
 - `serde`: Derivation of [serde::Serialize] and [serde::Deserialize] on various objects, and JSON support. Default.
 - `tls`: TLS support via [tokio-rustls](https://crates.io/crates/tokio-rustls).
 - `refinery`: Migrations via [refinery](https://crates.io/crates/refinery).

klickhouse/src/compression.rs

@@ -13,8 +13,8 @@ use crate::io::ClickhouseRead;
 use crate::protocol::CompressionMethod;
 use crate::{KlickhouseError, Result};
 
-pub async fn compress_block(block: Block, revision: u64) -> Result<(Vec<u8>, usize)> {
-    let mut raw = vec![];
+pub async fn lz4_compress_block(block: Block, revision: u64) -> Result<(Vec<u8>, usize)> {
+    let mut raw = Vec::new();
     block.write(&mut raw, revision).await?;
     let raw_len = raw.len();
     let mut compressed = Vec::<u8>::with_capacity(raw.len() + (raw.len() / 255) + 16 + 1);
@@ -39,7 +39,7 @@ pub async fn compress_block(block: Block, revision: u64) -> Result<(Vec<u8>, usi
     Ok((compressed, raw_len))
 }
 
-pub fn decompress_block(data: &[u8], decompressed_size: u32) -> Result<Vec<u8>> {
+pub fn lz4_decompress_block(data: &[u8], decompressed_size: u32) -> Result<Vec<u8>> {
     let mut output = Vec::with_capacity(decompressed_size as usize + 1);
 
     let out_len = unsafe {
@@ -63,21 +63,51 @@ pub fn decompress_block(data: &[u8], decompressed_size: u32) -> Result<Vec<u8>>
     Ok(output)
 }
 
-async fn read_compressed_blob(
-    reader: &mut impl ClickhouseRead,
-    compression: CompressionMethod,
-) -> Result<Vec<u8>> {
-    let checksum =
-        ((reader.read_u64_le().await? as u128) << 64u128) | (reader.read_u64_le().await? as u128);
+pub async fn zstd_compress_block(block: Block, revision: u64) -> Result<(Vec<u8>, usize)> {
+    let mut raw = Vec::new();
+    block.write(&mut raw, revision).await?;
+    let raw_len = raw.len();
+
+    let compression_level = 1;
+
+    match zstd::bulk::compress(&raw, compression_level) {
+        Ok(compressed) => Ok((compressed, raw_len)),
+        Err(e) => Err(KlickhouseError::ProtocolError(format!(
+            "ZSTD compression failed: {e}"
+        ))),
+    }
+}
+
+pub fn zstd_decompress_block(data: &[u8], decompressed_size: u32) -> Result<Vec<u8>> {
+    let mut output = vec![0u8; decompressed_size as usize];
+
+    match zstd::bulk::decompress_to_buffer(data, &mut output) {
+        Ok(out_len) => {
+            if out_len != decompressed_size as usize {
+                return Err(KlickhouseError::ProtocolError(format!(
+                    "ZSTD decompression size mismatch: expected {decompressed_size}, got {out_len}"
+                )));
+            }
+            Ok(output)
+        }
+        Err(e) => Err(KlickhouseError::ProtocolError(format!(
+            "malformed ZSTD compressed block: {e}"
+        ))),
+    }
+}
+
+async fn read_compressed_blob(reader: &mut impl ClickhouseRead) -> Result<Vec<u8>> {
+    let checksum_high = reader.read_u64_le().await? as u128;
+    let checksum_low = reader.read_u64_le().await? as u128;
+    let checksum = (checksum_high << 64) | checksum_low;
+
     let type_byte = reader.read_u8().await?;
-    if type_byte != compression.byte() {
+    let Ok(compression) = CompressionMethod::try_from(type_byte) else {
         return Err(KlickhouseError::ProtocolError(format!(
-            "unexpected compression algorithm identifier: '{:02X}', expected {:02X} ({:?})",
-            type_byte,
-            compression.byte(),
-            compression
+            "unexpected compression algorithm identifier: '{type_byte:02X}'",
         )));
-    }
+    };
+
     let compressed_size = reader.read_u32_le().await?;
     if compressed_size > MAX_COMPRESSION_SIZE {
         // 1 GB
@@ -89,37 +119,56 @@ async fn read_compressed_blob(
             "compressed payload too small! {compressed_size} < 9"
         )));
     }
+
     let decompressed_size = reader.read_u32_le().await?;
-    let mut compressed = vec![0u8; compressed_size as usize];
-    reader.read_exact(&mut compressed[9..]).await?;
-    compressed[0] = type_byte;
-    compressed[1..5].copy_from_slice(&compressed_size.to_le_bytes()[..]);
-    compressed[5..9].copy_from_slice(&decompressed_size.to_le_bytes()[..]);
-    let calc_checksum = cityhash_rs::cityhash_102_128(&compressed[..]);
+
+    let mut block = vec![0u8; compressed_size as usize];
+
+    block[0] = type_byte;
+    block[1..5].copy_from_slice(&compressed_size.to_le_bytes());
+    block[5..9].copy_from_slice(&decompressed_size.to_le_bytes());
+
+    reader.read_exact(&mut block[9..]).await?;
+
+    let calc_checksum = cityhash_rs::cityhash_102_128(&block[..]);
     if calc_checksum != checksum {
         return Err(KlickhouseError::ProtocolError(format!(
             "corrupt checksum from clickhouse '{calc_checksum:032X}' vs '{checksum:032X}'"
         )));
     }
-    let raw_block = crate::compression::decompress_block(&compressed[9..], decompressed_size)?;
+
+    let compressed_payload = &block[9..];
+
+    let raw_block = match compression {
+        CompressionMethod::LZ4 => {
+            crate::compression::lz4_decompress_block(compressed_payload, decompressed_size)?
+        }
+        CompressionMethod::ZSTD => {
+            crate::compression::zstd_decompress_block(compressed_payload, decompressed_size)?
+        }
+        _ => {
+            return Err(KlickhouseError::ProtocolError(format!(
+                "Compression method {compression:?} is not implemented yet for byte {type_byte}"
+            )));
+        }
+    };
+
     Ok(raw_block)
 }
 
 type BlockReadingFuture<R> =
     Pin<Box<dyn Future<Output = Result<(Vec<u8>, &'static mut R)>> + Send + Sync>>;
 
 pub struct DecompressionReader<'a, R: ClickhouseRead + 'static> {
-    mode: CompressionMethod,
     inner: Option<&'a mut R>,
     decompressed: Vec<u8>,
     position: usize,
     block_reading_future: Option<BlockReadingFuture<R>>,
 }
 
 impl<'a, R: ClickhouseRead + 'static> DecompressionReader<'a, R> {
-    pub fn new(mode: CompressionMethod, inner: &'a mut R) -> Self {
+    pub fn new(inner: &'a mut R) -> Self {
         Self {
-            mode,
             inner: Some(inner),
             decompressed: vec![],
             position: 0,
@@ -177,9 +226,8 @@ impl<R: ClickhouseRead + 'static> AsyncRead for DecompressionReader<'_, R> {
         while self.position >= self.decompressed.len() {
             let static_inner: &'static mut R =
                 unsafe { std::mem::transmute(self.inner.take().unwrap()) };
-            let mode = self.mode;
             self.block_reading_future = Some(Box::pin(async move {
-                let value = read_compressed_blob(static_inner, mode).await?;
+                let value = read_compressed_blob(static_inner).await?;
                 Ok((value, static_inner))
             }));
             match self.run_decompression(cx) {

klickhouse/src/internal_client_in.rs

@@ -50,9 +50,8 @@ impl<R: ClickhouseRead + 'static> InternalClientIn<R> {
     }
 
     #[cfg(feature = "compression")]
-    async fn decompress_data(&mut self, compression: CompressionMethod) -> Result<Block> {
-        let mut reader =
-            crate::compression::DecompressionReader::new(compression, &mut self.reader);
+    async fn decompress_data(&mut self, _compression: CompressionMethod) -> Result<Block> {
+        let mut reader = crate::compression::DecompressionReader::new(&mut self.reader);
 
         let block = Block::read(&mut reader, self.server_hello.revision_version).await?;
 

klickhouse/src/internal_client_out.rs

@@ -224,26 +224,50 @@ impl<W: ClickhouseWrite> InternalClientOut<W> {
     }
 
     #[cfg(feature = "compression")]
-    async fn compress_data(&mut self, byte: u8, block: Block) -> Result<()> {
-        let (out, decompressed_size) =
-            crate::compression::compress_block(block, self.server_hello.revision_version).await?;
-        let mut new_out = Vec::with_capacity(out.len() + 5);
-        new_out.push(byte);
-        new_out.extend_from_slice(&(out.len() as u32 + 9).to_le_bytes()[..]);
-        new_out.extend_from_slice(&(decompressed_size as u32).to_le_bytes()[..]);
+    async fn compress_data(&mut self, compression: CompressionMethod, block: Block) -> Result<()> {
+        let (out, decompressed_size) = match compression {
+            CompressionMethod::LZ4 => {
+                crate::compression::lz4_compress_block(block, self.server_hello.revision_version)
+                    .await?
+            }
+            CompressionMethod::ZSTD => {
+                crate::compression::zstd_compress_block(block, self.server_hello.revision_version)
+                    .await?
+            }
+            _ => {
+                unimplemented!(
+                    "Compression method {compression:?} is not implemented yet for revision version {}",
+                    self.server_hello.revision_version
+                )
+            }
+        };
+
+        let compressed_packet_len = out.len() + 9;
+
+        let mut new_out = Vec::with_capacity(compressed_packet_len);
+        new_out.push(compression.byte());
+
+        new_out.extend_from_slice(&(compressed_packet_len as u32).to_le_bytes());
+        new_out.extend_from_slice(&(decompressed_size as u32).to_le_bytes());
         new_out.extend(out);
 
         let hash = cityhash_rs::cityhash_102_128(&new_out[..]);
-        self.writer.write_u64_le((hash >> 64) as u64).await?;
-        self.writer.write_u64_le(hash as u64).await?;
-        // self.writer.write_u8(byte).await?;
-        // self.writer.write_u32_le(out.len() as u32).await?;
+        let hash_high = (hash >> 64) as u64;
+        let hash_low = hash as u64;
+
+        self.writer.write_u64_le(hash_high).await?;
+        self.writer.write_u64_le(hash_low).await?;
         self.writer.write_all(&new_out[..]).await?;
+
         Ok(())
     }
 
     #[cfg(not(feature = "compression"))]
-    async fn compress_data(&mut self, _byte: u8, _block: &Block) -> Result<()> {
+    async fn compress_data(
+        &mut self,
+        _compression: CompressionMethod,
+        _block: &Block,
+    ) -> Result<()> {
         panic!(
             "attempted to use compression when not compiled with `compression` feature in klickhouse"
         );
@@ -273,8 +297,10 @@ impl<W: ClickhouseWrite> InternalClientOut<W> {
                     .await?;
             }
             CompressionMethod::LZ4 => {
-                self.compress_data(CompressionMethod::LZ4.byte(), block)
-                    .await?;
+                self.compress_data(CompressionMethod::LZ4, block).await?;
+            }
+            CompressionMethod::ZSTD => {
+                self.compress_data(CompressionMethod::ZSTD, block).await?;
             }
         }
 

klickhouse/src/protocol.rs

@@ -178,19 +178,31 @@ pub enum ServerPacket {
 }
 
 #[derive(Clone, Copy, Debug, Default)]
-#[allow(unused)]
+#[repr(u8)]
 pub enum CompressionMethod {
     #[cfg_attr(not(feature = "compression"), default)]
-    None,
+    None = 0x02,
     #[cfg_attr(feature = "compression", default)]
-    LZ4,
+    LZ4 = 0x82,
+    ZSTD = 0x90,
 }
 
 impl CompressionMethod {
-    pub fn byte(&self) -> u8 {
-        match self {
-            CompressionMethod::None => 0x02,
-            CompressionMethod::LZ4 => 0x82,
-        }
+    #[inline]
+    pub const fn byte(self) -> u8 {
+        self as u8
+    }
+}
+
+impl TryFrom<u8> for CompressionMethod {
+    type Error = &'static str;
+
+    fn try_from(byte: u8) -> Result<Self, Self::Error> {
+        Ok(match byte {
+            0x02 => CompressionMethod::None,
+            0x82 => CompressionMethod::LZ4,
+            0x90 => CompressionMethod::ZSTD,
+            _ => return Err("Unknown compression method byte"),
+        })
     }
 }