Skip to content

Commit e38c7d2

Browse files
GuilucandGuilucand
committed
Add bz2 and xz output support
1 parent 7de2da2 commit e38c7d2

10 files changed

Lines changed: 194 additions & 25 deletions

File tree

‎Cargo.lock‎

Lines changed: 20 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/assembler_pipeline/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,8 @@ pub fn get_final_output_writer<
8181
"zst" | "zstd" => {
8282
W::new_compressed_zstd(&output_file, output_compression_level.min(22))
8383
}
84+
"bz2" => W::new_compressed_bz2(&output_file, output_compression_level.min(9)),
85+
"xz" => W::new_compressed_xz(&output_file, output_compression_level.min(9)),
8486
_ => W::new_plain(&output_file),
8587
},
8688
None => W::new_plain(&output_file),

‎crates/querier/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ parking_lot = "0.12.3"
3434
lz4 = "1.25.0"
3535
flate2 = "1.0.30"
3636
zstd = "0.13.2"
37+
bzip2 = "0.6.1"
38+
xz2 = "0.1.7"
3739
ggcat-logging = { version = "2.1.0", path = "../logging" }
3840
anyhow = "1.0.89"
3941
typenum = "1.18.0"

‎crates/querier/src/pipeline/colored_query_output.rs‎

Lines changed: 68 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -29,28 +29,65 @@ use std::path::PathBuf;
2929
use std::sync::atomic::{AtomicUsize, Ordering};
3030

3131
enum QueryOutputFileWriter {
32-
Plain(File),
33-
LZ4Compressed(lz4::Encoder<File>),
34-
GzipCompressed(flate2::write::GzEncoder<File>),
35-
ZstdCompressed(zstd::stream::write::Encoder<'static, File>),
32+
Plain(Option<File>),
33+
LZ4Compressed(Option<lz4::Encoder<File>>),
34+
GzipCompressed(Option<flate2::write::GzEncoder<File>>),
35+
ZstdCompressed(Option<zstd::stream::write::Encoder<'static, File>>),
36+
BZ2Compressed(Option<bzip2::write::BzEncoder<File>>),
37+
XZCompressed(Option<xz2::write::XzEncoder<File>>),
3638
}
3739

3840
impl Write for QueryOutputFileWriter {
3941
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
4042
match self {
41-
QueryOutputFileWriter::Plain(w) => w.write(buf),
42-
QueryOutputFileWriter::LZ4Compressed(w) => w.write(buf),
43-
QueryOutputFileWriter::GzipCompressed(w) => w.write(buf),
44-
QueryOutputFileWriter::ZstdCompressed(w) => w.write(buf),
43+
QueryOutputFileWriter::Plain(w) => w.as_mut().unwrap().write(buf),
44+
QueryOutputFileWriter::LZ4Compressed(w) => w.as_mut().unwrap().write(buf),
45+
QueryOutputFileWriter::GzipCompressed(w) => w.as_mut().unwrap().write(buf),
46+
QueryOutputFileWriter::ZstdCompressed(w) => w.as_mut().unwrap().write(buf),
47+
QueryOutputFileWriter::BZ2Compressed(w) => w.as_mut().unwrap().write(buf),
48+
QueryOutputFileWriter::XZCompressed(w) => w.as_mut().unwrap().write(buf),
4549
}
4650
}
4751

4852
fn flush(&mut self) -> std::io::Result<()> {
4953
match self {
50-
QueryOutputFileWriter::Plain(w) => w.flush(),
51-
QueryOutputFileWriter::LZ4Compressed(w) => w.flush(),
52-
QueryOutputFileWriter::GzipCompressed(w) => w.flush(),
53-
QueryOutputFileWriter::ZstdCompressed(w) => w.flush(),
54+
QueryOutputFileWriter::Plain(w) => w.as_mut().unwrap().flush(),
55+
QueryOutputFileWriter::LZ4Compressed(w) => w.as_mut().unwrap().flush(),
56+
QueryOutputFileWriter::GzipCompressed(w) => w.as_mut().unwrap().flush(),
57+
QueryOutputFileWriter::ZstdCompressed(w) => w.as_mut().unwrap().flush(),
58+
QueryOutputFileWriter::BZ2Compressed(w) => w.as_mut().unwrap().flush(),
59+
QueryOutputFileWriter::XZCompressed(w) => w.as_mut().unwrap().flush(),
60+
}
61+
}
62+
}
63+
64+
impl Drop for QueryOutputFileWriter {
65+
fn drop(&mut self) {
66+
match self {
67+
QueryOutputFileWriter::Plain(writer) => {
68+
writer.as_mut().unwrap().flush().unwrap();
69+
}
70+
QueryOutputFileWriter::LZ4Compressed(writer) => {
71+
let (mut file, err) = writer.take().unwrap().finish();
72+
err.unwrap();
73+
file.flush().unwrap();
74+
}
75+
QueryOutputFileWriter::GzipCompressed(writer) => {
76+
let mut file = writer.take().unwrap().finish().unwrap();
77+
file.flush().unwrap();
78+
}
79+
QueryOutputFileWriter::ZstdCompressed(writer) => {
80+
let mut file = writer.take().unwrap().finish().unwrap();
81+
file.flush().unwrap();
82+
}
83+
QueryOutputFileWriter::BZ2Compressed(writer) => {
84+
let mut file = writer.take().unwrap().finish().unwrap();
85+
file.flush().unwrap();
86+
}
87+
QueryOutputFileWriter::XZCompressed(writer) => {
88+
let mut file = writer.take().unwrap().finish().unwrap();
89+
file.flush().unwrap();
90+
}
5491
}
5592
}
5693
}
@@ -93,26 +130,37 @@ pub fn colored_query_output<MH: HashFunctionFactory, CX: ColorsManager>(
93130
let query_output = Mutex::new((
94131
BufWriter::new(
95132
match output_file.extension().map(|e| e.to_str()).flatten() {
96-
Some("lz4") => QueryOutputFileWriter::LZ4Compressed(
133+
Some("lz4") => QueryOutputFileWriter::LZ4Compressed(Some(
97134
lz4::EncoderBuilder::new()
98135
.level(output_compression_level.min(16))
99136
.build(query_output_file)
100137
.unwrap(),
101-
),
138+
)),
102139
Some("gz") => {
103-
QueryOutputFileWriter::GzipCompressed(flate2::GzBuilder::new().write(
104-
query_output_file,
105-
Compression::new(output_compression_level.min(9)),
140+
QueryOutputFileWriter::GzipCompressed(Some(
141+
flate2::GzBuilder::new().write(
142+
query_output_file,
143+
Compression::new(output_compression_level.min(9)),
144+
),
106145
))
107146
}
108-
Some("zst") | Some("zstd") => QueryOutputFileWriter::ZstdCompressed(
147+
Some("zst") | Some("zstd") => QueryOutputFileWriter::ZstdCompressed(Some(
109148
zstd::stream::write::Encoder::new(
110149
query_output_file,
111150
output_compression_level.min(22) as i32,
112151
)
113152
.unwrap(),
114-
),
115-
_ => QueryOutputFileWriter::Plain(query_output_file),
153+
)),
154+
Some("bz2") => QueryOutputFileWriter::BZ2Compressed(Some(
155+
bzip2::write::BzEncoder::new(
156+
query_output_file,
157+
bzip2::Compression::new(output_compression_level.min(9)),
158+
),
159+
)),
160+
Some("xz") => QueryOutputFileWriter::XZCompressed(Some(
161+
xz2::write::XzEncoder::new(query_output_file, output_compression_level.min(9)),
162+
)),
163+
_ => QueryOutputFileWriter::Plain(Some(query_output_file)),
116164
},
117165
),
118166
0,

‎crates/sequence_output/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ flate2 = "1.1.9"
2929
parking_lot = "0.12.5"
3030
serde = "1.0.228"
3131
zstd = "0.13.3"
32+
bzip2 = "0.6.1"
33+
xz2 = "0.1.7"
3234

3335

3436
[features]

‎crates/sequence_output/src/structured_sequences.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,14 @@ pub trait StructuredSequenceBackendInit: Sync + Send + Sized {
4545
unimplemented!()
4646
}
4747

48+
fn new_compressed_bz2(_path: impl AsRef<Path>, _level: u32) -> Self {
49+
unimplemented!()
50+
}
51+
52+
fn new_compressed_xz(_path: impl AsRef<Path>, _level: u32) -> Self {
53+
unimplemented!()
54+
}
55+
4856
fn new_plain(_path: impl AsRef<Path>) -> Self {
4957
unimplemented!()
5058
}

‎crates/sequence_output/src/structured_sequences/binary.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,14 @@ impl<CX: ColorsManager, LinksInfo: IdentSequenceWriter + SequenceExtraData>
9393
unimplemented!()
9494
}
9595

96+
fn new_compressed_bz2(_path: impl AsRef<Path>, _level: u32) -> Self {
97+
unimplemented!()
98+
}
99+
100+
fn new_compressed_xz(_path: impl AsRef<Path>, _level: u32) -> Self {
101+
unimplemented!()
102+
}
103+
96104
fn new_plain(_path: impl AsRef<Path>) -> Self {
97105
unimplemented!()
98106
}

‎crates/sequence_output/src/structured_sequences/fasta.rs‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
use crate::indirect_reads_extractor::{ReadExtractWorkData, indirect_read_extract_parts};
22
use crate::structured_sequences::{IdentSequenceWriter, StructuredSequenceBackend};
3+
use bzip2::Compression as BzipCompression;
34
use colors::colors_manager::ColorsManager;
45
use colors::colors_manager::color_types::PartialUnitigsColorStructure;
56
use config::{DEFAULT_OUTPUT_BUFFER_SIZE, DEFAULT_PER_CPU_BUFFER_SIZE};
67
use dynamic_dispatch::dynamic_dispatch;
7-
use flate2::Compression;
8+
use flate2::Compression as FlateCompression;
89
use flate2::write::GzEncoder;
910
use io::compressed_read::CompressedRead;
1011
use io::concurrent::temp_reads::extra_data::{SequenceExtraData, TempBuffer};
@@ -44,7 +45,7 @@ impl<CX: ColorsManager, LinksInfo: IdentSequenceWriter> StructuredSequenceBacken
4445
fn new_compressed_gzip(path: impl AsRef<Path>, level: u32) -> Self {
4546
let compress_stream = GzEncoder::new(
4647
BufWriter::with_capacity(DEFAULT_OUTPUT_BUFFER_SIZE, File::create(&path).unwrap()),
47-
Compression::new(level),
48+
FlateCompression::new(level),
4849
);
4950

5051
FastaWriter {
@@ -96,6 +97,38 @@ impl<CX: ColorsManager, LinksInfo: IdentSequenceWriter> StructuredSequenceBacken
9697
}
9798
}
9899

100+
fn new_compressed_bz2(path: impl AsRef<Path>, level: u32) -> Self {
101+
let compress_stream = bzip2::write::BzEncoder::new(
102+
BufWriter::with_capacity(DEFAULT_OUTPUT_BUFFER_SIZE, File::create(&path).unwrap()),
103+
BzipCompression::new(level),
104+
);
105+
106+
FastaWriter {
107+
writer: Box::new(SequencesWriterWrapper::new(BufWriter::with_capacity(
108+
DEFAULT_OUTPUT_BUFFER_SIZE,
109+
compress_stream,
110+
))),
111+
path: path.as_ref().to_path_buf(),
112+
_phantom: PhantomData,
113+
}
114+
}
115+
116+
fn new_compressed_xz(path: impl AsRef<Path>, level: u32) -> Self {
117+
let compress_stream = xz2::write::XzEncoder::new(
118+
BufWriter::with_capacity(DEFAULT_OUTPUT_BUFFER_SIZE, File::create(&path).unwrap()),
119+
level,
120+
);
121+
122+
FastaWriter {
123+
writer: Box::new(SequencesWriterWrapper::new(BufWriter::with_capacity(
124+
DEFAULT_OUTPUT_BUFFER_SIZE,
125+
compress_stream,
126+
))),
127+
path: path.as_ref().to_path_buf(),
128+
_phantom: PhantomData,
129+
}
130+
}
131+
99132
fn new_plain(path: impl AsRef<Path>) -> Self {
100133
FastaWriter {
101134
writer: Box::new(SequencesWriterWrapper::new(BufWriter::with_capacity(

‎crates/sequence_output/src/structured_sequences/gfa.rs‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
use crate::indirect_reads_extractor::{ReadExtractWorkData, indirect_read_extract_parts};
22
use crate::structured_sequences::{IdentSequenceWriter, StructuredSequenceBackend};
3+
use bzip2::Compression as BzipCompression;
34
use colors::colors_manager::ColorsManager;
45
use colors::colors_manager::color_types::PartialUnitigsColorStructure;
56
use config::{DEFAULT_OUTPUT_BUFFER_SIZE, DEFAULT_PER_CPU_BUFFER_SIZE};
67
use dynamic_dispatch::dynamic_dispatch;
7-
use flate2::Compression;
8+
use flate2::Compression as FlateCompression;
89
use flate2::write::GzEncoder;
910
use hashes::HashableSequence;
1011
use io::compressed_read::CompressedRead;
@@ -58,7 +59,7 @@ impl<const VERSION: u32, CX: ColorsManager, LinksInfo: IdentSequenceWriter>
5859
fn new_compressed_gzip(path: impl AsRef<Path>, level: u32) -> Self {
5960
let compress_stream = GzEncoder::new(
6061
BufWriter::with_capacity(DEFAULT_OUTPUT_BUFFER_SIZE, File::create(&path).unwrap()),
61-
Compression::new(level),
62+
FlateCompression::new(level),
6263
);
6364

6465
GFAWriter {
@@ -110,6 +111,38 @@ impl<const VERSION: u32, CX: ColorsManager, LinksInfo: IdentSequenceWriter>
110111
}
111112
}
112113

114+
fn new_compressed_bz2(path: impl AsRef<Path>, level: u32) -> Self {
115+
let compress_stream = bzip2::write::BzEncoder::new(
116+
BufWriter::with_capacity(DEFAULT_OUTPUT_BUFFER_SIZE, File::create(&path).unwrap()),
117+
BzipCompression::new(level),
118+
);
119+
120+
GFAWriter {
121+
writer: Box::new(SequencesWriterWrapper::new(BufWriter::with_capacity(
122+
DEFAULT_OUTPUT_BUFFER_SIZE,
123+
compress_stream,
124+
))),
125+
path: path.as_ref().to_path_buf(),
126+
_phantom: PhantomData,
127+
}
128+
}
129+
130+
fn new_compressed_xz(path: impl AsRef<Path>, level: u32) -> Self {
131+
let compress_stream = xz2::write::XzEncoder::new(
132+
BufWriter::with_capacity(DEFAULT_OUTPUT_BUFFER_SIZE, File::create(&path).unwrap()),
133+
level,
134+
);
135+
136+
GFAWriter {
137+
writer: Box::new(SequencesWriterWrapper::new(BufWriter::with_capacity(
138+
DEFAULT_OUTPUT_BUFFER_SIZE,
139+
compress_stream,
140+
))),
141+
path: path.as_ref().to_path_buf(),
142+
_phantom: PhantomData,
143+
}
144+
}
145+
113146
fn new_plain(path: impl AsRef<Path>) -> Self {
114147
GFAWriter {
115148
writer: Box::new(SequencesWriterWrapper::new(BufWriter::with_capacity(

0 commit comments

Comments
 (0)