Skip to content

Commit

Permalink
Implemented Elasticsearch's TermsQuery (#3747)
Browse files Browse the repository at this point in the history
Closes #3746
  • Loading branch information
fulmicoton committed Aug 15, 2023
1 parent d2afdbd commit e7e4ff1
Show file tree
Hide file tree
Showing 4 changed files with 155 additions and 3 deletions.
4 changes: 4 additions & 0 deletions quickwit/quickwit-query/src/elastic_query_dsl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ mod phrase_prefix_query;
mod query_string_query;
mod range_query;
mod term_query;
mod terms_query;

use bool_query::BoolQuery;
pub use one_field_map::OneFieldMap;
Expand All @@ -41,6 +42,7 @@ use crate::elastic_query_dsl::exists_query::ExistsQuery;
use crate::elastic_query_dsl::match_phrase_query::MatchPhraseQuery;
use crate::elastic_query_dsl::match_query::MatchQuery;
use crate::elastic_query_dsl::multi_match::MultiMatchQuery;
use crate::elastic_query_dsl::terms_query::TermsQuery;
use crate::not_nan_f32::NotNaNf32;
use crate::query_ast::QueryAst;

Expand All @@ -58,6 +60,7 @@ pub(crate) enum ElasticQueryDslInner {
QueryString(QueryStringQuery),
Bool(BoolQuery),
Term(TermQuery),
Terms(TermsQuery),
MatchAll(MatchAllQuery),
MatchNone(MatchNoneQuery),
Match(MatchQuery),
Expand Down Expand Up @@ -90,6 +93,7 @@ impl ConvertableToQueryAst for ElasticQueryDslInner {
Self::QueryString(query_string_query) => query_string_query.convert_to_query_ast(),
Self::Bool(bool_query) => bool_query.convert_to_query_ast(),
Self::Term(term_query) => term_query.convert_to_query_ast(),
Self::Terms(terms_query) => terms_query.convert_to_query_ast(),
Self::MatchAll(match_all_query) => {
if let Some(boost) = match_all_query.boost {
Ok(QueryAst::Boost {
Expand Down
4 changes: 1 addition & 3 deletions quickwit/quickwit-query/src/elastic_query_dsl/term_query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ pub struct TermQueryValue {
pub boost: Option<NotNaNf32>,
}

#[cfg(test)]
pub fn term_query_from_field_value(field: impl ToString, value: impl ToString) -> TermQuery {
TermQuery {
field: field.to_string(),
Expand Down Expand Up @@ -65,8 +64,7 @@ impl ConvertableToQueryAst for TermQuery {

#[cfg(test)]
mod tests {
use super::TermQuery;
use crate::elastic_query_dsl::term_query::term_query_from_field_value;
use super::*;

#[test]
fn test_term_query_simple() {
Expand Down
129 changes: 129 additions & 0 deletions quickwit/quickwit-query/src/elastic_query_dsl/terms_query.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
// Copyright (C) 2023 Quickwit, Inc.
//
// Quickwit is offered under the AGPL v3.0 and as commercial software.
// For commercial licensing, contact us at hello@quickwit.io.
//
// AGPL:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use serde::Deserialize;

use crate::elastic_query_dsl::bool_query::BoolQuery;
use crate::elastic_query_dsl::one_field_map::OneFieldMap;
use crate::elastic_query_dsl::term_query::term_query_from_field_value;
use crate::elastic_query_dsl::{ConvertableToQueryAst, ElasticQueryDslInner};
use crate::not_nan_f32::NotNaNf32;
use crate::query_ast::QueryAst;

#[derive(PartialEq, Eq, Debug, Deserialize, Clone)]
#[serde(try_from = "TermsQueryForSerialization")]
pub struct TermsQuery {
pub boost: Option<NotNaNf32>,
pub field: String,
pub values: Vec<String>,
}

#[derive(Deserialize)]
struct TermsQueryForSerialization {
#[serde(default)]
boost: Option<NotNaNf32>,
#[serde(flatten)]
capture_other: serde_json::Value,
}

#[derive(Deserialize)]
#[serde(untagged)]
enum OneOrMany {
One(String),
Many(Vec<String>),
}
impl From<OneOrMany> for Vec<String> {
fn from(one_or_many: OneOrMany) -> Vec<String> {
match one_or_many {
OneOrMany::One(one_value) => vec![one_value],
OneOrMany::Many(values) => values,
}
}
}

impl TryFrom<TermsQueryForSerialization> for TermsQuery {
type Error = serde_json::Error;

fn try_from(value: TermsQueryForSerialization) -> serde_json::Result<TermsQuery> {
let one_field: OneFieldMap<OneOrMany> = serde_json::from_value(value.capture_other)?;
let one_field_values: Vec<String> = one_field.value.into();
Ok(TermsQuery {
boost: value.boost,
field: one_field.field,
values: one_field_values,
})
}
}

impl ConvertableToQueryAst for TermsQuery {
fn convert_to_query_ast(self) -> anyhow::Result<QueryAst> {
let term_queries: Vec<ElasticQueryDslInner> = self
.values
.into_iter()
.map(|value| term_query_from_field_value(self.field.clone(), value))
.map(ElasticQueryDslInner::from)
.collect();
let mut union = BoolQuery::union(term_queries);
union.boost = self.boost;
union.convert_to_query_ast()
}
}

impl From<TermsQuery> for ElasticQueryDslInner {
fn from(term_query: TermsQuery) -> Self {
Self::Terms(term_query)
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_terms_query_simple() {
let terms_query_json = r#"{ "user.id": ["hello", "happy"] }"#;
let terms_query: TermsQuery = serde_json::from_str(terms_query_json).unwrap();
assert_eq!(&terms_query.field, "user.id");
assert_eq!(
&terms_query.values[..],
&["hello".to_string(), "happy".to_string()]
);
}

#[test]
fn test_terms_query_single_term_not_array() {
let terms_query_json = r#"{ "user.id": "hello"}"#;
let terms_query: TermsQuery = serde_json::from_str(terms_query_json).unwrap();
assert_eq!(&terms_query.field, "user.id");
assert_eq!(&terms_query.values[..], &["hello".to_string()]);
}

#[test]
fn test_terms_query_single_term_boost() {
let terms_query_json = r#"{ "user.id": ["hello", "happy"], "boost": 2 }"#;
let terms_query: TermsQuery = serde_json::from_str(terms_query_json).unwrap();
assert_eq!(&terms_query.field, "user.id");
assert_eq!(
&terms_query.values[..],
&["hello".to_string(), "happy".to_string()]
);
let boost: f32 = terms_query.boost.unwrap().into();
assert!((boost - 2.0f32).abs() < 0.0001f32);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
json:
query:
terms:
type:
- PushEvent
- CommitCommentEvent
expected:
hits:
total:
value: 0
---
json:
query:
terms:
type:
- pushevent
- commitcommentevent
expected:
hits:
total:
value: 61

0 comments on commit e7e4ff1

Please sign in to comment.