lindera_core/dictionary_builder/
dict.rs

1use std::borrow::Cow;
2use std::collections::BTreeMap;
3use std::fs::File;
4use std::io::Write;
5use std::io::{self, Read};
6use std::path::{Path, PathBuf};
7use std::str::FromStr;
8
9use anyhow::anyhow;
10use byteorder::{LittleEndian, WriteBytesExt};
11use csv::StringRecord;
12use derive_builder::Builder;
13use encoding_rs::{Encoding, UTF_8};
14use encoding_rs_io::DecodeReaderBytesBuilder;
15use glob::glob;
16use log::{debug, warn};
17use yada::builder::DoubleArrayBuilder;
18
19use crate::decompress::Algorithm;
20use crate::dictionary::word_entry::{WordEntry, WordId};
21use crate::dictionary_builder::utils::compress_write;
22use crate::error::LinderaErrorKind;
23use crate::LinderaResult;
24
25#[derive(Builder, Debug)]
26#[builder(name = "DictBuilderOptions")]
27#[builder(build_fn(name = "builder"))]
28pub struct DictBuilder {
29    #[builder(default = "true")]
30    flexible_csv: bool,
31    /* If set to UTF-8, it can also read UTF-16 files with BOM. */
32    #[builder(default = "\"UTF-8\".into()", setter(into))]
33    encoding: Cow<'static, str>,
34    #[builder(default = "Algorithm::Deflate")]
35    compress_algorithm: Algorithm,
36    #[builder(default = "false")]
37    normalize_details: bool,
38    #[builder(default = "false")]
39    skip_invalid_cost_or_id: bool,
40}
41
42impl DictBuilder {
43    pub fn build(&self, input_dir: &Path, output_dir: &Path) -> LinderaResult<()> {
44        let pattern = if let Some(path) = input_dir.to_str() {
45            format!("{}/*.csv", path)
46        } else {
47            return Err(
48                LinderaErrorKind::Io.with_error(anyhow::anyhow!("Failed to convert path to &str."))
49            );
50        };
51
52        let mut filenames: Vec<PathBuf> = Vec::new();
53        for entry in
54            glob(&pattern).map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?
55        {
56            match entry {
57                Ok(path) => {
58                    if let Some(filename) = path.file_name() {
59                        filenames.push(Path::new(input_dir).join(filename));
60                    } else {
61                        return Err(LinderaErrorKind::Io
62                            .with_error(anyhow::anyhow!("failed to get filename")));
63                    };
64                }
65                Err(err) => return Err(LinderaErrorKind::Content.with_error(anyhow!(err))),
66            }
67        }
68
69        let encoding = Encoding::for_label_no_replacement(self.encoding.as_bytes());
70        let encoding = encoding.ok_or_else(|| {
71            LinderaErrorKind::Decode.with_error(anyhow!("Invalid encoding: {}", self.encoding))
72        })?;
73
74        let mut rows: Vec<StringRecord> = vec![];
75        for filename in filenames {
76            debug!("reading {:?}", filename);
77
78            let file = File::open(filename)
79                .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
80            let reader: Box<dyn Read> = if encoding == UTF_8 {
81                Box::new(file)
82            } else {
83                Box::new(
84                    DecodeReaderBytesBuilder::new()
85                        .encoding(Some(encoding))
86                        .build(file),
87                )
88            };
89            let mut rdr = csv::ReaderBuilder::new()
90                .has_headers(false)
91                .flexible(self.flexible_csv)
92                .from_reader(reader);
93
94            for result in rdr.records() {
95                let record =
96                    result.map_err(|err| LinderaErrorKind::Content.with_error(anyhow!(err)))?;
97                rows.push(record);
98            }
99        }
100
101        if self.normalize_details {
102            rows.sort_by_key(|row| normalize(&row[0]));
103        } else {
104            rows.sort_by(|a, b| a[0].cmp(&b[0]))
105        }
106
107        let wtr_da_path = output_dir.join(Path::new("dict.da"));
108        let mut wtr_da = io::BufWriter::new(
109            File::create(wtr_da_path)
110                .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?,
111        );
112
113        let wtr_vals_path = output_dir.join(Path::new("dict.vals"));
114        let mut wtr_vals = io::BufWriter::new(
115            File::create(wtr_vals_path)
116                .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?,
117        );
118
119        let mut word_entry_map: BTreeMap<String, Vec<WordEntry>> = BTreeMap::new();
120
121        for (row_id, row) in rows.iter().enumerate() {
122            let word_cost = match i16::from_str(row[3].trim()) {
123                Ok(wc) => wc,
124                Err(_err) => {
125                    if self.skip_invalid_cost_or_id {
126                        warn!("failed to parse word_cost: {:?}", row);
127                        continue;
128                    } else {
129                        return Err(LinderaErrorKind::Parse
130                            .with_error(anyhow::anyhow!("failed to parse word_cost")));
131                    }
132                }
133            };
134            let left_id = match u16::from_str(row[1].trim()) {
135                Ok(lid) => lid,
136                Err(_err) => {
137                    if self.skip_invalid_cost_or_id {
138                        warn!("failed to parse left_id: {:?}", row);
139                        continue;
140                    } else {
141                        return Err(LinderaErrorKind::Parse
142                            .with_error(anyhow::anyhow!("failed to parse left_id")));
143                    }
144                }
145            };
146            let right_id = match u16::from_str(row[2].trim()) {
147                Ok(rid) => rid,
148                Err(_err) => {
149                    if self.skip_invalid_cost_or_id {
150                        warn!("failed to parse right_id: {:?}", row);
151                        continue;
152                    } else {
153                        return Err(LinderaErrorKind::Parse
154                            .with_error(anyhow::anyhow!("failed to parse right_id")));
155                    }
156                }
157            };
158            let key = if self.normalize_details {
159                normalize(&row[0])
160            } else {
161                row[0].to_string()
162            };
163            word_entry_map.entry(key).or_default().push(WordEntry {
164                word_id: WordId(row_id as u32, true),
165                word_cost,
166                left_id,
167                right_id,
168            });
169        }
170
171        let wtr_words_path = output_dir.join(Path::new("dict.words"));
172        let mut wtr_words = io::BufWriter::new(
173            File::create(wtr_words_path)
174                .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?,
175        );
176
177        let wtr_words_idx_path = output_dir.join(Path::new("dict.wordsidx"));
178        let mut wtr_words_idx = io::BufWriter::new(
179            File::create(wtr_words_idx_path)
180                .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?,
181        );
182
183        let mut words_buffer = Vec::new();
184        let mut words_idx_buffer = Vec::new();
185        for row in rows.iter() {
186            let offset = words_buffer.len();
187            words_idx_buffer
188                .write_u32::<LittleEndian>(offset as u32)
189                .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
190
191            let joined_details = if self.normalize_details {
192                row.iter()
193                    .skip(4)
194                    .map(normalize)
195                    .collect::<Vec<String>>()
196                    .join("\0")
197            } else {
198                row.iter().skip(4).collect::<Vec<&str>>().join("\0")
199            };
200            let joined_details_len = u32::try_from(joined_details.as_bytes().len())
201                .map_err(|err| LinderaErrorKind::Serialize.with_error(anyhow::anyhow!(err)))?;
202            words_buffer
203                .write_u32::<LittleEndian>(joined_details_len)
204                .map_err(|err| LinderaErrorKind::Serialize.with_error(anyhow::anyhow!(err)))?;
205            words_buffer
206                .write_all(joined_details.as_bytes())
207                .map_err(|err| LinderaErrorKind::Serialize.with_error(anyhow::anyhow!(err)))?;
208        }
209
210        compress_write(&words_buffer, self.compress_algorithm, &mut wtr_words)?;
211        compress_write(
212            &words_idx_buffer,
213            self.compress_algorithm,
214            &mut wtr_words_idx,
215        )?;
216
217        wtr_words
218            .flush()
219            .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
220        wtr_words_idx
221            .flush()
222            .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
223
224        let mut id = 0u32;
225
226        let mut keyset: Vec<(&[u8], u32)> = vec![];
227        for (key, word_entries) in &word_entry_map {
228            let len = word_entries.len() as u32;
229            let val = (id << 5) | len; // 27bit for word ID, 5bit for different parts of speech on the same surface.
230            keyset.push((key.as_bytes(), val));
231            id += len;
232        }
233
234        let da_bytes = DoubleArrayBuilder::build(&keyset).ok_or_else(|| {
235            LinderaErrorKind::Io.with_error(anyhow::anyhow!("DoubleArray build error."))
236        })?;
237
238        compress_write(&da_bytes, self.compress_algorithm, &mut wtr_da)?;
239
240        let mut vals_buffer = Vec::new();
241        for word_entries in word_entry_map.values() {
242            for word_entry in word_entries {
243                word_entry
244                    .serialize(&mut vals_buffer)
245                    .map_err(|err| LinderaErrorKind::Serialize.with_error(anyhow::anyhow!(err)))?;
246            }
247        }
248
249        compress_write(&vals_buffer, self.compress_algorithm, &mut wtr_vals)?;
250
251        wtr_vals
252            .flush()
253            .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
254
255        Ok(())
256    }
257}
258
259fn normalize(text: &str) -> String {
260    text.to_string().replace('―', "—").replace('～', "〜")
261}
lindera_core/dictionary_builder/dict.rs

lindera_core/dictionary_builder/
dict.rs
