Issue #459: Add Mustache Scripting Support for knn_vector Fields

Signed-off-by: Abhinav Tripathi <[email protected]>

Author: night-owl-1709

CHANGELOG.md

@@ -7,3 +7,5 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 3.3](https://github.com/opensearch-project/k-NN/compare/main...HEAD)
 ### Refactoring
 * Refactored the KNN Stat files for better readability.
+### Enhancements
+* Added mustache template support for knn [#2807](https://github.com/opensearch-project/k-NN/pull/2807)

build.gradle

@@ -349,6 +349,8 @@ dependencies {
     compileOnly "org.opensearch:protobufs:0.6.0"
     compileOnly group: 'com.google.guava', name: 'failureaccess', version:'1.0.2'
     compileOnly group: 'com.google.guava', name: 'guava', version:'33.2.1-jre'
+    // Mustache support - only core mustache library, OpenSearch classes available at runtime
+    compileOnly "com.github.spullara.mustache.java:compiler:0.9.14"
     api group: 'commons-lang', name: 'commons-lang', version: '2.6'
     testFixturesImplementation "org.opensearch.test:framework:${opensearch_version}"
     // Add Guava for test configurations since it's needed for testing but excluded from runtime

src/main/java/org/opensearch/knn/plugin/KNNPlugin.java

@@ -59,6 +59,7 @@
 import org.opensearch.knn.plugin.rest.RestSearchModelHandler;
 import org.opensearch.knn.plugin.rest.RestTrainModelHandler;
 import org.opensearch.knn.plugin.script.KNNScoringScriptEngine;
+import org.opensearch.knn.plugin.script.KNNMustacheEngine;
 import org.opensearch.knn.plugin.search.KNNConcurrentSearchRequestDecider;
 import org.opensearch.knn.plugin.stats.KNNStats;
 import org.opensearch.knn.plugin.transport.ClearCacheAction;
@@ -107,6 +108,7 @@
 import org.opensearch.rest.RestHandler;
 import org.opensearch.script.ScriptContext;
 import org.opensearch.script.ScriptEngine;
+import org.opensearch.script.TemplateScript;
 import org.opensearch.script.ScriptService;
 import org.opensearch.search.deciders.ConcurrentSearchRequestDecider;
 import org.opensearch.threadpool.ExecutorBuilder;
@@ -380,6 +382,11 @@ public void onIndexModule(IndexModule indexModule) {
      */
     @Override
     public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
+        // Check if this is specifically a template script request
+        if (contexts.size() == 1 && contexts.contains(TemplateScript.CONTEXT)) {
+            return new KNNMustacheEngine();
+        }
+        // Default to scoring engine for other contexts
         return new KNNScoringScriptEngine();
     }
 

src/main/java/org/opensearch/knn/plugin/script/KNNMustacheEngine.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.plugin.script;
+
+import com.github.mustachejava.DefaultMustacheFactory;
+import com.github.mustachejava.Mustache;
+import com.github.mustachejava.MustacheException;
+import org.opensearch.script.GeneralScriptException;
+import org.opensearch.script.ScriptContext;
+import org.opensearch.script.ScriptEngine;
+import org.opensearch.script.ScriptException;
+import org.opensearch.script.TemplateScript;
+import org.opensearch.knn.plugin.stats.KNNCounter;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.io.StringReader;
+import java.io.StringWriter;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Mustache script engine for k-NN query templating.
+ */
+public class KNNMustacheEngine implements ScriptEngine {
+
+    public static final String NAME = "knn-mustache";
+    private static final Logger logger = LogManager.getLogger(KNNMustacheEngine.class);
+
+    @Override
+    public String getType() {
+        return NAME;
+    }
+
+    @Override
+    public <T> T compile(String templateName, String templateSource, ScriptContext<T> context, Map<String, String> options) {
+        KNNCounter.SCRIPT_COMPILATIONS.increment();
+
+        if (!context.instanceClazz.equals(TemplateScript.class)) {
+            KNNCounter.SCRIPT_COMPILATION_ERRORS.increment();
+            throw new IllegalArgumentException("knn-mustache engine only supports TemplateScript context, got [" + context.name + "]");
+        }
+
+        try (StringReader reader = new StringReader(templateSource)) {
+            DefaultMustacheFactory factory = new DefaultMustacheFactory();
+            Mustache template = factory.compile(reader, templateName != null ? templateName : "knn-template");
+
+            TemplateScript.Factory compiled = params -> new KNNMustacheTemplateScript(template, params);
+            return context.factoryClazz.cast(compiled);
+        } catch (MustacheException ex) {
+            KNNCounter.SCRIPT_COMPILATION_ERRORS.increment();
+            logger.error("Failed to compile k-NN Mustache template [{}]: {}", templateName, ex.getMessage());
+            throw new ScriptException(ex.getMessage(), ex, Collections.emptyList(), templateSource, NAME);
+        }
+    }
+
+    @Override
+    public Set<ScriptContext<?>> getSupportedContexts() {
+        return Collections.singleton(TemplateScript.CONTEXT);
+    }
+
+    /**
+     * Template script implementation for k-NN query generation.
+     */
+    private static class KNNMustacheTemplateScript extends TemplateScript {
+        private final Mustache template;
+        private final Map<String, Object> params;
+
+        KNNMustacheTemplateScript(Mustache template, Map<String, Object> params) {
+            super(params);
+            this.template = template;
+            this.params = params;
+        }
+
+        @Override
+        public String execute() {
+            StringWriter writer = new StringWriter();
+            try {
+                template.execute(writer, params);
+                return writer.toString();
+            } catch (Exception e) {
+                logger.error("Error executing k-NN mustache template: {}", e.getMessage());
+                throw new GeneralScriptException("Error executing k-NN mustache template: " + e.getMessage(), e);
+            }
+        }
+    }
+}