Skip to content

Commit

Permalink
Support boost in Range and Bool query
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevchuang
kevchuang committed Feb 8, 2023
1 parent 1e01903 commit bed428e
Show file tree
Hide file tree
Showing 3 changed files with 164 additions and 22 deletions.
24 changes: 22 additions & 2 deletions modules/library/src/main/scala/zio/elasticsearch/Boost.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,17 @@

package zio.elasticsearch

import zio.elasticsearch.ElasticQuery.{ElasticPrimitive, MatchAllQuery, TermQuery, WildcardQuery}
import zio.elasticsearch.ElasticQueryType.{MatchAll, Term, Wildcard}
import zio.elasticsearch.ElasticQuery.{
BoolQuery,
ElasticPrimitive,
LowerBound,
MatchAllQuery,
RangeQuery,
TermQuery,
UpperBound,
WildcardQuery
}
import zio.elasticsearch.ElasticQueryType.{Bool, MatchAll, Range, Term, Wildcard}

object Boost {

Expand All @@ -26,11 +35,22 @@ object Boost {
}

object WithBoost {
implicit val boolWithBoost: WithBoost[Bool] = (query: ElasticQuery[Bool], value: Double) =>
query match {
case q: BoolQuery => q.copy(boost = Some(value))
}

implicit val matchAllWithBoost: WithBoost[MatchAll] = (query: ElasticQuery[MatchAll], value: Double) =>
query match {
case q: MatchAllQuery => q.copy(boost = Some(value))
}

implicit def rangeWithBoost[A, LB <: LowerBound, UB <: UpperBound]: WithBoost[Range[A, LB, UB]] =
(query: ElasticQuery[Range[A, LB, UB]], value: Double) =>
query match {
case q: RangeQuery[A, LB, UB] => q.copy(boost = Some(value))
}

implicit def termWithBoost[A: ElasticPrimitive]: WithBoost[Term[A]] =
(query: ElasticQuery[Term[A]], value: Double) =>
query match {
Expand Down
45 changes: 25 additions & 20 deletions modules/library/src/main/scala/zio/elasticsearch/ElasticQuery.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,17 +114,18 @@ object ElasticQuery {
private[elasticsearch] final case class BoolQuery(
filter: List[ElasticQuery[_]],
must: List[ElasticQuery[_]],
should: List[ElasticQuery[_]]
should: List[ElasticQuery[_]],
boost: Option[Double] = None
) extends ElasticQuery[Bool] { self =>

override def toJson: Json =
Obj(
"bool" -> Obj(
"filter" -> Arr(filter.map(_.toJson): _*),
"must" -> Arr(must.map(_.toJson): _*),
"should" -> Arr(should.map(_.toJson): _*)
)
)
override def toJson: Json = {
val boolFields =
Some("filter" -> Arr(filter.map(_.toJson): _*)) ++
Some("must" -> Arr(must.map(_.toJson): _*)) ++
Some("should" -> Arr(should.map(_.toJson): _*)) ++
boost.map("boost" -> Num(_))
Obj("bool" -> Obj(boolFields.toList: _*))
}

def filter(queries: ElasticQuery[_]*): BoolQuery =
self.copy(filter = filter ++ queries)
Expand Down Expand Up @@ -181,11 +182,12 @@ object ElasticQuery {
override def toJson: Option[(String, Json)] = None
}

private[elasticsearch] final case class RangeQuery[A, LB <: LowerBound, UB <: UpperBound] private (
private[elasticsearch] final case class RangeQuery[A, LB <: LowerBound, UB <: UpperBound](
field: String,
lower: LB,
upper: UB
) extends ElasticQuery[Range] { self =>
upper: UB,
boost: Option[Double] = None
) extends ElasticQuery[Range[A, LB, UB]] { self =>

def gt[B <: A: ElasticPrimitive](value: B)(implicit
@unused ev: LB =:= Unbounded.type
Expand All @@ -207,7 +209,10 @@ object ElasticQuery {
): RangeQuery[B, LB, LessThanOrEqualTo[B]] =
self.copy(upper = LessThanOrEqualTo(value))

override def toJson: Json = Obj("range" -> Obj(field -> Obj(List(lower.toJson, upper.toJson).flatten: _*)))
override def toJson: Json = {
val rangeFields = Some(field -> Obj(List(lower.toJson, upper.toJson).flatten: _*)) ++ boost.map("boost" -> Num(_))
Obj("range" -> Obj(rangeFields.toList: _*))
}
}

private[elasticsearch] object RangeQuery {
Expand Down Expand Up @@ -248,11 +253,11 @@ object ElasticQuery {
sealed trait ElasticQueryType

object ElasticQueryType {
trait Bool extends ElasticQueryType
trait Exists extends ElasticQueryType
trait Match extends ElasticQueryType
trait MatchAll extends ElasticQueryType
trait Range extends ElasticQueryType
trait Term[A] extends ElasticQueryType
trait Wildcard extends ElasticQueryType
trait Bool extends ElasticQueryType
trait Exists extends ElasticQueryType
trait Match extends ElasticQueryType
trait MatchAll extends ElasticQueryType
trait Range[A, LB, UB] extends ElasticQueryType
trait Term[A] extends ElasticQueryType
trait Wildcard extends ElasticQueryType
}
117 changes: 117 additions & 0 deletions modules/library/src/test/scala/zio/elasticsearch/QueryDSLSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,11 @@ object QueryDSLSpec extends ZIOSpecDefault {

assert(query)(equalTo(MatchQuery(field = "name.keyword", value = "Name")))
},
test("successfully create Bool Query with boost") {
val query = boolQuery().boost(1.0)

assert(query)(equalTo(BoolQuery(filter = Nil, must = Nil, should = Nil, boost = Some(1.0))))
},
test("successfully create `Filter` query from two Match queries") {
val query = boolQuery()
.filter(
Expand Down Expand Up @@ -182,6 +187,24 @@ object QueryDSLSpec extends ZIOSpecDefault {
)
)
},
test("successfully create `Filter/Must/Should` mixed query with boost") {
val query = boolQuery()
.filter(matches(field = "customer_id", value = 1))
.must(matches(field = "customer_age", value = 23))
.should(matches(field = "day_of_week", value = "Monday"))
.boost(1.0)

assert(query)(
equalTo(
BoolQuery(
filter = List(MatchQuery(field = "customer_id", value = 1)),
must = List(MatchQuery(field = "customer_age", value = 23)),
should = List(MatchQuery(field = "day_of_week", value = "Monday")),
boost = Some(1.0)
)
)
)
},
test("successfully create Exists Query") {
val query = exists(field = "day_of_week")

Expand Down Expand Up @@ -421,6 +444,24 @@ object QueryDSLSpec extends ZIOSpecDefault {

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Bool Query with boost") {
val query = boolQuery().boost(1.0)
val expected =
"""
|{
| "query": {
| "bool": {
| "filter": [],
| "must": [],
| "should": [],
| "boost": 1.0
| }
| }
|}
|""".stripMargin

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Bool Query with Filter containing `Match` leaf query") {
val query = boolQuery().filter(matches(field = "day_of_week", value = "Monday"))
val expected =
Expand Down Expand Up @@ -528,6 +569,46 @@ object QueryDSLSpec extends ZIOSpecDefault {

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Bool Query with Filter, Must and Should containing `Match` leaf query and with boost") {
val query = boolQuery()
.filter(matches(field = "customer_age", value = 23))
.must(matches(field = "customer_id", value = 1))
.should(matches(field = "day_of_week", value = "Monday"))
.boost(1.0)
val expected =
"""
|{
| "query": {
| "bool": {
| "filter": [
| {
| "match": {
| "customer_age": 23
| }
| }
| ],
| "must": [
| {
| "match": {
| "customer_id": 1
| }
| }
| ],
| "should": [
| {
| "match": {
| "day_of_week": "Monday"
| }
| }
| ],
| "boost": 1.0
| }
| }
|}
|""".stripMargin

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Exists Query") {
val query = exists(field = "day_of_week")
val expected =
Expand Down Expand Up @@ -587,6 +668,23 @@ object QueryDSLSpec extends ZIOSpecDefault {

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Unbounded Range Query with boost") {
val query = range(field = "field").boost(1.0)
val expected =
"""
|{
| "query": {
| "range": {
| "field": {
| },
| "boost": 1.0
| }
| }
|}
|""".stripMargin

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Range Query with Lower Bound") {
val query = range(field = "customer_age").gt(23)
val expected =
Expand Down Expand Up @@ -673,6 +771,25 @@ object QueryDSLSpec extends ZIOSpecDefault {

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Range Query with both Upper and Lower Bound with boost") {
val query = range(field = "customer_age").gte(10).lt(100).boost(1.0)
val expected =
"""
|{
| "query": {
| "range": {
| "customer_age": {
| "gte": 10,
| "lt": 100
| },
| "boost": 1.0
| }
| }
|}
|""".stripMargin

assert(query.toJsonBody)(equalTo(expected.toJson))
},
test("properly encode Term query") {
val query = term(field = "day_of_week", value = true)
val expected =
Expand Down

0 comments on commit bed428e

Please sign in to comment.