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                Ok(QueryExpr::ProbabilisticCommand(
180                    ProbabilisticCommand::CreateHll {
181                        name,
182                        if_not_exists,
183                    },
184                ))
185            }
186            Token::Ident(ref name) if name.eq_ignore_ascii_case("SKETCH") => {
187                self.advance()?;
188                let if_not_exists = self.match_if_not_exists()?;
189                let name = self.expect_ident()?;
190                // Optional WIDTH and DEPTH. `DEPTH` lexes as the
191                // reserved `Token::Depth` keyword (used by graph
192                // traversal modifiers), so `consume_ident_ci("DEPTH")`
193                // would never match — fall through to the typed
194                // consumer for that arm.
195                let mut width = 1000usize;
196                let mut depth = 5usize;
197                for _ in 0..2 {
198                    if self.consume_ident_ci("WIDTH")? {
199                        width = self.parse_positive_integer("WIDTH")? as usize;
200                    } else if self.consume(&Token::Depth)? {
201                        depth = self.parse_positive_integer("DEPTH")? as usize;
202                    } else {
203                        break;
204                    }
205                }
206                Ok(QueryExpr::ProbabilisticCommand(
207                    ProbabilisticCommand::CreateSketch {
208                        name,
209                        width,
210                        depth,
211                        if_not_exists,
212                    },
213                ))
214            }
215            Token::Ident(ref name) if name.eq_ignore_ascii_case("FILTER") => {
216                self.advance()?;
217                let if_not_exists = self.match_if_not_exists()?;
218                let name = self.expect_ident()?;
219                let capacity = if self.consume_ident_ci("CAPACITY")? {
220                    self.parse_positive_integer("CAPACITY")? as usize
221                } else {
222                    100_000
223                };
224                Ok(QueryExpr::ProbabilisticCommand(
225                    ProbabilisticCommand::CreateFilter {
226                        name,
227                        capacity,
228                        if_not_exists,
229                    },
230                ))
231            }
232            _ => Err(ParseError::expected(
233                vec!["HLL", "SKETCH", "FILTER"],
234                self.peek(),
235                self.position(),
236            )),
237        }
238    }
239
240    /// Parse DROP HLL|SKETCH|FILTER ... (called after DROP has been consumed)
241    pub fn parse_drop_probabilistic(&mut self) -> Result<QueryExpr, ParseError> {
242        match self.peek().clone() {
243            Token::Ident(ref name) if name.eq_ignore_ascii_case("HLL") => {
244                self.advance()?;
245                let if_exists = self.match_if_exists()?;
246                let name = self.expect_ident()?;
247                Ok(QueryExpr::ProbabilisticCommand(
248                    ProbabilisticCommand::DropHll { name, if_exists },
249                ))
250            }
251            Token::Ident(ref name) if name.eq_ignore_ascii_case("SKETCH") => {
252                self.advance()?;
253                let if_exists = self.match_if_exists()?;
254                let name = self.expect_ident()?;
255                Ok(QueryExpr::ProbabilisticCommand(
256                    ProbabilisticCommand::DropSketch { name, if_exists },
257                ))
258            }
259            Token::Ident(ref name) if name.eq_ignore_ascii_case("FILTER") => {
260                self.advance()?;
261                let if_exists = self.match_if_exists()?;
262                let name = self.expect_ident()?;
263                Ok(QueryExpr::ProbabilisticCommand(
264                    ProbabilisticCommand::DropFilter { name, if_exists },
265                ))
266            }
267            _ => Err(ParseError::expected(
268                vec!["HLL", "SKETCH", "FILTER"],
269                self.peek(),
270                self.position(),
271            )),
272        }
273    }
274}