Skip to main content

lindera_analysis/token_filter/
mapping.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3
4use daachorse::DoubleArrayAhoCorasick;
5use daachorse::DoubleArrayAhoCorasickBuilder;
6use daachorse::MatchKind;
7use serde_json::Value;
8
9use crate::token_filter::TokenFilter;
10use lindera::LinderaResult;
11use lindera::error::LinderaErrorKind;
12use lindera::token::Token;
13
14pub const MAPPING_TOKEN_FILTER_NAME: &str = "mapping";
15
16pub type MappingTokenFilterConfig = Value;
17
18/// Replace characters with the specified character mappings.
19///
20#[derive(Clone)]
21pub struct MappingTokenFilter {
22    mapping: HashMap<String, String>,
23    trie: DoubleArrayAhoCorasick<u32>,
24}
25
26impl MappingTokenFilter {
27    /// Create a new `MappingTokenFilter` from a surface-to-replacement mapping.
28    ///
29    /// # Arguments
30    ///
31    /// * `mapping` - A map from surface text to its replacement text. Keys must be
32    ///   non-empty; an empty key would match at every byte position under the
33    ///   leftmost-longest search strategy used by `apply`, which is never a
34    ///   meaningful mapping and is therefore rejected.
35    ///
36    /// # Returns
37    ///
38    /// A `MappingTokenFilter`, or an error if `mapping` contains an empty key or
39    /// the underlying Aho-Corasick automaton fails to build.
40    pub fn new(mapping: HashMap<String, String>) -> LinderaResult<Self> {
41        if mapping.keys().any(|key| key.is_empty()) {
42            return Err(LinderaErrorKind::Args
43                .with_error(anyhow::anyhow!("mapping key must not be empty.")));
44        }
45
46        let mut keyset: Vec<(&[u8], u32)> = Vec::new();
47        let mut keys = mapping.keys().collect::<Vec<_>>();
48        keys.sort();
49        for (value, key) in keys.into_iter().enumerate() {
50            keyset.push((key.as_bytes(), value as u32));
51        }
52
53        let trie = DoubleArrayAhoCorasickBuilder::new()
54            .match_kind(MatchKind::LeftmostLongest)
55            .build_with_values(keyset)
56            .map_err(|err| LinderaErrorKind::Build.with_error(anyhow::anyhow!(err)))?;
57
58        Ok(Self { mapping, trie })
59    }
60
61    pub fn from_config(config: &MappingTokenFilterConfig) -> LinderaResult<Self> {
62        let mapping = config
63            .get("mapping")
64            .and_then(Value::as_object)
65            .ok_or_else(|| {
66                LinderaErrorKind::Parse.with_error(anyhow::anyhow!("mapping must be an object."))
67            })?
68            .iter()
69            .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
70            .collect::<HashMap<String, String>>();
71
72        Self::new(mapping)
73    }
74}
75
76impl TokenFilter for MappingTokenFilter {
77    fn name(&self) -> &'static str {
78        MAPPING_TOKEN_FILTER_NAME
79    }
80
81    /// Apply the filter to each token's surface, in place.
82    ///
83    /// Performs a single leftmost-longest pass over each token's surface with the
84    /// underlying Aho-Corasick automaton: at each position the longest configured
85    /// key wins, matches never overlap, and scanning resumes immediately after
86    /// each match.
87    ///
88    /// # Arguments
89    ///
90    /// * `tokens` - The tokens to filter; each token's `surface` is replaced in
91    ///   place with the mapped text.
92    fn apply(&self, tokens: &mut Vec<Token<'_>>) -> LinderaResult<()> {
93        for token in tokens.iter_mut() {
94            let mut result = String::with_capacity(token.surface.len());
95
96            {
97                let source = token.surface.as_ref();
98                let mut cursor = 0_usize;
99
100                for m in self.trie.leftmost_find_iter(source) {
101                    // Keys are validated non-empty in `new`, and all keys are
102                    // valid UTF-8, so matches are non-empty, non-overlapping,
103                    // strictly ascending, and always land on char boundaries.
104                    debug_assert!(m.start() >= cursor && m.end() > m.start());
105
106                    result.push_str(&source[cursor..m.start()]);
107                    result.push_str(&self.mapping[&source[m.start()..m.end()]]);
108                    cursor = m.end();
109                }
110
111                result.push_str(&source[cursor..]);
112            }
113
114            token.surface = Cow::Owned(result);
115        }
116
117        Ok(())
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use std::collections::HashMap;
124
125    use crate::token_filter::mapping::{MappingTokenFilter, MappingTokenFilterConfig};
126
127    #[test]
128    fn test_mapping_token_filter_empty_key_rejected() {
129        let mut mapping = HashMap::new();
130        mapping.insert(String::new(), "x".to_string());
131        assert!(MappingTokenFilter::new(mapping).is_err());
132
133        let mut mapping = HashMap::new();
134        mapping.insert("a".to_string(), "b".to_string());
135        assert!(MappingTokenFilter::new(mapping).is_ok());
136    }
137
138    #[test]
139    fn test_mapping_token_filter_config() {
140        let config_str = r#"
141        {
142            "mapping": {
143                "ア": "ア",
144                "イ": "イ",
145                "ウ": "ウ",
146                "エ": "エ",
147                "オ": "オ"
148            }
149        }
150        "#;
151        let result: Result<MappingTokenFilterConfig, _> = serde_json::from_str(config_str);
152        assert!(result.is_ok());
153    }
154
155    #[test]
156    fn test_mapping_token_filter() {
157        let config_str = r#"
158        {
159            "mapping": {
160                "ア": "ア",
161                "イ": "イ",
162                "ウ": "ウ",
163                "エ": "エ",
164                "オ": "オ"
165            }
166        }
167        "#;
168        let config = serde_json::from_str::<MappingTokenFilterConfig>(config_str).unwrap();
169
170        let result = MappingTokenFilter::from_config(&config);
171        assert!(result.is_ok());
172    }
173
174    #[test]
175    #[cfg(feature = "embed-ipadic")]
176    fn test_mapping_token_filter_apply_ipadic() {
177        use std::borrow::Cow;
178
179        use crate::token_filter::TokenFilter;
180        use lindera::dictionary::{DictionaryKind, WordId, load_embedded_dictionary};
181        use lindera::token::Token;
182        use lindera_dictionary::viterbi::LexType;
183
184        let config_str = r#"
185        {
186            "mapping": {
187                "籠": "篭"
188            }
189        }
190        "#;
191        let config = serde_json::from_str::<MappingTokenFilterConfig>(config_str).unwrap();
192
193        let filter = MappingTokenFilter::from_config(&config).unwrap();
194
195        let dictionary = load_embedded_dictionary(DictionaryKind::IPADIC).unwrap();
196
197        let mut tokens: Vec<Token> = vec![
198            Token {
199                surface: Cow::Borrowed("籠原"),
200                byte_start: 0,
201                byte_end: 6,
202                position: 0,
203                position_length: 1,
204                word_id: WordId::new(LexType::System, 312630),
205                dictionary: &dictionary,
206                user_dictionary: None,
207                details: Some(vec![
208                    Cow::Borrowed("名詞"),
209                    Cow::Borrowed("固有名詞"),
210                    Cow::Borrowed("一般"),
211                    Cow::Borrowed("*"),
212                    Cow::Borrowed("*"),
213                    Cow::Borrowed("*"),
214                    Cow::Borrowed("籠原"),
215                    Cow::Borrowed("カゴハラ"),
216                    Cow::Borrowed("カゴハラ"),
217                ]),
218            },
219            Token {
220                surface: Cow::Borrowed("駅"),
221                byte_start: 6,
222                byte_end: 9,
223                position: 1,
224                position_length: 1,
225                word_id: WordId::new(LexType::System, 383791),
226                dictionary: &dictionary,
227                user_dictionary: None,
228                details: Some(vec![
229                    Cow::Borrowed("名詞"),
230                    Cow::Borrowed("接尾"),
231                    Cow::Borrowed("地域"),
232                    Cow::Borrowed("*"),
233                    Cow::Borrowed("*"),
234                    Cow::Borrowed("*"),
235                    Cow::Borrowed("駅"),
236                    Cow::Borrowed("エキ"),
237                    Cow::Borrowed("エキ"),
238                ]),
239            },
240        ];
241
242        filter.apply(&mut tokens).unwrap();
243
244        assert_eq!(tokens.len(), 2);
245        assert_eq!(&tokens[0].surface, "篭原");
246        assert_eq!(&tokens[1].surface, "駅");
247    }
248
249    #[test]
250    #[cfg(feature = "embed-ipadic")]
251    fn test_mapping_token_filter_apply_longest_match() {
252        use std::borrow::Cow;
253
254        use crate::token_filter::TokenFilter;
255        use lindera::dictionary::{DictionaryKind, WordId, load_embedded_dictionary};
256        use lindera::token::Token;
257        use lindera_dictionary::viterbi::LexType;
258
259        let config_str = r#"
260        {
261            "mapping": {
262                "ab": "X",
263                "abc": "YY"
264            }
265        }
266        "#;
267        let config = serde_json::from_str::<MappingTokenFilterConfig>(config_str).unwrap();
268
269        let filter = MappingTokenFilter::from_config(&config).unwrap();
270
271        let dictionary = load_embedded_dictionary(DictionaryKind::IPADIC).unwrap();
272
273        let mut tokens: Vec<Token> = vec![Token {
274            surface: Cow::Borrowed("abcab"),
275            byte_start: 0,
276            byte_end: 5,
277            position: 0,
278            position_length: 1,
279            word_id: WordId::new(LexType::System, 0),
280            dictionary: &dictionary,
281            user_dictionary: None,
282            details: None,
283        }];
284
285        filter.apply(&mut tokens).unwrap();
286
287        // "abc" (longest match at 0) wins over "ab"; "ab" (longest at 3) wins.
288        assert_eq!(tokens.len(), 1);
289        assert_eq!(&tokens[0].surface, "YYX");
290    }
291}