Add zstd output compression and configurable output level. Contribution from malfoy <antoine.limasset@gmail.com> pr #71

Author: Guilucand

Cargo.lock

@@ -4,11 +4,12 @@ use std::{
     any::Any,
     path::{Path, PathBuf},
     sync::Arc,
+    sync::atomic::Ordering,
 };
 
 use ::dynamic_dispatch::dynamic_dispatch;
 use colors::colors_manager::{ColorsManager, color_types::PartialUnitigsColorStructure};
-use config::{SwapPriority, get_compression_level_info, get_memory_mode};
+use config::{OUTPUT_COMPRESSION_LEVEL, SwapPriority, get_compression_level_info, get_memory_mode};
 use hashes::HashFunctionFactory;
 use io::{
     concurrent::temp_reads::extra_data::{
@@ -71,10 +72,15 @@ pub fn get_final_output_writer<
 >(
     output_file: &Path,
 ) -> W {
+    let output_compression_level = OUTPUT_COMPRESSION_LEVEL.load(Ordering::Relaxed);
+
     match output_file.extension() {
         Some(ext) => match ext.to_string_lossy().to_string().as_str() {
-            "lz4" => W::new_compressed_lz4(&output_file, 2),
-            "gz" => W::new_compressed_gzip(&output_file, 2),
+            "lz4" => W::new_compressed_lz4(&output_file, output_compression_level.min(16)),
+            "gz" => W::new_compressed_gzip(&output_file, output_compression_level.min(9)),
+            "zst" | "zstd" => {
+                W::new_compressed_zstd(&output_file, output_compression_level.min(22))
+            }
             _ => W::new_plain(&output_file),
         },
         None => W::new_plain(&output_file),

crates/assembler_pipeline/src/lib.rs

@@ -15,6 +15,7 @@ byteorder = "1.5.0"
 itertools = "0.13.0"
 lazy_static = "1.5.0"
 lz4 = "1.25.0"
+zstd = "0.13.2"
 rayon = "1.10.0"
 serde = "1.0.203"
 hashbrown = "0.15.3"

crates/cmdline/Cargo.toml

@@ -11,6 +11,7 @@ use colors::DefaultColorsSerializer;
 use colors::colors_manager::ColorMapReader;
 use colors::storage::deserializer::ColorsDeserializer;
 use config::ColorIndexType;
+use config::OUTPUT_COMPRESSION_LEVEL;
 use ggcat_api::{ExtraElaboration, GGCATConfig, GGCATInstance, GfaVersion};
 use ggcat_logging::UnrecoverableErrorLogging;
 use io::sequences_stream::general::GeneralSequenceBlockData;
@@ -136,6 +137,10 @@ struct CommonArgs {
     )]
     pub intermediate_compression_level: Option<u32>,
 
+    /// Compression level for final output files (applies to .lz4/.gz/.zst/.zstd outputs)
+    #[arg(long = "output-compression-level", default_value = "2")]
+    pub output_compression_level: u32,
+
     #[arg(hide = true, long = "only-bstats")]
     pub only_bstats: bool,
 }
@@ -313,6 +318,7 @@ fn initialize(args: &CommonArgs, out_file: &PathBuf) -> &'static GGCATInstance {
     ggcat_api::debug::DEBUG_KEEP_FILES.store(args.keep_temp_files, Ordering::Relaxed);
     *ggcat_api::debug::BUCKETS_COUNT_LOG_FORCE.lock() = args.buckets_count_log;
     ggcat_api::debug::DEBUG_ONLY_BSTATS.store(args.only_bstats, Ordering::Relaxed);
+    OUTPUT_COMPRESSION_LEVEL.store(args.output_compression_level, Ordering::Relaxed);
     *ggcat_api::debug::DEBUG_HASH_TYPE.lock() = match args.hash_type {
         HashType::Auto => ggcat_api::HashType::Auto,
         HashType::SeqHash => ggcat_api::HashType::SeqHash,

crates/cmdline/src/main.rs

@@ -149,6 +149,7 @@ impl SwapPriority {
 pub static KEEP_FILES: AtomicBool = AtomicBool::new(false);
 pub static INTERMEDIATE_COMPRESSION_LEVEL_SLOW: AtomicU32 = AtomicU32::new(3);
 pub static INTERMEDIATE_COMPRESSION_LEVEL_FAST: AtomicU32 = AtomicU32::new(0);
+pub static OUTPUT_COMPRESSION_LEVEL: AtomicU32 = AtomicU32::new(2);
 pub static PREFER_MEMORY: AtomicBool = AtomicBool::new(false);
 
 pub fn get_memory_mode(swap_priority: usize) -> MemoryFileMode {

crates/config/src/lib.rs

@@ -33,6 +33,7 @@ csv = "1.3.0"
 parking_lot = "0.12.3"
 lz4 = "1.25.0"
 flate2 = "1.0.30"
+zstd = "0.13.2"
 ggcat-logging = { version = "2.1.0", path = "../logging" }
 anyhow = "1.0.89"
 typenum = "1.18.0"

crates/querier/Cargo.toml

@@ -3,8 +3,8 @@ use crate::structs::query_colored_counters::{ColorsRange, QueryColoredCountersSe
 use colors::colors_manager::ColorMapReader;
 use colors::colors_manager::{ColorsManager, ColorsMergeManager};
 use config::{
-    ColorIndexType, KEEP_FILES, QUERIES_COUNT_MIN_BATCH, SwapPriority, get_compression_level_info,
-    get_memory_mode,
+    ColorIndexType, KEEP_FILES, OUTPUT_COMPRESSION_LEVEL, QUERIES_COUNT_MIN_BATCH, SwapPriority,
+    get_compression_level_info, get_memory_mode,
 };
 use flate2::Compression;
 use ggcat_logging::UnrecoverableErrorLogging;
@@ -32,6 +32,7 @@ enum QueryOutputFileWriter {
     Plain(File),
     LZ4Compressed(lz4::Encoder<File>),
     GzipCompressed(flate2::write::GzEncoder<File>),
+    ZstdCompressed(zstd::stream::write::Encoder<'static, File>),
 }
 
 impl Write for QueryOutputFileWriter {
@@ -40,6 +41,7 @@ impl Write for QueryOutputFileWriter {
             QueryOutputFileWriter::Plain(w) => w.write(buf),
             QueryOutputFileWriter::LZ4Compressed(w) => w.write(buf),
             QueryOutputFileWriter::GzipCompressed(w) => w.write(buf),
+            QueryOutputFileWriter::ZstdCompressed(w) => w.write(buf),
         }
     }
 
@@ -48,6 +50,7 @@ impl Write for QueryOutputFileWriter {
             QueryOutputFileWriter::Plain(w) => w.flush(),
             QueryOutputFileWriter::LZ4Compressed(w) => w.flush(),
             QueryOutputFileWriter::GzipCompressed(w) => w.flush(),
+            QueryOutputFileWriter::ZstdCompressed(w) => w.flush(),
         }
     }
 }
@@ -85,17 +88,29 @@ pub fn colored_query_output<MH: HashFunctionFactory, CX: ColorsManager>(
     let query_output_file = File::create(&output_file)
         .log_unrecoverable_error_with_data("Cannot create output file", output_file.display())?;
 
+    let output_compression_level = OUTPUT_COMPRESSION_LEVEL.load(Ordering::Relaxed);
+
     let query_output = Mutex::new((
         BufWriter::new(
             match output_file.extension().map(|e| e.to_str()).flatten() {
                 Some("lz4") => QueryOutputFileWriter::LZ4Compressed(
                     lz4::EncoderBuilder::new()
-                        .level(4)
+                        .level(output_compression_level.min(16))
                         .build(query_output_file)
                         .unwrap(),
                 ),
-                Some("gz") => QueryOutputFileWriter::GzipCompressed(
-                    flate2::GzBuilder::new().write(query_output_file, Compression::default()),
+                Some("gz") => {
+                    QueryOutputFileWriter::GzipCompressed(flate2::GzBuilder::new().write(
+                        query_output_file,
+                        Compression::new(output_compression_level.min(9)),
+                    ))
+                }
+                Some("zst") | Some("zstd") => QueryOutputFileWriter::ZstdCompressed(
+                    zstd::stream::write::Encoder::new(
+                        query_output_file,
+                        output_compression_level.min(22) as i32,
+                    )
+                    .unwrap(),
                 ),
                 _ => QueryOutputFileWriter::Plain(query_output_file),
             },

crates/querier/src/pipeline/colored_query_output.rs

@@ -28,6 +28,7 @@ byteorder = "1.5.0"
 flate2 = "1.1.9"
 parking_lot = "0.12.5"
 serde = "1.0.228"
+zstd = "0.13.3"
 
 
 [features]

crates/sequence_output/Cargo.toml

@@ -41,6 +41,10 @@ pub trait StructuredSequenceBackendInit: Sync + Send + Sized {
         unimplemented!()
     }
 
+    fn new_compressed_zstd(_path: impl AsRef<Path>, _level: u32) -> Self {
+        unimplemented!()
+    }
+
     fn new_plain(_path: impl AsRef<Path>) -> Self {
         unimplemented!()
     }

crates/sequence_output/src/structured_sequences.rs

@@ -89,6 +89,10 @@ impl<CX: ColorsManager, LinksInfo: IdentSequenceWriter + SequenceExtraData>
         unimplemented!()
     }
 
+    fn new_compressed_zstd(_path: impl AsRef<Path>, _level: u32) -> Self {
+        unimplemented!()
+    }
+
     fn new_plain(_path: impl AsRef<Path>) -> Self {
         unimplemented!()
     }