├── .gitignore ├── LICENSE.txt ├── NOTICE.txt ├── README.md ├── VERSION.txt ├── build.gradle ├── docker-compose.yml ├── gradle.properties ├── gradle └── wrapper │ ├── gradle-wrapper.jar │ └── gradle-wrapper.properties ├── gradlew ├── gradlew.bat ├── settings.gradle └── src ├── main ├── java │ └── com │ │ └── bigdataboutique │ │ └── elasticsearch │ │ └── plugin │ │ ├── Config.java │ │ ├── RedisRescoreBuilder.java │ │ └── RedisRescorePlugin.java └── plugin-metadata │ └── plugin-security.policy └── test ├── java └── com │ └── bigdataboutique │ └── elasticsearch │ └── plugin │ ├── ConfigTests.java │ ├── RedisRescoreBuilderTests.java │ └── RedisRescoreClientYamlTestSuiteIT.java └── resources └── rest-api-spec └── test └── redis-rescore ├── 10_basic.yml ├── 10_config.yml └── 20_score.yml /.gitignore: -------------------------------------------------------------------------------- 1 | .gradle/ 2 | build/ 3 | .idea/ 4 | 5 | # ignore all the files created by java extension pack vscode extension 6 | bin 7 | .classpath 8 | .project 9 | .settings 10 | .vscode 11 | out 12 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Source code in this repository is covered by one of three licenses: (i) the 2 | Apache License 2.0 (ii) an Apache License 2.0 compatible license (iii) the 3 | Elastic License. The default license throughout the repository is Apache License 4 | 2.0 unless the header specifies another license. Elastic Licensed code is found 5 | only in the x-pack directory. 6 | 7 | The build produces two sets of binaries - one set that falls under the Elastic 8 | License and another set that falls under Apache License 2.0. The binaries that 9 | contain `-oss` in the artifact name are licensed under Apache License 2.0 and 10 | these binaries do not package any code from the x-pack directory. 11 | -------------------------------------------------------------------------------- /NOTICE.txt: -------------------------------------------------------------------------------- 1 | Elasticsearch Redis Rescore plugin 2 | Copyright 2020 BigData Boutique 3 | 4 | This product includes software developed by The Apache Software 5 | Foundation (http://www.apache.org/). -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Elasticsearch Redis Rescore Plugin 2 | 3 | Ever wanted to use data from external systems to influence scoring in Elasticsearch? now you can. 4 | 5 | The idea is simple: query a low-latency service per document during search and get a numeric value for that specific document, then use that number as a multiplier for influencing the document score, up or down. 6 | 7 | This plugin uses Redis to rescore top ranking results in Elasticsearch. It's your job to make sure the query gets the basic filtering and scoring right. Then, you can use the plugin to [rescore](https://www.elastic.co/guide/en/elasticsearch/reference/7.5/search-request-body.html#request-body-search-rescore) the top N results. 8 | 9 | Rescoring with this plugin assumes Redis contains keys that can be correlated with the data in the documents, such that for every doc D there exists a value in a predefined field that also exists in Redis as a key whose value is numeric. 10 | 11 | Documents that do not contain this field, or no value in this field, or that value does not exist in Redis as a key - are left untouched. 12 | 13 | See an example below. 14 | 15 | ## Installation 16 | 17 | Follow the standard plugin installation instructions, with a zip version of the compiled plugin: [https://www.elastic.co/guide/en/elasticsearch/plugins/current/installation.html](https://www.elastic.co/guide/en/elasticsearch/plugins/current/installation.html) 18 | 19 | ## Usage 20 | 21 | ```json 22 | { 23 | "query": { "match_all": {} }, 24 | "rescore": { 25 | "redis":{ 26 | "key_field": "productId.keyword", 27 | "key_prefix": "mystore-" 28 | } 29 | } 30 | } 31 | ``` 32 | 33 | In this example, we are expecting each hit to contain a field `productId` (of keyword type). The value of that field will be looked up in Redis as a key (for example, Redis key `mystore-abc123` will be looked-up for a document with productId abc123; the `mystore-` key prefix is configurable in query time). 34 | 35 | The value which will be found under that Redis key, if exists and of numeric type, will be multiplied by the current document score to produce a new score. 36 | 37 | You can use `0` to demote results (e.g. mark as unavailable in stock), `1` to leave unchanged, or any other value to produce positive or negative boosts based on your business logic. -------------------------------------------------------------------------------- /VERSION.txt: -------------------------------------------------------------------------------- 1 | 7.4.2 -------------------------------------------------------------------------------- /build.gradle: -------------------------------------------------------------------------------- 1 | buildscript { 2 | ext { 3 | elasticsearchVersion = rootProject.file('VERSION.txt').text.trim() 4 | } 5 | repositories { 6 | jcenter() 7 | } 8 | dependencies { 9 | classpath "com.avast.gradle:gradle-docker-compose-plugin:0.10.7" 10 | classpath "org.elasticsearch.gradle:build-tools:${rootProject.file('VERSION.txt').text.trim()}" 11 | } 12 | } 13 | 14 | apply plugin: 'docker-compose' 15 | apply plugin: 'java' 16 | apply plugin: 'idea' 17 | apply plugin: 'elasticsearch.esplugin' 18 | 19 | group 'com.bigdataboutique' 20 | version '1.0-SNAPSHOT' 21 | description 'Elasticsearch scoring plugin.' 22 | 23 | licenseFile = rootProject.file('LICENSE.txt') 24 | noticeFile = rootProject.file('NOTICE.txt') 25 | 26 | version = rootProject.file('VERSION.txt').text.trim() 27 | 28 | sourceCompatibility = 1.8 29 | 30 | ext { 31 | elasticsearch = rootProject.file('VERSION.txt').text.trim() 32 | } 33 | 34 | repositories { 35 | jcenter() 36 | } 37 | 38 | configurations.all { 39 | resolutionStrategy { 40 | force 'org.slf4j:slf4j-api:1.7.21' 41 | } 42 | } 43 | 44 | dependencies { 45 | compile group: 'org.elasticsearch', name: 'elasticsearch', version: elasticsearch 46 | compile group: 'redis.clients', name: 'jedis', version: '3.2.0' 47 | testCompile group: 'org.elasticsearch.test', name: 'framework', version: elasticsearch 48 | // compileClasspath group: 'org.codelibs.elasticsearch.lib', name: 'plugin-classloader', version: elasticsearch 49 | // compile group: 'org.apache.logging.log4j', name: 'log4j-core', version: '2.11.0' 50 | testCompile group: 'junit', name: 'junit', version: '4.12' 51 | testCompile group: 'org.slf4j', name: 'slf4j-simple', version: '1.7.30' 52 | testImplementation('org.testcontainers:testcontainers:1.12.5') { 53 | exclude group: 'junit', module: 'junit' 54 | exclude group: 'net.java.dev.jna', module: 'jna' 55 | } 56 | } 57 | 58 | esplugin { 59 | classname 'com.bigdataboutique.elasticsearch.plugin.RedisRescorePlugin' 60 | name project.name 61 | version project.version 62 | description project.description 63 | } 64 | 65 | dependencyLicenses.enabled = false 66 | thirdPartyAudit.enabled = false 67 | licenseHeaders.enabled = false 68 | 69 | dockerCompose.isRequiredBy(integTestRunner) 70 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: "3.3" 2 | services: 3 | redis: 4 | image: redis:3.2.0 5 | ports: 6 | - "6379:6379" 7 | network_mode: "host" -------------------------------------------------------------------------------- /gradle.properties: -------------------------------------------------------------------------------- 1 | org.gradle.daemon=true 2 | org.gradle.jvmargs=-Xmx2g 3 | options.forkOptions.memoryMaximumSize=2g 4 | -------------------------------------------------------------------------------- /gradle/wrapper/gradle-wrapper.jar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BigDataBoutique/elasticsearch-rescore-redis/7b75a9f4808cdacb4f079ff0147c33a3da791e34/gradle/wrapper/gradle-wrapper.jar -------------------------------------------------------------------------------- /gradle/wrapper/gradle-wrapper.properties: -------------------------------------------------------------------------------- 1 | #Fri Jan 17 14:33:46 IST 2020 2 | distributionUrl=https\://services.gradle.org/distributions/gradle-6.0.1-all.zip 3 | distributionBase=GRADLE_USER_HOME 4 | distributionPath=wrapper/dists 5 | zipStorePath=wrapper/dists 6 | zipStoreBase=GRADLE_USER_HOME 7 | -------------------------------------------------------------------------------- /gradlew: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env sh 2 | 3 | # 4 | # Copyright 2015 the original author or authors. 5 | # 6 | # Licensed under the Apache License, Version 2.0 (the "License"); 7 | # you may not use this file except in compliance with the License. 8 | # You may obtain a copy of the License at 9 | # 10 | # https://www.apache.org/licenses/LICENSE-2.0 11 | # 12 | # Unless required by applicable law or agreed to in writing, software 13 | # distributed under the License is distributed on an "AS IS" BASIS, 14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 15 | # See the License for the specific language governing permissions and 16 | # limitations under the License. 17 | # 18 | 19 | ############################################################################## 20 | ## 21 | ## Gradle start up script for UN*X 22 | ## 23 | ############################################################################## 24 | 25 | # Attempt to set APP_HOME 26 | # Resolve links: $0 may be a link 27 | PRG="$0" 28 | # Need this for relative symlinks. 29 | while [ -h "$PRG" ] ; do 30 | ls=`ls -ld "$PRG"` 31 | link=`expr "$ls" : '.*-> \(.*\)$'` 32 | if expr "$link" : '/.*' > /dev/null; then 33 | PRG="$link" 34 | else 35 | PRG=`dirname "$PRG"`"/$link" 36 | fi 37 | done 38 | SAVED="`pwd`" 39 | cd "`dirname \"$PRG\"`/" >/dev/null 40 | APP_HOME="`pwd -P`" 41 | cd "$SAVED" >/dev/null 42 | 43 | APP_NAME="Gradle" 44 | APP_BASE_NAME=`basename "$0"` 45 | 46 | # Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. 47 | DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"' 48 | 49 | # Use the maximum available, or set MAX_FD != -1 to use that value. 50 | MAX_FD="maximum" 51 | 52 | warn () { 53 | echo "$*" 54 | } 55 | 56 | die () { 57 | echo 58 | echo "$*" 59 | echo 60 | exit 1 61 | } 62 | 63 | # OS specific support (must be 'true' or 'false'). 64 | cygwin=false 65 | msys=false 66 | darwin=false 67 | nonstop=false 68 | case "`uname`" in 69 | CYGWIN* ) 70 | cygwin=true 71 | ;; 72 | Darwin* ) 73 | darwin=true 74 | ;; 75 | MINGW* ) 76 | msys=true 77 | ;; 78 | NONSTOP* ) 79 | nonstop=true 80 | ;; 81 | esac 82 | 83 | CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar 84 | 85 | # Determine the Java command to use to start the JVM. 86 | if [ -n "$JAVA_HOME" ] ; then 87 | if [ -x "$JAVA_HOME/jre/sh/java" ] ; then 88 | # IBM's JDK on AIX uses strange locations for the executables 89 | JAVACMD="$JAVA_HOME/jre/sh/java" 90 | else 91 | JAVACMD="$JAVA_HOME/bin/java" 92 | fi 93 | if [ ! -x "$JAVACMD" ] ; then 94 | die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME 95 | 96 | Please set the JAVA_HOME variable in your environment to match the 97 | location of your Java installation." 98 | fi 99 | else 100 | JAVACMD="java" 101 | which java >/dev/null 2>&1 || die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 102 | 103 | Please set the JAVA_HOME variable in your environment to match the 104 | location of your Java installation." 105 | fi 106 | 107 | # Increase the maximum file descriptors if we can. 108 | if [ "$cygwin" = "false" -a "$darwin" = "false" -a "$nonstop" = "false" ] ; then 109 | MAX_FD_LIMIT=`ulimit -H -n` 110 | if [ $? -eq 0 ] ; then 111 | if [ "$MAX_FD" = "maximum" -o "$MAX_FD" = "max" ] ; then 112 | MAX_FD="$MAX_FD_LIMIT" 113 | fi 114 | ulimit -n $MAX_FD 115 | if [ $? -ne 0 ] ; then 116 | warn "Could not set maximum file descriptor limit: $MAX_FD" 117 | fi 118 | else 119 | warn "Could not query maximum file descriptor limit: $MAX_FD_LIMIT" 120 | fi 121 | fi 122 | 123 | # For Darwin, add options to specify how the application appears in the dock 124 | if $darwin; then 125 | GRADLE_OPTS="$GRADLE_OPTS \"-Xdock:name=$APP_NAME\" \"-Xdock:icon=$APP_HOME/media/gradle.icns\"" 126 | fi 127 | 128 | # For Cygwin or MSYS, switch paths to Windows format before running java 129 | if [ "$cygwin" = "true" -o "$msys" = "true" ] ; then 130 | APP_HOME=`cygpath --path --mixed "$APP_HOME"` 131 | CLASSPATH=`cygpath --path --mixed "$CLASSPATH"` 132 | JAVACMD=`cygpath --unix "$JAVACMD"` 133 | 134 | # We build the pattern for arguments to be converted via cygpath 135 | ROOTDIRSRAW=`find -L / -maxdepth 1 -mindepth 1 -type d 2>/dev/null` 136 | SEP="" 137 | for dir in $ROOTDIRSRAW ; do 138 | ROOTDIRS="$ROOTDIRS$SEP$dir" 139 | SEP="|" 140 | done 141 | OURCYGPATTERN="(^($ROOTDIRS))" 142 | # Add a user-defined pattern to the cygpath arguments 143 | if [ "$GRADLE_CYGPATTERN" != "" ] ; then 144 | OURCYGPATTERN="$OURCYGPATTERN|($GRADLE_CYGPATTERN)" 145 | fi 146 | # Now convert the arguments - kludge to limit ourselves to /bin/sh 147 | i=0 148 | for arg in "$@" ; do 149 | CHECK=`echo "$arg"|egrep -c "$OURCYGPATTERN" -` 150 | CHECK2=`echo "$arg"|egrep -c "^-"` ### Determine if an option 151 | 152 | if [ $CHECK -ne 0 ] && [ $CHECK2 -eq 0 ] ; then ### Added a condition 153 | eval `echo args$i`=`cygpath --path --ignore --mixed "$arg"` 154 | else 155 | eval `echo args$i`="\"$arg\"" 156 | fi 157 | i=`expr $i + 1` 158 | done 159 | case $i in 160 | 0) set -- ;; 161 | 1) set -- "$args0" ;; 162 | 2) set -- "$args0" "$args1" ;; 163 | 3) set -- "$args0" "$args1" "$args2" ;; 164 | 4) set -- "$args0" "$args1" "$args2" "$args3" ;; 165 | 5) set -- "$args0" "$args1" "$args2" "$args3" "$args4" ;; 166 | 6) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" ;; 167 | 7) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" ;; 168 | 8) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" ;; 169 | 9) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" "$args8" ;; 170 | esac 171 | fi 172 | 173 | # Escape application args 174 | save () { 175 | for i do printf %s\\n "$i" | sed "s/'/'\\\\''/g;1s/^/'/;\$s/\$/' \\\\/" ; done 176 | echo " " 177 | } 178 | APP_ARGS=`save "$@"` 179 | 180 | # Collect all arguments for the java command, following the shell quoting and substitution rules 181 | eval set -- $DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS "\"-Dorg.gradle.appname=$APP_BASE_NAME\"" -classpath "\"$CLASSPATH\"" org.gradle.wrapper.GradleWrapperMain "$APP_ARGS" 182 | 183 | exec "$JAVACMD" "$@" 184 | -------------------------------------------------------------------------------- /gradlew.bat: -------------------------------------------------------------------------------- 1 | @rem 2 | @rem Copyright 2015 the original author or authors. 3 | @rem 4 | @rem Licensed under the Apache License, Version 2.0 (the "License"); 5 | @rem you may not use this file except in compliance with the License. 6 | @rem You may obtain a copy of the License at 7 | @rem 8 | @rem https://www.apache.org/licenses/LICENSE-2.0 9 | @rem 10 | @rem Unless required by applicable law or agreed to in writing, software 11 | @rem distributed under the License is distributed on an "AS IS" BASIS, 12 | @rem WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | @rem See the License for the specific language governing permissions and 14 | @rem limitations under the License. 15 | @rem 16 | 17 | @if "%DEBUG%" == "" @echo off 18 | @rem ########################################################################## 19 | @rem 20 | @rem Gradle startup script for Windows 21 | @rem 22 | @rem ########################################################################## 23 | 24 | @rem Set local scope for the variables with windows NT shell 25 | if "%OS%"=="Windows_NT" setlocal 26 | 27 | set DIRNAME=%~dp0 28 | if "%DIRNAME%" == "" set DIRNAME=. 29 | set APP_BASE_NAME=%~n0 30 | set APP_HOME=%DIRNAME% 31 | 32 | @rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. 33 | set DEFAULT_JVM_OPTS="-Xmx64m" "-Xms64m" 34 | 35 | @rem Find java.exe 36 | if defined JAVA_HOME goto findJavaFromJavaHome 37 | 38 | set JAVA_EXE=java.exe 39 | %JAVA_EXE% -version >NUL 2>&1 40 | if "%ERRORLEVEL%" == "0" goto init 41 | 42 | echo. 43 | echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 44 | echo. 45 | echo Please set the JAVA_HOME variable in your environment to match the 46 | echo location of your Java installation. 47 | 48 | goto fail 49 | 50 | :findJavaFromJavaHome 51 | set JAVA_HOME=%JAVA_HOME:"=% 52 | set JAVA_EXE=%JAVA_HOME%/bin/java.exe 53 | 54 | if exist "%JAVA_EXE%" goto init 55 | 56 | echo. 57 | echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 58 | echo. 59 | echo Please set the JAVA_HOME variable in your environment to match the 60 | echo location of your Java installation. 61 | 62 | goto fail 63 | 64 | :init 65 | @rem Get command-line arguments, handling Windows variants 66 | 67 | if not "%OS%" == "Windows_NT" goto win9xME_args 68 | 69 | :win9xME_args 70 | @rem Slurp the command line arguments. 71 | set CMD_LINE_ARGS= 72 | set _SKIP=2 73 | 74 | :win9xME_args_slurp 75 | if "x%~1" == "x" goto execute 76 | 77 | set CMD_LINE_ARGS=%* 78 | 79 | :execute 80 | @rem Setup the command line 81 | 82 | set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar 83 | 84 | @rem Execute Gradle 85 | "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS% 86 | 87 | :end 88 | @rem End local scope for the variables with windows NT shell 89 | if "%ERRORLEVEL%"=="0" goto mainEnd 90 | 91 | :fail 92 | rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of 93 | rem the _cmd.exe /c_ return code! 94 | if not "" == "%GRADLE_EXIT_CONSOLE%" exit 1 95 | exit /b 1 96 | 97 | :mainEnd 98 | if "%OS%"=="Windows_NT" endlocal 99 | 100 | :omega 101 | -------------------------------------------------------------------------------- /settings.gradle: -------------------------------------------------------------------------------- 1 | rootProject.name = 'elasticsearch-rescore-redis' 2 | 3 | -------------------------------------------------------------------------------- /src/main/java/com/bigdataboutique/elasticsearch/plugin/Config.java: -------------------------------------------------------------------------------- 1 | 2 | package com.bigdataboutique.elasticsearch.plugin; 3 | 4 | import org.elasticsearch.common.settings.Setting; 5 | import org.elasticsearch.common.settings.Setting.Property; 6 | import org.elasticsearch.common.settings.Settings; 7 | 8 | import java.net.URI; 9 | /** 10 | * {@link Config} contains the settings values and their static declarations. 11 | */ 12 | public class Config { 13 | static final Setting REDIS_URL = new Setting("redisRescore.redisUrl", "localhost", v -> { 14 | try { 15 | new URI(v); 16 | return v; 17 | } catch (Exception e) { 18 | throw new IllegalArgumentException("Setting is not a valid URI"); 19 | } 20 | }, Property.NodeScope, Property.Dynamic); 21 | 22 | private final String redisUrl; 23 | 24 | public Config(final Settings settings) { 25 | this.redisUrl = REDIS_URL.get(settings); 26 | } 27 | 28 | public String getRedisUrl() { 29 | return redisUrl; 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /src/main/java/com/bigdataboutique/elasticsearch/plugin/RedisRescoreBuilder.java: -------------------------------------------------------------------------------- 1 | package com.bigdataboutique.elasticsearch.plugin; 2 | 3 | import org.apache.logging.log4j.LogManager; 4 | import org.apache.logging.log4j.Logger; 5 | import org.apache.lucene.index.LeafReaderContext; 6 | import org.apache.lucene.index.SortedNumericDocValues; 7 | import org.apache.lucene.index.SortedSetDocValues; 8 | import org.apache.lucene.search.Explanation; 9 | import org.apache.lucene.search.IndexSearcher; 10 | import org.apache.lucene.search.TopDocs; 11 | import org.elasticsearch.common.Nullable; 12 | import org.elasticsearch.common.ParseField; 13 | import org.elasticsearch.common.io.stream.StreamInput; 14 | import org.elasticsearch.common.io.stream.StreamOutput; 15 | import org.elasticsearch.common.xcontent.ConstructingObjectParser; 16 | import org.elasticsearch.common.xcontent.XContentBuilder; 17 | import org.elasticsearch.common.xcontent.XContentParser; 18 | import org.elasticsearch.index.fielddata.AtomicFieldData; 19 | import org.elasticsearch.index.fielddata.AtomicNumericFieldData; 20 | import org.elasticsearch.index.fielddata.IndexFieldData; 21 | import org.elasticsearch.index.fielddata.plain.SortedSetDVBytesAtomicFieldData; 22 | import org.elasticsearch.index.query.QueryRewriteContext; 23 | import org.elasticsearch.index.query.QueryShardContext; 24 | import org.elasticsearch.search.rescore.RescoreContext; 25 | import org.elasticsearch.search.rescore.Rescorer; 26 | import org.elasticsearch.search.rescore.RescorerBuilder; 27 | import redis.clients.jedis.Jedis; 28 | 29 | import java.io.IOException; 30 | import java.security.AccessController; 31 | import java.security.PrivilegedAction; 32 | import java.util.Arrays; 33 | import java.util.Iterator; 34 | import java.util.Objects; 35 | 36 | import static java.util.Collections.singletonList; 37 | import static org.elasticsearch.common.xcontent.ConstructingObjectParser.constructorArg; 38 | import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg; 39 | 40 | public class RedisRescoreBuilder extends RescorerBuilder { 41 | public static final String NAME = "redis"; 42 | 43 | protected static final Logger log = LogManager.getLogger(RedisRescoreBuilder.class); 44 | 45 | private final String keyField; 46 | private final String keyPrefix; 47 | 48 | private static Jedis jedis; 49 | public static void setJedis(Jedis j) { 50 | jedis = j; 51 | } 52 | 53 | public RedisRescoreBuilder(final String keyField, @Nullable String keyPrefix) { 54 | this.keyField = keyField; 55 | this.keyPrefix = keyPrefix; 56 | } 57 | 58 | public RedisRescoreBuilder(StreamInput in) throws IOException { 59 | super(in); 60 | keyField = in.readString(); 61 | keyPrefix = in.readOptionalString(); 62 | } 63 | 64 | @Override 65 | protected void doWriteTo(StreamOutput out) throws IOException { 66 | out.writeString(keyField); 67 | out.writeOptionalString(keyPrefix); 68 | } 69 | 70 | @Override 71 | public String getWriteableName() { 72 | return NAME; 73 | } 74 | 75 | @Override 76 | public RescorerBuilder rewrite(QueryRewriteContext ctx) throws IOException { 77 | return this; 78 | } 79 | 80 | private static final ParseField KEY_FIELD = new ParseField("key_field"); 81 | private static final ParseField KEY_PREFIX = new ParseField("key_prefix"); 82 | @Override 83 | protected void doXContent(XContentBuilder builder, Params params) throws IOException { 84 | builder.field(KEY_FIELD.getPreferredName(), keyField); 85 | if (keyPrefix != null) { 86 | builder.field(KEY_PREFIX.getPreferredName(), keyPrefix); 87 | } 88 | } 89 | 90 | private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>(NAME, 91 | args -> new RedisRescoreBuilder((String) args[0], (String) args[1])); 92 | static { 93 | PARSER.declareString(constructorArg(), KEY_FIELD); 94 | PARSER.declareString(optionalConstructorArg(), KEY_PREFIX); 95 | } 96 | public static RedisRescoreBuilder fromXContent(XContentParser parser) { 97 | return PARSER.apply(parser, null); 98 | } 99 | 100 | @Override 101 | public RescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException { 102 | IndexFieldData keyField = 103 | this.keyField == null ? null : context.getForField(context.fieldMapper(this.keyField)); 104 | return new RedisRescoreContext(windowSize, keyPrefix, keyField); 105 | } 106 | 107 | @Override 108 | public boolean equals(Object obj) { 109 | if (!super.equals(obj)) { 110 | return false; 111 | } 112 | RedisRescoreBuilder other = (RedisRescoreBuilder) obj; 113 | return keyField.equals(other.keyField) 114 | && Objects.equals(keyPrefix, other.keyPrefix); 115 | } 116 | 117 | @Override 118 | public int hashCode() { 119 | return Objects.hash(super.hashCode(), keyField, keyPrefix); 120 | } 121 | 122 | String keyField() { 123 | return keyField; 124 | } 125 | 126 | @Nullable 127 | String keyPrefix() { 128 | return keyPrefix; 129 | } 130 | 131 | private static class RedisRescoreContext extends RescoreContext { 132 | private final String keyPrefix; 133 | @Nullable 134 | private final IndexFieldData keyField; 135 | 136 | RedisRescoreContext(int windowSize, String keyPrefix, @Nullable IndexFieldData keyField) { 137 | super(windowSize, RedisRescorer.INSTANCE); 138 | this.keyPrefix = keyPrefix; 139 | this.keyField = keyField; 140 | } 141 | } 142 | 143 | private static class RedisRescorer implements Rescorer { 144 | 145 | private static final RedisRescorer INSTANCE = new RedisRescorer(); 146 | 147 | private static String getTermFromFieldData(int topLevelDocId, AtomicFieldData fd, 148 | LeafReaderContext leaf, String fieldName) throws IOException { 149 | String term = null; 150 | if (fd instanceof SortedSetDVBytesAtomicFieldData) { 151 | final SortedSetDocValues data = ((SortedSetDVBytesAtomicFieldData) fd).getOrdinalsValues(); 152 | if (data != null) { 153 | if (data.advanceExact(topLevelDocId - leaf.docBase)) { 154 | // document does have data for the field 155 | term = data.lookupOrd(data.nextOrd()).utf8ToString(); 156 | } 157 | } 158 | } else if (fd instanceof AtomicNumericFieldData) { 159 | final SortedNumericDocValues data = ((AtomicNumericFieldData) fd).getLongValues(); 160 | if (data != null) { 161 | if (!data.advanceExact(topLevelDocId - leaf.docBase)) { 162 | throw new IllegalArgumentException("document [" + topLevelDocId 163 | + "] does not have the field [" + fieldName + "]"); 164 | } 165 | if (data.docValueCount() > 1) { 166 | throw new IllegalArgumentException("document [" + topLevelDocId 167 | + "] has more than one value for [" + fieldName + "]"); 168 | } 169 | term = String.valueOf(data.nextValue()); 170 | } 171 | } 172 | return term; 173 | } 174 | 175 | @Override 176 | public TopDocs rescore(TopDocs topDocs, IndexSearcher searcher, RescoreContext rescoreContext) throws IOException { 177 | assert rescoreContext != null; 178 | if (topDocs == null || topDocs.scoreDocs.length == 0) { 179 | return topDocs; 180 | } 181 | 182 | final RedisRescoreContext context = (RedisRescoreContext) rescoreContext; 183 | 184 | if (context.keyField != null) { 185 | /* 186 | * Since this example looks up a single field value it should 187 | * access them in docId order because that is the order in 188 | * which they are stored on disk and we want reads to be 189 | * forwards and close together if possible. 190 | * 191 | * If accessing multiple fields we'd be better off accessing 192 | * them in (reader, field, docId) order because that is the 193 | * order they are on disk. 194 | */ 195 | 196 | final Iterator leaves = searcher.getIndexReader().leaves().iterator(); 197 | LeafReaderContext leaf = null; 198 | 199 | final int end = Math.min(topDocs.scoreDocs.length, rescoreContext.getWindowSize()); 200 | SortedSetDocValues docValues = null; 201 | SortedNumericDocValues numericDocValues = null; 202 | int endDoc = 0; 203 | for (int i = 0; i < end; i++) { 204 | if (topDocs.scoreDocs[i].doc >= endDoc) { 205 | do { 206 | leaf = leaves.next(); 207 | endDoc = leaf.docBase + leaf.reader().maxDoc(); 208 | } while (topDocs.scoreDocs[i].doc >= endDoc); 209 | 210 | final AtomicFieldData fd = context.keyField.load(leaf); 211 | if (fd instanceof SortedSetDVBytesAtomicFieldData) { 212 | docValues = ((SortedSetDVBytesAtomicFieldData) fd).getOrdinalsValues(); 213 | } else if (fd instanceof AtomicNumericFieldData) { 214 | numericDocValues = ((AtomicNumericFieldData) fd).getLongValues(); 215 | } 216 | } 217 | if (docValues != null) { 218 | if (docValues.advanceExact(topDocs.scoreDocs[i].doc - leaf.docBase)) { 219 | // document does have data for the field 220 | final String term = docValues.lookupOrd(docValues.nextOrd()).utf8ToString(); 221 | topDocs.scoreDocs[i].score *= getScoreFactor(term, context.keyPrefix); 222 | } 223 | } else if (numericDocValues != null) { 224 | if (!numericDocValues.advanceExact(topDocs.scoreDocs[i].doc - leaf.docBase)) { 225 | throw new IllegalArgumentException("document [" + topDocs.scoreDocs[i].doc 226 | + "] does not have the field [" + context.keyField.getFieldName() + "]"); 227 | } 228 | if (numericDocValues.docValueCount() > 1) { 229 | throw new IllegalArgumentException("document [" + topDocs.scoreDocs[i].doc 230 | + "] has more than one value for [" + context.keyField.getFieldName() + "]"); 231 | } 232 | 233 | topDocs.scoreDocs[i].score *= getScoreFactor(String.valueOf(numericDocValues.nextValue()), 234 | context.keyPrefix); 235 | 236 | } 237 | } 238 | } 239 | 240 | // Sort by score descending, then docID ascending, just like lucene's QueryRescorer 241 | Arrays.sort(topDocs.scoreDocs, (a, b) -> { 242 | if (a.score > b.score) { 243 | return -1; 244 | } 245 | if (a.score < b.score) { 246 | return 1; 247 | } 248 | // Safe because doc ids >= 0 249 | return a.doc - b.doc; 250 | }); 251 | return topDocs; 252 | } 253 | 254 | private static float getScoreFactor(final String key, @Nullable final String keyPrefix) { 255 | assert key != null; 256 | 257 | return AccessController.doPrivileged((PrivilegedAction) () -> { 258 | final String fullKey = fullKey(key, keyPrefix); 259 | final String factor = jedis.get(fullKey); 260 | if (factor == null) { 261 | log.debug("Redis rescore factor null for key " + keyPrefix + key); 262 | return 1.0f; 263 | } 264 | 265 | try { 266 | return Float.parseFloat(factor); 267 | } catch (NumberFormatException ignored_e) { 268 | log.warn("Redis rescore factor NumberFormatException for key " + fullKey); 269 | return 1.0f; 270 | } 271 | }); 272 | } 273 | 274 | private static String fullKey(final String key, @Nullable final String keyPrefix) { 275 | if (keyPrefix == null) { 276 | return key; 277 | } else { 278 | return keyPrefix + key; 279 | } 280 | } 281 | 282 | @Override 283 | public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreContext rescoreContext, 284 | Explanation sourceExplanation) throws IOException { 285 | final RedisRescoreContext context = (RedisRescoreContext) rescoreContext; 286 | final Iterator leaves = searcher.getIndexReader().leaves().iterator(); 287 | LeafReaderContext leaf = null; 288 | int endDoc = 0; 289 | do { 290 | leaf = leaves.next(); 291 | endDoc = leaf.docBase + leaf.reader().maxDoc(); 292 | } while (topLevelDocId >= endDoc); 293 | 294 | AtomicFieldData fd = context.keyField.load(leaf); 295 | String fieldName = context.keyField.getFieldName(); 296 | String term = getTermFromFieldData(topLevelDocId, fd, leaf, fieldName); 297 | if (term != null) { 298 | float score = getScoreFactor(term, context.keyPrefix); 299 | return Explanation.match(score, fieldName, singletonList(sourceExplanation)); 300 | } else { 301 | return Explanation.noMatch(fieldName, sourceExplanation); 302 | } 303 | } 304 | } 305 | } 306 | -------------------------------------------------------------------------------- /src/main/java/com/bigdataboutique/elasticsearch/plugin/RedisRescorePlugin.java: -------------------------------------------------------------------------------- 1 | package com.bigdataboutique.elasticsearch.plugin; 2 | 3 | import org.elasticsearch.common.settings.Setting; 4 | import org.elasticsearch.common.settings.Settings; 5 | import org.elasticsearch.plugins.Plugin; 6 | import org.elasticsearch.plugins.SearchPlugin; 7 | import redis.clients.jedis.Jedis; 8 | 9 | import java.util.Arrays; 10 | import java.util.List; 11 | import static java.util.Collections.singletonList; 12 | import java.nio.file.Path; 13 | 14 | public class RedisRescorePlugin extends Plugin implements SearchPlugin { 15 | private final Config config; 16 | 17 | public RedisRescorePlugin(final Settings settings, final Path configPath) { 18 | this.config = new Config(settings); 19 | RedisRescoreBuilder.setJedis(new Jedis(config.getRedisUrl())); 20 | } 21 | 22 | /** 23 | * @return the plugin's custom settings 24 | */ 25 | @Override 26 | public List> getSettings() { 27 | return Arrays.asList(Config.REDIS_URL); 28 | } 29 | 30 | @Override 31 | public Settings additionalSettings() { 32 | final Settings.Builder builder = Settings.builder(); 33 | 34 | // Exposes REDIS_URL as a node setting 35 | builder.put(Config.REDIS_URL.getKey(), config.getRedisUrl()); 36 | 37 | return builder.build(); 38 | } 39 | 40 | @Override 41 | public List> getRescorers() { 42 | return singletonList( 43 | new RescorerSpec<>(RedisRescoreBuilder.NAME, RedisRescoreBuilder::new, RedisRescoreBuilder::fromXContent)); 44 | } 45 | } -------------------------------------------------------------------------------- /src/main/plugin-metadata/plugin-security.policy: -------------------------------------------------------------------------------- 1 | grant { 2 | permission java.net.SocketPermission "127.0.0.1:6379", "connect,resolve"; 3 | }; 4 | -------------------------------------------------------------------------------- /src/test/java/com/bigdataboutique/elasticsearch/plugin/ConfigTests.java: -------------------------------------------------------------------------------- 1 | package com.bigdataboutique.elasticsearch.plugin; 2 | 3 | import org.elasticsearch.common.settings.Settings; 4 | import org.elasticsearch.test.ESTestCase; 5 | 6 | import static com.bigdataboutique.elasticsearch.plugin.Config.REDIS_URL; 7 | 8 | /** 9 | * {@link ConfigTests} is a unit test class for {@link Config}. 10 | *

11 | * It's a JUnit test class that extends {@link ESTestCase} which provides useful methods for testing. 12 | *

13 | * The tests can be executed in the IDE or using the command: ./gradlew test 14 | */ 15 | public class ConfigTests extends ESTestCase { 16 | 17 | public void testValidatedSetting() { 18 | final String expected = "http://localhost:6379"; 19 | final String actual = REDIS_URL.get(Settings.builder().put(REDIS_URL.getKey(), expected).build()); 20 | assertEquals(expected, actual); 21 | 22 | final IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> 23 | REDIS_URL.get(Settings.builder().put(REDIS_URL.getKey(), "not an URI").build())); 24 | assertEquals("Setting is not a valid URI", exception.getMessage()); 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /src/test/java/com/bigdataboutique/elasticsearch/plugin/RedisRescoreBuilderTests.java: -------------------------------------------------------------------------------- 1 | package com.bigdataboutique.elasticsearch.plugin; 2 | 3 | import org.elasticsearch.common.io.stream.Writeable; 4 | import org.elasticsearch.test.AbstractWireSerializingTestCase; 5 | 6 | import java.io.IOException; 7 | 8 | public class RedisRescoreBuilderTests extends AbstractWireSerializingTestCase { 9 | @Override 10 | protected RedisRescoreBuilder createTestInstance() { 11 | String factorField = randomBoolean() ? null : randomAlphaOfLength(5); 12 | return new RedisRescoreBuilder("prefix-", factorField).windowSize(between(0, Integer.MAX_VALUE)); 13 | } 14 | 15 | @Override 16 | protected Writeable.Reader instanceReader() { 17 | return RedisRescoreBuilder::new; 18 | } 19 | 20 | @Override 21 | protected RedisRescoreBuilder mutateInstance(RedisRescoreBuilder instance) throws IOException { 22 | return new RedisRescoreBuilder(instance.keyField(), instance.keyPrefix()); 23 | } 24 | 25 | public void testRewrite() throws IOException { 26 | RedisRescoreBuilder builder = createTestInstance(); 27 | assertSame(builder, builder.rewrite(null)); 28 | } 29 | 30 | // public void testRescore() throws IOException { 31 | // // Always use a factor > 1 so rescored fields are sorted in front of the unrescored fields. 32 | // float factor = (float) randomDoubleBetween(1.0d, Float.MAX_VALUE, false); 33 | // RedisRescoreBuilder builder = new RedisRescoreBuilder("productId", "mystore-").windowSize(2); 34 | // RescoreContext context = builder.buildContext(null); 35 | // TopDocs docs = new TopDocs(new TotalHits(10, TotalHits.Relation.EQUAL_TO), new ScoreDoc[3]); 36 | // docs.scoreDocs[0] = new ScoreDoc(0, 1.0f); 37 | // docs.scoreDocs[1] = new ScoreDoc(1, 1.0f); 38 | // docs.scoreDocs[2] = new ScoreDoc(2, 1.0f); 39 | // context.rescorer().rescore(docs, null, context); 40 | // assertEquals(factor, docs.scoreDocs[0].score, 0.0f); 41 | // assertEquals(factor, docs.scoreDocs[1].score, 0.0f); 42 | // assertEquals(1.0f, docs.scoreDocs[2].score, 0.0f); 43 | // } 44 | } 45 | -------------------------------------------------------------------------------- /src/test/java/com/bigdataboutique/elasticsearch/plugin/RedisRescoreClientYamlTestSuiteIT.java: -------------------------------------------------------------------------------- 1 | package com.bigdataboutique.elasticsearch.plugin; 2 | 3 | import java.security.AccessController; 4 | import java.security.PrivilegedAction; 5 | 6 | import com.carrotsearch.randomizedtesting.annotations.Name; 7 | import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; 8 | import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate; 9 | import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase; 10 | import org.junit.After; 11 | import org.junit.Before; 12 | import redis.clients.jedis.Jedis; 13 | 14 | public class RedisRescoreClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase { 15 | final Jedis jedis = new Jedis("localhost"); 16 | 17 | @Before 18 | public void startRedis() { 19 | AccessController.doPrivileged((PrivilegedAction) () -> { 20 | jedis.set("foo", "5.0"); 21 | jedis.set("bar", "3.0"); 22 | 23 | jedis.set("mystore-foo", "4.0"); 24 | jedis.set("mystore-bar", "6.0"); 25 | 26 | return true; 27 | }); 28 | } 29 | 30 | @After 31 | public void closeRedisConnection() { 32 | jedis.close(); 33 | } 34 | 35 | public RedisRescoreClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) { 36 | super(testCandidate); 37 | } 38 | 39 | @ParametersFactory 40 | public static Iterable parameters() throws Exception { 41 | return ESClientYamlSuiteTestCase.createParameters(); 42 | } 43 | } -------------------------------------------------------------------------------- /src/test/resources/rest-api-spec/test/redis-rescore/10_basic.yml: -------------------------------------------------------------------------------- 1 | # Integration tests for the redis rescoring plugin 2 | # 3 | "Plugin loaded": 4 | - skip: 5 | reason: "contains is a newly added assertion" 6 | features: contains 7 | - do: 8 | cluster.state: {} 9 | 10 | # Get master node id 11 | - set: { master_node: master } 12 | 13 | - do: 14 | nodes.info: {} 15 | 16 | - contains: { nodes.$master.plugins: { name: elasticsearch-rescore-redis } } 17 | -------------------------------------------------------------------------------- /src/test/resources/rest-api-spec/test/redis-rescore/10_config.yml: -------------------------------------------------------------------------------- 1 | "Test plugin settings": 2 | 3 | # Use the Get Cluster Settings API to list the settings including the default ones 4 | - do: 5 | cluster.get_settings: 6 | include_defaults: true 7 | 8 | - match: { defaults.redisRescore.redisUrl: "localhost" } 9 | 10 | # Use the Cluster Update Settings API to update some custom settings 11 | - do: 12 | cluster.put_settings: 13 | body: 14 | transient: 15 | redisRescore: 16 | redisUrl: "http://example.com" 17 | 18 | # Use the Get Cluster Settings API to list the settings again 19 | - do: 20 | cluster.get_settings: {} 21 | 22 | - match: { transient.redisRescore.redisUrl: "http://example.com" } 23 | 24 | # Try to update the "redisUrl" setting with an invalid value 25 | - do: 26 | catch: bad_request 27 | cluster.put_settings: 28 | body: 29 | transient: 30 | redisRescore: 31 | redisUrl: "this is not a URI!" 32 | 33 | # Reset the settings to their default values 34 | - do: 35 | cluster.put_settings: 36 | body: 37 | transient: 38 | redisRescore: 39 | redisUrl: "localhost" 40 | -------------------------------------------------------------------------------- /src/test/resources/rest-api-spec/test/redis-rescore/20_score.yml: -------------------------------------------------------------------------------- 1 | --- 2 | setup: 3 | - do: 4 | indices.create: 5 | index: test 6 | body: 7 | settings: 8 | number_of_shards: 1 9 | number_of_replicas: 1 10 | 11 | - do: 12 | index: 13 | index: test 14 | id: 1 15 | body: { "test": 1, "productId": "foo" } 16 | - do: 17 | index: 18 | index: test 19 | id: 2 20 | body: { "test": 2, "productId": "bar" } 21 | - do: 22 | index: 23 | index: test 24 | id: 3 25 | body: { "test": 3 } 26 | - do: 27 | indices.refresh: {} 28 | 29 | --- 30 | "no prefix": 31 | - do: 32 | search: 33 | rest_total_hits_as_int: true 34 | index: test 35 | body: 36 | rescore: 37 | redis: 38 | key_field: "productId.keyword" 39 | - length: { hits.hits: 3 } 40 | - match: { hits.hits.0._score: 5 } 41 | - match: { hits.hits.1._score: 3 } 42 | - match: { hits.hits.2._score: 1 } 43 | --- 44 | "with key prefix": 45 | - do: 46 | search: 47 | rest_total_hits_as_int: true 48 | index: test 49 | body: 50 | rescore: 51 | redis: 52 | key_field: "productId.keyword" 53 | key_prefix: "mystore-" 54 | - length: { hits.hits: 3 } 55 | - match: { hits.hits.0._score: 6 } 56 | - match: { hits.hits.1._score: 4 } 57 | - match: { hits.hits.2._score: 1 } 58 | --- 59 | "key is numeric field": 60 | - do: 61 | search: 62 | rest_total_hits_as_int: true 63 | index: test 64 | body: 65 | rescore: 66 | redis: 67 | key_field: "test" 68 | - length: { hits.hits: 3 } 69 | - match: { hits.hits.0._score: 1 } 70 | - match: { hits.hits.1._score: 1 } 71 | - match: { hits.hits.2._score: 1 } 72 | --------------------------------------------------------------------------------