deepbiop_fq/encode/
parquet.rs

1use std::{
2    fmt::Display,
3    path::{Path, PathBuf},
4    sync::Arc,
5};
6
7use arrow::array::{Array, Int32Builder, ListBuilder, RecordBatch, StringBuilder};
8use arrow::datatypes::{DataType, Field, Schema};
9
10use bstr::BString;
11use derive_builder::Builder;
12use log::{debug, info};
13use serde::{Deserialize, Serialize};
14
15use crate::types::Element;
16use deepbiop_utils::io::write_parquet;
17
18use super::{traits::Encoder, EncoderOption, RecordData};
19use anyhow::{Context, Result};
20use pyo3::prelude::*;
21use rayon::prelude::*;
22
23use pyo3_stub_gen::derive::*;
24
25#[derive(Debug, Builder, Default)]
26pub struct ParquetData {
27    pub id: BString,        // id
28    pub seq: BString,       // kmer_seq
29    pub qual: Vec<Element>, // kmer_qual
30}
31
32#[gen_stub_pyclass]
33#[pyclass(module = "deepbiop.fq")]
34#[derive(Debug, Builder, Default, Clone, Serialize, Deserialize)]
35pub struct ParquetEncoder {
36    pub option: EncoderOption,
37}
38
39impl ParquetEncoder {
40    pub fn new(option: EncoderOption) -> Self {
41        Self { option }
42    }
43
44    fn generate_schema(&self) -> Arc<Schema> {
45        Arc::new(Schema::new(vec![
46            Field::new("id", DataType::Utf8, false),
47            Field::new("seq", DataType::Utf8, false),
48            Field::new(
49                "qual",
50                DataType::List(Box::new(Field::new("item", DataType::Int32, true)).into()),
51                false,
52            ),
53        ]))
54    }
55
56    fn generate_batches(
57        &self,
58        records: &[RecordData],
59        schema: &Arc<Schema>,
60    ) -> Result<Vec<RecordBatch>> {
61        // Process smaller batches to avoid 2GB limit
62        const BATCH_SIZE: usize = 10000; // Adjust this value based on your data size
63        let all_batches: Vec<_> = records
64            .par_chunks(BATCH_SIZE)
65            .map(|chunk| {
66                let _capacity = chunk.len();
67
68                let mut id_builder = StringBuilder::new();
69                let mut seq_builder = StringBuilder::new();
70                let mut qual_builder = ListBuilder::new(Int32Builder::new());
71
72                for data in chunk {
73                    let record = self
74                        .encode_record(data.id.as_ref(), data.seq.as_ref(), data.qual.as_ref())
75                        .context(format!(
76                            "encode fq read id {} error",
77                            String::from_utf8_lossy(data.id.as_ref())
78                        ))
79                        .unwrap();
80                    id_builder.append_value(record.id.to_string());
81                    seq_builder.append_value(record.seq.to_string());
82                    for qual in record.qual {
83                        qual_builder.values().append_value(qual);
84                    }
85                    qual_builder.append(true);
86                }
87
88                RecordBatch::try_new(
89                    schema.clone(),
90                    vec![
91                        Arc::new(id_builder.finish()) as Arc<dyn Array>,
92                        Arc::new(seq_builder.finish()) as Arc<dyn Array>,
93                        Arc::new(qual_builder.finish()) as Arc<dyn Array>,
94                    ],
95                )
96                .unwrap()
97            })
98            .collect();
99        debug!("all batches: {}", all_batches.len());
100        Ok(all_batches)
101    }
102
103    fn generate_batch(&self, records: &[RecordData], schema: &Arc<Schema>) -> Result<RecordBatch> {
104        let all_batches = self.generate_batches(records, schema)?;
105        // Concatenate all batches
106        arrow::compute::concat_batches(schema, &all_batches)
107            .context("Failed to concatenate record batches")
108    }
109
110    pub fn encode_chunk<P: AsRef<Path>>(
111        &mut self,
112        path: P,
113        chunk_size: usize,
114        parallel: bool,
115    ) -> Result<()> {
116        let schema = self.generate_schema();
117        let records = self.fetch_records(&path)?;
118        info!("Encoding records with chunk size {} ", chunk_size);
119
120        // create a folder for the chunk parquet files
121        let file_name = path.as_ref().file_name().unwrap().to_str().unwrap();
122        let chunks_folder = path
123            .as_ref()
124            .parent()
125            .unwrap()
126            .join(format!("{}_{}", file_name, "chunks"))
127            .to_path_buf();
128        // create the folder
129        std::fs::create_dir_all(&chunks_folder).context("Failed to create folder for chunks")?;
130
131        if parallel {
132            records
133                // .chunks(chunk_size)
134                .par_chunks(chunk_size)
135                .enumerate()
136                .for_each(|(idx, record)| {
137                    let record_batch = self
138                        .generate_batch(record, &schema)
139                        .context(format!("Failed to generate record batch for chunk {}", idx))
140                        .unwrap();
141                    let parquet_path = chunks_folder.join(format!("{}_{}.parquet", file_name, idx));
142                    write_parquet(parquet_path, record_batch, schema.clone())
143                        .context(format!("Failed to write parquet file for chunk {}", idx))
144                        .unwrap();
145                });
146        } else {
147            records
148                .chunks(chunk_size)
149                .enumerate()
150                .for_each(|(idx, record)| {
151                    let record_batch = self
152                        .generate_batch(record, &schema)
153                        .context(format!("Failed to generate record batch for chunk {}", idx))
154                        .unwrap();
155                    let parquet_path = chunks_folder.join(format!("{}_{}.parquet", file_name, idx));
156                    write_parquet(parquet_path, record_batch, schema.clone())
157                        .context(format!("Failed to write parquet file for chunk {}", idx))
158                        .unwrap();
159                });
160        }
161
162        Ok(())
163    }
164}
165
166impl Display for ParquetEncoder {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        write!(f, "FqEncoder {{ option: {} }}", self.option)
169    }
170}
171
172impl Encoder for ParquetEncoder {
173    type RecordOutput = Result<ParquetData>;
174    type EncodeOutput = Result<(Vec<RecordBatch>, Arc<Schema>)>;
175
176    fn encode_qual(&self, qual: &[u8], qual_offset: u8) -> Vec<Element> {
177        // input is quality of fastq
178        // 1. convert the quality to a score
179        // 2. return the score
180        let encoded_qual: Vec<Element> = qual
181            .par_iter()
182            .map(|&q| {
183                // Convert ASCII to Phred score for Phred+33 encoding
184                (q - qual_offset) as Element
185            })
186            .collect();
187        encoded_qual
188    }
189
190    fn encode_record(&self, id: &[u8], seq: &[u8], qual: &[u8]) -> Self::RecordOutput {
191        // encode the quality
192        let encoded_qual = self.encode_qual(qual, self.option.qual_offset);
193
194        let result = ParquetDataBuilder::default()
195            .id(id.into())
196            .seq(seq.into())
197            .qual(encoded_qual)
198            .build()
199            .context("Failed to build parquet data")?;
200        Ok(result)
201    }
202
203    fn encode<P: AsRef<Path>>(&mut self, path: P) -> Self::EncodeOutput {
204        // Define the schema of the data (one column of integers)
205        let schema = self.generate_schema();
206        let records = self.fetch_records(path)?;
207        let record_batch = self.generate_batches(&records, &schema)?;
208        Ok((record_batch, schema))
209    }
210
211    fn encode_multiple(&mut self, _paths: &[PathBuf], _parallel: bool) -> Self::EncodeOutput {
212        todo!()
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use deepbiop_utils::io::write_parquet_for_batches;
219
220    use crate::encode::EncoderOptionBuilder;
221
222    use super::*;
223    #[test]
224    fn test_encode_fq_for_parquet() {
225        let option = EncoderOptionBuilder::default().build().unwrap();
226
227        let mut encoder = ParquetEncoderBuilder::default()
228            .option(option)
229            .build()
230            .unwrap();
231        let (record_batch, scheme) = encoder.encode("tests/data/one_record.fq").unwrap();
232        write_parquet_for_batches("test.parquet", &record_batch, scheme).unwrap();
233        // remove test.parquet
234        std::fs::remove_file("test.parquet").unwrap();
235    }
236}
deepbiop_fq/encode/parquet.rs

deepbiop_fq/encode/
parquet.rs
