Refactor parsing of KNNQueryBuilder (opensearch-project#1824)

Refactors parsing of KNNQueryBuilder. First, it moves parsing logic to a separate class. Next, it uses ObjectParser instead of parsing manually by hand. Lastly, it also moves out the streaming. To test, it contains a simple jmh benchmark for testing as well as new unit tests. Signed-off-by: John Mazanec <[email protected]> (cherry picked from commit 52636c4)
jmazanec15 · Jul 29, 2024 · dd0d142 · dd0d142
1 parent c1450cd
commit dd0d142
Show file tree

Hide file tree

Showing 9 changed files with 878 additions and 538 deletions.
diff --git a/.idea/copyright/SPDX_ALv2.xml b/.idea/copyright/SPDX_ALv2.xml
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,4 +20,5 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Infrastructure
 ### Documentation
 ### Maintenance
-### Refactoring
+### Refactoring
+* Clean up parsing for query [#1824](https://github.com/opensearch-project/k-NN/pull/1824)
diff --git a/micro-benchmarks/src/main/java/org/opensearch/knn/QueryParsingBenchmarks.java b/micro-benchmarks/src/main/java/org/opensearch/knn/QueryParsingBenchmarks.java
@@ -0,0 +1,109 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+import org.opensearch.cluster.ClusterModule;
+import org.opensearch.common.xcontent.LoggingDeprecationHandler;
+import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.common.xcontent.json.JsonXContent;
+import org.opensearch.core.common.bytes.BytesArray;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.index.query.TermQueryBuilder;
+import org.opensearch.knn.index.query.KNNQueryBuilder;
+import org.opensearch.knn.index.query.parser.KNNQueryBuilderParser;
+import org.opensearch.plugins.SearchPlugin;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * Benchmarks for impact of changes around query parsing
+ */
+@Warmup(iterations = 5, time = 10)
+@Measurement(iterations = 3, time = 10)
+@Fork(3)
+@State(Scope.Benchmark)
+public class QueryParsingBenchmarks {
+    private static final TermQueryBuilder TERM_QUERY = QueryBuilders.termQuery("field", "value");
+    private static final NamedXContentRegistry NAMED_X_CONTENT_REGISTRY = xContentRegistry();
+
+    @Param({ "128", "1024" })
+    private int dimension;
+    @Param({ "basic", "filter" })
+    private String type;
+
+    private BytesReference bytesReference;
+
+    @Setup
+    public void setup() throws IOException {
+        XContentBuilder builder = XContentFactory.jsonBuilder();
+        builder.startObject();
+        builder.startObject("test");
+        builder.field(KNNQueryBuilder.VECTOR_FIELD.getPreferredName(), generateVectorWithOnes(dimension));
+        builder.field(KNNQueryBuilder.K_FIELD.getPreferredName(), 1);
+        if (type.equals("filter")) {
+            builder.field(KNNQueryBuilder.FILTER_FIELD.getPreferredName(), TERM_QUERY);
+        }
+        builder.endObject();
+        builder.endObject();
+        bytesReference = BytesReference.bytes(builder);
+    }
+
+    @Benchmark
+    public void fromXContent(final Blackhole bh) throws IOException {
+        XContentParser xContentParser = createParser();
+        bh.consume(KNNQueryBuilderParser.fromXContent(xContentParser));
+    }
+
+    private XContentParser createParser() throws IOException {
+        XContentParser contentParser = createParser(bytesReference);
+        contentParser.nextToken();
+        return contentParser;
+    }
+
+    private float[] generateVectorWithOnes(final int dimensions) {
+        float[] vector = new float[dimensions];
+        Arrays.fill(vector, (float) 1);
+        return vector;
+    }
+
+    private XContentParser createParser(final BytesReference data) throws IOException {
+        BytesArray array = (BytesArray) data;
+        return JsonXContent.jsonXContent.createParser(
+            NAMED_X_CONTENT_REGISTRY,
+            LoggingDeprecationHandler.INSTANCE,
+            array.array(),
+            array.offset(),
+            array.length()
+        );
+    }
+
+    private static NamedXContentRegistry xContentRegistry() {
+        List<NamedXContentRegistry.Entry> list = ClusterModule.getNamedXWriteables();
+        SearchPlugin.QuerySpec<?> spec = new SearchPlugin.QuerySpec<>(
+            TermQueryBuilder.NAME,
+            TermQueryBuilder::new,
+            TermQueryBuilder::fromXContent
+        );
+        list.add(new NamedXContentRegistry.Entry(QueryBuilder.class, spec.getName(), (p, c) -> spec.getParser().fromXContent(p)));
+        return new NamedXContentRegistry(list);
+    }
+}