diff --git a/search-client/src/commonTest/kotlin/com/jillesvangurp/ktsearch/SearchTest.kt b/search-client/src/commonTest/kotlin/com/jillesvangurp/ktsearch/SearchTest.kt index d44555f5..fe491b17 100644 --- a/search-client/src/commonTest/kotlin/com/jillesvangurp/ktsearch/SearchTest.kt +++ b/search-client/src/commonTest/kotlin/com/jillesvangurp/ktsearch/SearchTest.kt @@ -217,6 +217,44 @@ class SearchTest : SearchTestBase() { response.responses shouldHaveSize 2 } + @Test + fun shouldApplyRescore() = coRun { + val indexName = testDocumentIndex() + client.bulk(target = indexName, refresh = Refresh.WaitFor) { + index(TestDocument("doc 1", tags = listOf("rescore")).json()) + index(TestDocument("doc 2", tags = listOf("nope")).json()) + index(TestDocument("doc 3", tags = listOf("another")).json()) + } + val response = client.search(indexName, explain = true) { + query = matchAll() + val firstRescoreQuery = constantScore { + filter = match(TestDocument::tags, "rescore") + } + + val secondRescoreQuery = constantScore { + filter = match(TestDocument::tags, "another") + } + rescore( + rescorer(3) { + scoreMode = RescoreScoreMode.total + rescoreQueryWeight = 20.0 + queryWeight = 2.0 + rescoreQuery = firstRescoreQuery + }, + rescorer(3) { + scoreMode = RescoreScoreMode.multiply + rescoreQueryWeight = 5.0 + queryWeight = 1.0 + rescoreQuery = secondRescoreQuery + } + ) + } + + response.hits!!.hits shouldHaveSize 3 + response.parseHits().map(TestDocument::name) shouldBe listOf("doc 1", "doc 3", "doc 2") + response.hits!!.hits.map(SearchResponse.Hit::score) shouldBe listOf(22.0, 10.0, 2.0) + } + @Test fun shouldExposeSeqNo() = coRun { val indexName = testDocumentIndex() diff --git a/search-dsls/src/commonMain/kotlin/com/jillesvangurp/searchdsls/querydsl/search-dsl.kt b/search-dsls/src/commonMain/kotlin/com/jillesvangurp/searchdsls/querydsl/search-dsl.kt index cdf5aff6..3838db98 100644 --- a/search-dsls/src/commonMain/kotlin/com/jillesvangurp/searchdsls/querydsl/search-dsl.kt +++ b/search-dsls/src/commonMain/kotlin/com/jillesvangurp/searchdsls/querydsl/search-dsl.kt @@ -255,3 +255,80 @@ fun SearchDSL.dotted(vararg elements: Any) = elements.joinToString(".") { pathCo else -> pathComponent.toString() } } + +fun SearchDSL.rescore(vararg rescores: Rescorer) = getOrCreateMutableList("rescore").addAll(rescores) + +/** + * Rescoring can help to improve precision by reordering just the top (e.g. 100 - 500) documents + * returned by the query and post_filter phases, using a secondary (usually more costly) algorithm, + * instead of applying the costly algorithm to all documents in the index. + */ +class Rescorer : JsonDsl() { + /** + * The number of docs which will be examined on each shard + */ + var windowSize by property() + + /** + * Second query excuted only on the Top-K results returned by the query and post_filter phases. + */ + var query by property() +} + +class RescoreQuery : JsonDsl() { + /** + * Query to apply + */ + var rescoreQuery by esQueryProperty() + + /** + * The relative importance of the original query + */ + var queryWeight by property() + + /** + * The relative importance of the rescore query + */ + var rescoreQueryWeight by property() + + /** + * way the original score and rescore score are combined + */ + var scoreMode by property() +} + +/** + * Controls the way the original score and rescore score are combined + */ +enum class RescoreScoreMode { + /** + * Average the original score and the rescore query score. + */ + avg, + + /** + * Take the min of the original score and the rescore query score. + */ + min, + + /** + * Take the max of original score and the rescore query score. + */ + max, + + /** + * Add the original score and the rescore query score. The default. + */ + total, + + /** + * Multiply the original score by the rescore query score. + */ + multiply +} + +fun SearchDSL.rescorer(windowSize: Int, queryBlock: RescoreQuery.() -> Unit) = + Rescorer().apply { + this.windowSize = windowSize + this.query = RescoreQuery().apply(queryBlock) + }