Skip to main content

reddb_server/storage/query/parser/
probabilistic_commands.rs

1//! Parser for probabilistic data structure commands: HLL, SKETCH, FILTER
2
3use super::super::ast::{ProbabilisticCommand, QueryExpr};
4use super::super::lexer::Token;
5use super::error::ParseError;
6use super::Parser;
7
8impl<'a> Parser<'a> {
9    /// Parse HLL subcommand: HLL ADD|COUNT|MERGE|INFO name ...
10    pub fn parse_hll_command(&mut self) -> Result<QueryExpr, ParseError> {
11        self.advance()?; // consume HLL
12
13        match self.peek().clone() {
14            Token::Add => {
15                self.advance()?;
16                let name = self.expect_ident()?;
17                let mut elements = Vec::new();
18                loop {
19                    match self.peek() {
20                        Token::String(_) => elements.push(self.parse_string()?),
21                        Token::Eof | Token::Semi => break,
22                        _ => break,
23                    }
24                }
25                Ok(QueryExpr::ProbabilisticCommand(
26                    ProbabilisticCommand::HllAdd { name, elements },
27                ))
28            }
29            Token::Count => {
30                self.advance()?;
31                let mut names = Vec::new();
32                while let Token::Ident(_) = self.peek() {
33                    names.push(self.expect_ident()?);
34                }
35                Ok(QueryExpr::ProbabilisticCommand(
36                    ProbabilisticCommand::HllCount { names },
37                ))
38            }
39            Token::Ident(ref name) if name.eq_ignore_ascii_case("MERGE") => {
40                self.advance()?;
41                let dest = self.expect_ident()?;
42                let mut sources = Vec::new();
43                while let Token::Ident(_) = self.peek() {
44                    sources.push(self.expect_ident()?);
45                }
46                Ok(QueryExpr::ProbabilisticCommand(
47                    ProbabilisticCommand::HllMerge { dest, sources },
48                ))
49            }
50            Token::Ident(ref name) if name.eq_ignore_ascii_case("INFO") => {
51                self.advance()?;
52                let name = self.expect_ident()?;
53                Ok(QueryExpr::ProbabilisticCommand(
54                    ProbabilisticCommand::HllInfo { name },
55                ))
56            }
57            _ => Err(ParseError::expected(
58                vec!["ADD", "COUNT", "MERGE", "INFO"],
59                self.peek(),
60                self.position(),
61            )),
62        }
63    }
64
65    /// Parse SKETCH subcommand: SKETCH ADD|COUNT|MERGE|INFO name ...
66    pub fn parse_sketch_command(&mut self) -> Result<QueryExpr, ParseError> {
67        self.advance()?; // consume SKETCH
68
69        match self.peek().clone() {
70            Token::Add => {
71                self.advance()?;
72                let name = self.expect_ident()?;
73                let element = self.parse_string()?;
74                let count = if matches!(self.peek(), Token::Integer(_)) {
75                    self.parse_integer()? as u64
76                } else {
77                    1
78                };
79                Ok(QueryExpr::ProbabilisticCommand(
80                    ProbabilisticCommand::SketchAdd {
81                        name,
82                        element,
83                        count,
84                    },
85                ))
86            }
87            Token::Count => {
88                self.advance()?;
89                let name = self.expect_ident()?;
90                let element = self.parse_string()?;
91                Ok(QueryExpr::ProbabilisticCommand(
92                    ProbabilisticCommand::SketchCount { name, element },
93                ))
94            }
95            Token::Ident(ref name) if name.eq_ignore_ascii_case("MERGE") => {
96                self.advance()?;
97                let dest = self.expect_ident()?;
98                let mut sources = Vec::new();
99                while let Token::Ident(_) = self.peek() {
100                    sources.push(self.expect_ident()?);
101                }
102                Ok(QueryExpr::ProbabilisticCommand(
103                    ProbabilisticCommand::SketchMerge { dest, sources },
104                ))
105            }
106            Token::Ident(ref name) if name.eq_ignore_ascii_case("INFO") => {
107                self.advance()?;
108                let name = self.expect_ident()?;
109                Ok(QueryExpr::ProbabilisticCommand(
110                    ProbabilisticCommand::SketchInfo { name },
111                ))
112            }
113            _ => Err(ParseError::expected(
114                vec!["ADD", "COUNT", "MERGE", "INFO"],
115                self.peek(),
116                self.position(),
117            )),
118        }
119    }
120
121    /// Parse FILTER subcommand: FILTER ADD|CHECK|DELETE|COUNT|INFO name ...
122    pub fn parse_filter_command(&mut self) -> Result<QueryExpr, ParseError> {
123        self.advance()?; // consume FILTER
124
125        match self.peek().clone() {
126            Token::Add => {
127                self.advance()?;
128                let name = self.expect_ident()?;
129                let element = self.parse_string()?;
130                Ok(QueryExpr::ProbabilisticCommand(
131                    ProbabilisticCommand::FilterAdd { name, element },
132                ))
133            }
134            Token::Ident(ref name) if name.eq_ignore_ascii_case("CHECK") => {
135                self.advance()?;
136                let name = self.expect_ident()?;
137                let element = self.parse_string()?;
138                Ok(QueryExpr::ProbabilisticCommand(
139                    ProbabilisticCommand::FilterCheck { name, element },
140                ))
141            }
142            Token::Delete => {
143                self.advance()?;
144                let name = self.expect_ident()?;
145                let element = self.parse_string()?;
146                Ok(QueryExpr::ProbabilisticCommand(
147                    ProbabilisticCommand::FilterDelete { name, element },
148                ))
149            }
150            Token::Count => {
151                self.advance()?;
152                let name = self.expect_ident()?;
153                Ok(QueryExpr::ProbabilisticCommand(
154                    ProbabilisticCommand::FilterCount { name },
155                ))
156            }
157            Token::Ident(ref name) if name.eq_ignore_ascii_case("INFO") => {
158                self.advance()?;
159                let name = self.expect_ident()?;
160                Ok(QueryExpr::ProbabilisticCommand(
161                    ProbabilisticCommand::FilterInfo { name },
162                ))
163            }
164            _ => Err(ParseError::expected(
165                vec!["ADD", "CHECK", "DELETE", "COUNT", "INFO"],
166                self.peek(),
167                self.position(),
168            )),
169        }
170    }
171
172    /// Parse CREATE HLL|SKETCH|FILTER ... (called after CREATE has been consumed)
173    pub fn parse_create_probabilistic(&mut self) -> Result<QueryExpr, ParseError> {
174        match self.peek().clone() {
175            Token::Ident(ref name) if name.eq_ignore_ascii_case("HLL") => {
176                self.advance()?;
177                let if_not_exists = self.match_if_not_exists()?;
178                let name = self.expect_ident()?;
179                let precision = if self.consume_ident_ci("PRECISION")? {
180                    self.parse_positive_integer("PRECISION")? as u8
181                } else {
182                    14
183                };
184                Ok(QueryExpr::ProbabilisticCommand(
185                    ProbabilisticCommand::CreateHll {
186                        name,
187                        precision,
188                        if_not_exists,
189                    },
190                ))
191            }
192            Token::Ident(ref name) if name.eq_ignore_ascii_case("SKETCH") => {
193                self.advance()?;
194                let if_not_exists = self.match_if_not_exists()?;
195                let name = self.expect_ident()?;
196                // Optional WIDTH and DEPTH. `DEPTH` lexes as the
197                // reserved `Token::Depth` keyword (used by graph
198                // traversal modifiers), so `consume_ident_ci("DEPTH")`
199                // would never match — fall through to the typed
200                // consumer for that arm.
201                let mut width = 1000usize;
202                let mut depth = 5usize;
203                for _ in 0..2 {
204                    if self.consume_ident_ci("WIDTH")? {
205                        width = self.parse_positive_integer("WIDTH")? as usize;
206                    } else if self.consume(&Token::Depth)? {
207                        depth = self.parse_positive_integer("DEPTH")? as usize;
208                    } else {
209                        break;
210                    }
211                }
212                Ok(QueryExpr::ProbabilisticCommand(
213                    ProbabilisticCommand::CreateSketch {
214                        name,
215                        width,
216                        depth,
217                        if_not_exists,
218                    },
219                ))
220            }
221            Token::Ident(ref name) if name.eq_ignore_ascii_case("FILTER") => {
222                self.advance()?;
223                let if_not_exists = self.match_if_not_exists()?;
224                let name = self.expect_ident()?;
225                let capacity = if self.consume_ident_ci("CAPACITY")? {
226                    self.parse_positive_integer("CAPACITY")? as usize
227                } else {
228                    100_000
229                };
230                Ok(QueryExpr::ProbabilisticCommand(
231                    ProbabilisticCommand::CreateFilter {
232                        name,
233                        capacity,
234                        if_not_exists,
235                    },
236                ))
237            }
238            _ => Err(ParseError::expected(
239                vec!["HLL", "SKETCH", "FILTER"],
240                self.peek(),
241                self.position(),
242            )),
243        }
244    }
245
246    /// Parse DROP HLL|SKETCH|FILTER ... (called after DROP has been consumed)
247    pub fn parse_drop_probabilistic(&mut self) -> Result<QueryExpr, ParseError> {
248        match self.peek().clone() {
249            Token::Ident(ref name) if name.eq_ignore_ascii_case("HLL") => {
250                self.advance()?;
251                let if_exists = self.match_if_exists()?;
252                let name = self.expect_ident()?;
253                Ok(QueryExpr::ProbabilisticCommand(
254                    ProbabilisticCommand::DropHll { name, if_exists },
255                ))
256            }
257            Token::Ident(ref name) if name.eq_ignore_ascii_case("SKETCH") => {
258                self.advance()?;
259                let if_exists = self.match_if_exists()?;
260                let name = self.expect_ident()?;
261                Ok(QueryExpr::ProbabilisticCommand(
262                    ProbabilisticCommand::DropSketch { name, if_exists },
263                ))
264            }
265            Token::Ident(ref name) if name.eq_ignore_ascii_case("FILTER") => {
266                self.advance()?;
267                let if_exists = self.match_if_exists()?;
268                let name = self.expect_ident()?;
269                Ok(QueryExpr::ProbabilisticCommand(
270                    ProbabilisticCommand::DropFilter { name, if_exists },
271                ))
272            }
273            _ => Err(ParseError::expected(
274                vec!["HLL", "SKETCH", "FILTER"],
275                self.peek(),
276                self.position(),
277            )),
278        }
279    }
280}