From a1fd3e17981de636dddfe114c7592644e3dd2f87 Mon Sep 17 00:00:00 2001 From: Negezor Date: Mon, 26 Oct 2026 17:16:15 +1100 Subject: [PATCH 1/5] chore(deps): add zstd dependency --- Cargo.lock | 35 +++++++++++++++++++++++++++++++++++ klickhouse/Cargo.toml | 3 ++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index f993166..c37b5b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -722,6 +722,7 @@ dependencies = [ "tokio-rustls", "tokio-stream", "uuid", + "zstd", ] [[package]] @@ -895,6 +896,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + [[package]] name = "portable-atomic" version = "1.13.1" @@ -2018,3 +2025,31 @@ name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db3ae9bf90af6e40197bef767eeefc46ad6397420fe873830bd85e5fbbfb3c70" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "8.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ded520f16b29e1a9b771c722221b2b0a0398c2316d9bb3c42b458a5414d7434b" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.1.2+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "481ada9c543b2323cbb7bbdbb5249177fbf0188ed959ed2c15e4b22421c84ed3" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/klickhouse/Cargo.toml b/klickhouse/Cargo.toml index e571597..57c9718 100644 --- a/klickhouse/Cargo.toml +++ b/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"] From 0fbe6d9d6812373d54c5816b522cd7b0a239d9a3 Mon Sep 17 00:00:00 2001 From: Negezor Date: Mon, 26 Oct 2026 17:51:02 +1100 Subject: [PATCH 2/5] feat: implement zstd compression method --- klickhouse/src/compression.rs | 81 +++++++++++++++++++++++---- klickhouse/src/internal_client_out.rs | 54 +++++++++++++----- klickhouse/src/protocol.rs | 4 +- 3 files changed, 112 insertions(+), 27 deletions(-) diff --git a/klickhouse/src/compression.rs b/klickhouse/src/compression.rs index 8bec27b..90e1613 100644 --- a/klickhouse/src/compression.rs +++ b/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, usize)> { - let mut raw = vec![]; +pub async fn lz4_compress_block(block: Block, revision: u64) -> Result<(Vec, usize)> { + let mut raw = Vec::new(); block.write(&mut raw, revision).await?; let raw_len = raw.len(); let mut compressed = Vec::::with_capacity(raw.len() + (raw.len() / 255) + 16 + 1); @@ -39,7 +39,7 @@ pub async fn compress_block(block: Block, revision: u64) -> Result<(Vec, usi Ok((compressed, raw_len)) } -pub fn decompress_block(data: &[u8], decompressed_size: u32) -> Result> { +pub fn lz4_decompress_block(data: &[u8], decompressed_size: u32) -> Result> { let mut output = Vec::with_capacity(decompressed_size as usize + 1); let out_len = unsafe { @@ -63,12 +63,47 @@ pub fn decompress_block(data: &[u8], decompressed_size: u32) -> Result> Ok(output) } +pub async fn zstd_compress_block(block: Block, revision: u64) -> Result<(Vec, 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> { + 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, compression: CompressionMethod, ) -> Result> { - let checksum = - ((reader.read_u64_le().await? as u128) << 64u128) | (reader.read_u64_le().await? as u128); + 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() { return Err(KlickhouseError::ProtocolError(format!( @@ -78,6 +113,7 @@ async fn read_compressed_blob( compression ))); } + let compressed_size = reader.read_u32_le().await?; if compressed_size > MAX_COMPRESSION_SIZE { // 1 GB @@ -89,19 +125,40 @@ 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) } diff --git a/klickhouse/src/internal_client_out.rs b/klickhouse/src/internal_client_out.rs index 456a22c..7a823de 100644 --- a/klickhouse/src/internal_client_out.rs +++ b/klickhouse/src/internal_client_out.rs @@ -224,26 +224,50 @@ impl InternalClientOut { } #[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 InternalClientOut { .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?; } } diff --git a/klickhouse/src/protocol.rs b/klickhouse/src/protocol.rs index 3b7ac91..e83023c 100644 --- a/klickhouse/src/protocol.rs +++ b/klickhouse/src/protocol.rs @@ -182,8 +182,9 @@ pub enum ServerPacket { pub enum CompressionMethod { #[cfg_attr(not(feature = "compression"), default)] None, - #[cfg_attr(feature = "compression", default)] LZ4, + #[cfg_attr(feature = "compression", default)] + ZSTD, } impl CompressionMethod { @@ -191,6 +192,7 @@ impl CompressionMethod { match self { CompressionMethod::None => 0x02, CompressionMethod::LZ4 => 0x82, + CompressionMethod::ZSTD => 0x90, } } } From 71ccdbca84fbb4904cc752048b3b4f4b54d72456 Mon Sep 17 00:00:00 2001 From: Negezor Date: Mon, 26 Oct 2026 18:04:53 +1100 Subject: [PATCH 3/5] refactor: simplify protocol compression byte --- klickhouse/src/protocol.rs | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/klickhouse/src/protocol.rs b/klickhouse/src/protocol.rs index e83023c..f1c6094 100644 --- a/klickhouse/src/protocol.rs +++ b/klickhouse/src/protocol.rs @@ -178,21 +178,31 @@ pub enum ServerPacket { } #[derive(Clone, Copy, Debug, Default)] -#[allow(unused)] +#[repr(u8)] pub enum CompressionMethod { #[cfg_attr(not(feature = "compression"), default)] - None, - LZ4, + None = 0x02, #[cfg_attr(feature = "compression", default)] - ZSTD, + LZ4 = 0x82, + ZSTD = 0x90, } impl CompressionMethod { - pub fn byte(&self) -> u8 { - match self { - CompressionMethod::None => 0x02, - CompressionMethod::LZ4 => 0x82, - CompressionMethod::ZSTD => 0x90, - } + #[inline] + pub const fn byte(self) -> u8 { + self as u8 + } +} + +impl TryFrom for CompressionMethod { + type Error = &'static str; + + fn try_from(byte: u8) -> Result { + Ok(match byte { + 0x02 => CompressionMethod::None, + 0x82 => CompressionMethod::LZ4, + 0x90 => CompressionMethod::ZSTD, + _ => return Err("Unknown compression method byte"), + }) } } From 331ea955932119a23f64a1ec1b63d2671161b690 Mon Sep 17 00:00:00 2001 From: Negezor Date: Mon, 26 Oct 2026 18:14:17 +1100 Subject: [PATCH 4/5] feat(compression): implement dynamic determine compression --- klickhouse/src/compression.rs | 21 ++++++--------------- klickhouse/src/internal_client_in.rs | 5 ++--- 2 files changed, 8 insertions(+), 18 deletions(-) diff --git a/klickhouse/src/compression.rs b/klickhouse/src/compression.rs index 90e1613..981975f 100644 --- a/klickhouse/src/compression.rs +++ b/klickhouse/src/compression.rs @@ -96,23 +96,17 @@ pub fn zstd_decompress_block(data: &[u8], decompressed_size: u32) -> Result Result> { +async fn read_compressed_blob(reader: &mut impl ClickhouseRead) -> Result> { 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 { @@ -166,7 +160,6 @@ type BlockReadingFuture = Pin, &'static mut R)>> + Send + Sync>>; pub struct DecompressionReader<'a, R: ClickhouseRead + 'static> { - mode: CompressionMethod, inner: Option<&'a mut R>, decompressed: Vec, position: usize, @@ -174,9 +167,8 @@ pub struct DecompressionReader<'a, R: ClickhouseRead + 'static> { } 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, @@ -234,9 +226,8 @@ impl 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) { diff --git a/klickhouse/src/internal_client_in.rs b/klickhouse/src/internal_client_in.rs index 73fb833..9233b0c 100644 --- a/klickhouse/src/internal_client_in.rs +++ b/klickhouse/src/internal_client_in.rs @@ -50,9 +50,8 @@ impl InternalClientIn { } #[cfg(feature = "compression")] - async fn decompress_data(&mut self, compression: CompressionMethod) -> Result { - let mut reader = - crate::compression::DecompressionReader::new(compression, &mut self.reader); + async fn decompress_data(&mut self, _compression: CompressionMethod) -> Result { + let mut reader = crate::compression::DecompressionReader::new(&mut self.reader); let block = Block::read(&mut reader, self.server_hello.revision_version).await?; From d4ef5bc97f852c27becb572f8acd5bcbc1506940 Mon Sep 17 00:00:00 2001 From: Negezor Date: Mon, 26 Oct 2026 18:21:41 +1100 Subject: [PATCH 5/5] docs: add zstd compression mention --- klickhouse/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/klickhouse/README.md b/klickhouse/README.md index ac46413..5655dfb 100644 --- a/klickhouse/README.md +++ b/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).