add support for TermSetQuery in query parser (#1683)

src/query/query_parser/logical_ast.rs

@@ -15,6 +15,11 @@ pub enum LogicalLiteral {
         lower: Bound<Term>,
         upper: Bound<Term>,
     },
+    Set {
+        field: Field,
+        value_type: Type,
+        elements: Vec<Term>,
+    },
     All,
 }
 
@@ -87,6 +92,27 @@ impl fmt::Debug for LogicalLiteral {
                 ref upper,
                 ..
             } => write!(formatter, "({:?} TO {:?})", lower, upper),
+            LogicalLiteral::Set { ref elements, .. } => {
+                const MAX_DISPLAYED: usize = 10;
+
+                write!(formatter, "IN [")?;
+                for (i, element) in elements.iter().enumerate() {
+                    if i == 0 {
+                        write!(formatter, "{:?}", element)?;
+                    } else if i == MAX_DISPLAYED - 1 {
+                        write!(
+                            formatter,
+                            ", {:?}, ... ({} more)",
+                            element,
+                            elements.len() - i - 1
+                        )?;
+                        break;
+                    } else {
+                        write!(formatter, ", {:?}", element)?;
+                    }
+                }
+                write!(formatter, "]")
+            }
             LogicalLiteral::All => write!(formatter, "*"),
         }
     }

src/query/query_parser/query_parser.rs

@@ -13,7 +13,7 @@ use crate::indexer::{
 };
 use crate::query::{
     AllQuery, BooleanQuery, BoostQuery, EmptyQuery, Occur, PhraseQuery, Query, RangeQuery,
-    TermQuery,
+    TermQuery, TermSetQuery,
 };
 use crate::schema::{
     Facet, FacetParseError, Field, FieldType, IndexRecordOption, IntoIpv6Addr, Schema, Term, Type,
@@ -685,6 +685,31 @@ impl QueryParser {
                 }));
                 Ok(logical_ast)
             }
+            UserInputLeaf::Set {
+                field: full_field_opt,
+                elements,
+            } => {
+                let full_path = full_field_opt.ok_or_else(|| {
+                    QueryParserError::UnsupportedQuery(
+                        "Set query need to target a specific field.".to_string(),
+                    )
+                })?;
+                let (field, json_path) = self
+                    .split_full_path(&full_path)
+                    .ok_or_else(|| QueryParserError::FieldDoesNotExist(full_path.clone()))?;
+                let field_entry = self.schema.get_field_entry(field);
+                let value_type = field_entry.field_type().value_type();
+                let logical_ast = LogicalAst::Leaf(Box::new(LogicalLiteral::Set {
+                    elements: elements
+                        .into_iter()
+                        .map(|element| self.compute_boundary_term(field, json_path, &element))
+                        .collect::<Result<Vec<_>, _>>()?,
+
+                    field,
+                    value_type,
+                }));
+                Ok(logical_ast)
+            }
         }
     }
 }
@@ -703,6 +728,7 @@ fn convert_literal_to_query(logical_literal: LogicalLiteral) -> Box<dyn Query> {
         } => Box::new(RangeQuery::new_term_bounds(
             field, value_type, &lower, &upper,
         )),
+        LogicalLiteral::Set { elements, .. } => Box::new(TermSetQuery::new(elements)),
         LogicalLiteral::All => Box::new(AllQuery),
     }
 }
@@ -1563,4 +1589,29 @@ mod test {
             false,
         );
     }
+
+    #[test]
+    pub fn test_term_set_query() {
+        test_parse_query_to_logical_ast_helper(
+            "title: IN [a b cd]",
+            r#"IN [Term(type=Str, field=0, "a"), Term(type=Str, field=0, "b"), Term(type=Str, field=0, "cd")]"#,
+            false,
+        );
+        test_parse_query_to_logical_ast_helper(
+            "bytes: IN [AA== ABA= ABCD]",
+            r#"IN [Term(type=Bytes, field=12, [0]), Term(type=Bytes, field=12, [0, 16]), Term(type=Bytes, field=12, [0, 16, 131])]"#,
+            false,
+        );
+        test_parse_query_to_logical_ast_helper(
+            "signed: IN [1 2 -3]",
+            r#"IN [Term(type=I64, field=2, 1), Term(type=I64, field=2, 2), Term(type=I64, field=2, -3)]"#,
+            false,
+        );
+
+        test_parse_query_to_logical_ast_helper(
+            "float: IN [1.1 2.2 -3.3]",
+            r#"IN [Term(type=F64, field=10, 1.1), Term(type=F64, field=10, 2.2), Term(type=F64, field=10, -3.3)]"#,
+            false,
+        );
+    }
 }

src/query/set_query.rs

@@ -101,9 +101,8 @@ impl Automaton for SetDfaWrapper {
 
 #[cfg(test)]
 mod tests {
-
     use crate::collector::TopDocs;
-    use crate::query::TermSetQuery;
+    use crate::query::{QueryParser, TermSetQuery};
     use crate::schema::{Schema, TEXT};
     use crate::{assert_nearly_equals, Index, Term};
 
@@ -215,4 +214,31 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_term_set_query_parser() -> crate::Result<()> {
+        let mut schema_builder = Schema::builder();
+        schema_builder.add_text_field("field", TEXT);
+        let schema = schema_builder.build();
+        let index = Index::create_in_ram(schema.clone());
+        let mut index_writer = index.writer_for_tests()?;
+        let field = schema.get_field("field").unwrap();
+        index_writer.add_document(doc!(
+          field => "val1",
+        ))?;
+        index_writer.add_document(doc!(
+          field => "val2",
+        ))?;
+        index_writer.add_document(doc!(
+          field => "val3",
+        ))?;
+        index_writer.commit()?;
+        let reader = index.reader()?;
+        let searcher = reader.searcher();
+        let query_parser = QueryParser::for_index(&index, vec![]);
+        let query = query_parser.parse_query("field: IN [val1 val2]")?;
+        let top_docs = searcher.search(&query, &TopDocs::with_limit(3))?;
+        assert_eq!(top_docs.len(), 2);
+        Ok(())
+    }
 }