Issue #459: Add Mustache Scripting Support for knn_vector Fields

[night-owl-1709]

Summary

This PR introduces Mustache templating support for k-NN queries by implementing a dedicated KNNMustacheEngine that enables dynamic query generation through the Search Template API. This enhancement allows users to create reusable k-NN query templates with parameterized field names, vectors, and search parameters.

Issues

Implements #459

Changes Made

1. New Script Engine Implementation

• File: src/main/java/org/opensearch/knn/plugin/script/KNNMustacheEngine.java
• Created a new script engine that implements ScriptEngine interface
• Supports only TemplateScript.CONTEXT for query templating
• Integrates with existing k-NN statistics system through KNNCounter
• Uses standard Mustache Java library for template compilation and execution

2. Plugin Integration

• Modified: Integration in KNNPlugin.getScriptEngine() method
• Added logic to return KNNMustacheEngine when TemplateScript.CONTEXT is requested
• Maintains backward compatibility with existing KNNScoringScriptEngine
• Added import for the new engine class

3. Build Configuration

• Modified: build.gradle
• Added Mustache Java library dependency: com.github.spullara.mustache.java:compiler:0.9.10

How Mustache Integration Works

The integration follows OpenSearch's plugin architecture patterns and works through these components:

1. Script Engine Registration: The engine is registered through KNNPlugin.getScriptEngine() and responds to TemplateScript.CONTEXT requests
2. Template Compilation: Uses DefaultMustacheFactory to compile template strings into executable Mustache objects
3. Parameter Substitution: Replaces template variables (e.g., {{field}}, {{vector}}, {{k}}) with actual values from the params map
4. Query Generation: Produces complete k-NN queries that are processed by the standard k-NN query pipeline

Usage Examples

Basic k-NN Query Template

POST /products/_search/template
{
  "source": {
    "knn": {
      "{{field_name}}": {
        "vector": {{query_vector}},
        "k": {{k_value}}
      }
    }
  },
  "params": {
    "field_name": "product_embedding",
    "query_vector": [0.2, 0.8, 0.1, 0.9, 0.3],
    "k_value": 10
  }
}

k-NN Query with Conditional Parameters

POST /products/_search/template
{
  "source": {
    "knn": {
      "{{field_name}}": {
        "vector": {{query_vector}},
        "k": {{k_value}}{{#max_distance}},
        "max_distance": {{max_distance}}{{/max_distance}}
      }
    }
  },
  "params": {
    "field_name": "product_embedding",
    "query_vector": [0.2, 0.8, 0.1, 0.9, 0.3],
    "k_value": 10,
    "max_distance": 0.8
  }
}

k-NN Query with Filter Template

POST /products/_search/template
{
  "source": {
    "query": {
      "knn": {
        "{{vector_field}}": {
          "vector": {{query_vector}},
          "k": {{k_value}},
          "filter": {
            "term": {
              "{{filter_field}}": "{{filter_value}}"
            }
          }
        }
      }
    }
  },
  "params": {
    "vector_field": "product_embedding",
    "query_vector": [0.2, 0.8, 0.1, 0.9, 0.3],
    "k_value": 10,
    "filter_field": "category",
    "filter_value": "electronics"
  }
}

Benefits

1. Dynamic Query Generation: Enables parameterized k-NN queries for different use cases
2. Template Reusability: Single templates can be used with different parameters
3. Integration with Existing Architecture: Leverages OpenSearch's Search Template API
4. Statistics Tracking: Integrates with k-NN plugin's existing statistics system
5. Error Handling: Provides comprehensive error handling for template compilation and execution

Testing

• Unit tests for KNNMustacheEngine compilation and execution
• Integration tests with actual k-NN queries using templates
• Error handling tests for malformed templates and invalid parameters

Backward Compatibility

This change is fully backward compatible:

• Existing k-NN queries continue to work unchanged
• The KNNScoringScriptEngine remains the default for script scoring contexts
• No changes to existing k-NN query processing pipeline

[jmazanec15]

Thanks @night-owl-1709 - will take a look!

[jmazanec15]

@night-owl-1709 One question: for painless, we extend the lang-painless plugin and use SPI to extend it. Is that not possible with expression and mustache?

[jmazanec15]

Instead of treating knn_l2_distance like a value (and the other distance methods as values), is it possible to treat them like functions in painless? For example:

GET my-knn-index-2/_search
{
  "size": 2,
  "query": {
    "script_score": {
      "query": {
        "bool": {
          "filter": {
            "term": {
              "color": "BLUE"
            }
          }
        }
      },
      "script": {
        "source": "1.0 + cosineSimilarity(params.query_value, doc[params.field])",
        "params": {
          "field": "my_vector",
          "query_value": [9.9, 9.9]
        }
      }
    }
  }
}

Im not sure if this is possible, but I feel the distance functions should be somewhat standardized, similar to operations like "+".

[jmazanec15]

Just want to

[jmazanec15]

@Vikasht34, @shatejas to review as well

[shatejas]

I was wondering if plugin should reuse the existing implementation of mustache engine. The plugin will probably need a dependency on lang-mustache module. @night-owl-1709 have you explored the approach? were there any problems with it?

[night-owl-1709]

I checked but we do not have lang-mustache support yet like we have for painless in opensearch.

[shatejas]

Not sure if I understand this but did you try adding lang-mustache here https://github.com/opensearch-project/k-NN/blob/main/build.gradle#L318

[shatejas]

So the implementation here is similar to https://github.com/opensearch-project/OpenSearch/blob/main/modules/lang-mustache/src/main/java/org/opensearch/script/mustache/MustacheScriptEngine.java#L67 with minor changes.

• This uses default factory compared to custom factory. The compile implementation of both is the same as custom factory inherits from default factory.
• There are some node metrics that are being added, if core implementation is used, I don't think its needed.

so instead of returning knn mustache engine we can return core implementation of it.

I would prefer using core implementation to be consistent and less maintenance on the plugin. Let me know if I am missing something here

[shatejas]

So the implementation here is similar to https://github.com/opensearch-project/OpenSearch/blob/main/modules/lang-mustache/src/main/java/org/opensearch/script/mustache/MustacheScriptEngine.java#L67 with minor changes.

• This uses default factory compared to custom factory. The compile implementation of both is the same as custom factory inherits from default factory.
• There are some node metrics that are being added, if core implementation is used, I don't think its needed.

so instead of returning knn mustache engine we can return core implementation of it.

I would prefer using core implementation to be consistent and less maintenance on the plugin. Let me know if I am missing something here

[shatejas]

Does this need privileged access? https://github.com/opensearch-project/OpenSearch/blob/main/modules/lang-mustache/src/main/java/org/opensearch/script/mustache/MustacheScriptEngine.java#L67