Refactor parsing logic of the Query Builder (second try) #1824

jmazanec15 · 2024-07-13T00:19:34Z

Description

A continuation of #1793, but decided to just focus on parsing. KNNQueryBuilder parsing has gotten pretty complex and this changes it to use ObjectParser in order to simplify it.

I moved XContent related parsing and streaming into its own class (KNNQueryParser).

In addition, I added a JMH test to test the parsing. I ran locally on my mac, and got the following results. From the results it seems there is no statistically significant difference:

## Before
Benchmark                            (dimension)  (type)   Mode  Cnt       Score       Error  Units
QueryParsingBenchmarks.fromXContent          128   basic  thrpt    9  112177.006 ± 11012.866  ops/s
QueryParsingBenchmarks.fromXContent          128  filter  thrpt    9  100488.178 ± 24078.112  ops/s
QueryParsingBenchmarks.fromXContent         1024   basic  thrpt    9   16597.576 ±   887.830  ops/s
QueryParsingBenchmarks.fromXContent         1024  filter  thrpt    9   15989.327 ±   826.738  ops/s

## After
Benchmark                            (dimension)  (type)   Mode  Cnt       Score      Error  Units
QueryParsingBenchmarks.fromXContent          128   basic  thrpt    9  111009.726 ± 7980.943  ops/s
QueryParsingBenchmarks.fromXContent          128  filter  thrpt    9  108396.222 ± 5215.041  ops/s
QueryParsingBenchmarks.fromXContent         1024   basic  thrpt    9   15589.731 ±  731.359  ops/s
QueryParsingBenchmarks.fromXContent         1024  filter  thrpt    9   15042.023 ± 1330.928  ops/s

Issues Resolved

[List any issues this PR will resolve]

Check List

Commits are signed as per the DCO using --signoff

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

shatejas

Overall I am content with how this refactoring is going. A few small comments

For method parameter parser, I ended up moving stream in and stream out logic in the parser as well since it was in a way parsing the bytes stream. Wanted to know your thoughts on moving the logic here

shatejas · 2024-07-15T17:22:52Z

src/main/java/org/opensearch/knn/index/query/KNNQueryXContentHelper.java

+ * Helper class responsible for converting XContent to/from KNNQueryBuilders
+ */
+@Log4j2
+public final class KNNQueryXContentHelper {


nit:

I tend to avoid generic suffixes like helper it opens up the possibility for adding things other than what the class is meant for. Consider using KNNQueryParser, this will limit the responsibility of the class to parsing the query and not leak XContent which is implementation detail related to paring in the name of the class

Hmm Im actually okay limiting to XContent only because the toXContent method and fromXContent explicitly require XContent related functions. I thought about KNNQueryParser, but then I thought it wouldnt make sense to include the toXContent code in here, because its not actually parsing.

A parser can have a reverse parsing capability. Its generally an inverse function. something like a serialize and deserialize. Not too hung up on this, will leave it up to you

Sure, I renamed to Parser and updated.

shatejas · 2024-07-15T17:30:41Z

src/main/java/org/opensearch/knn/index/query/KNNQueryXContentHelper.java

+            } else {
+                throwParsingExceptionOnMultipleFields(parser.getTokenLocation(), fieldName, parser.currentName());
+                fieldName = parser.currentName();
+                vector = parser.list();


Is there a need to check if the fieldName.equals(VECTOR_FIELD.getName())

Good catch, that makes sense to have.

Wait, nvm. We do not do this now. I dont want to change behavior in this PR. See https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java#L409

shatejas · 2024-07-15T17:43:17Z

src/test/java/org/opensearch/knn/index/query/KNNQueryXContentHelperTests.java

+import static org.opensearch.knn.index.KNNClusterTestUtils.mockClusterService;
+import static org.opensearch.knn.index.query.KNNQueryXContentHelper.EF_SEARCH_FIELD;
+
+public class KNNQueryXContentHelperTests extends KNNTestCase {


Lets make sure we have the below case. This should hit the final else part in the while loop of fromXcontent

{ knn : { myvector : { vector : [1, 2, 3] } } }

Do we need tests for toXContent() ?

There is a place in the code where it throws for multiple fields, is that case covered in any of the tests?

nit: Consider using a parameterized test, makes it easy to figure out what tests are covered and not covered IMO

Ill take a look. It was mostly just moving the tests from KNNQueryBuilder parsing to here. And Im not sure why there were no tests added for toXContent()

For (1) Interestingly, the above fails because: https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java#L222. This is actually a bug with our defaults (see https://opensearch.org/docs/latest/query-dsl/specialized/neural/). @junqiu-lei Can you take a look around this? Im going to hold off adding the test for now until that is settled on.

Also, I believe the above is the last line in the while loop:

{ knn : { myvector : [1, 2, 3] } }

This also fails with

[knn] requires exactly one of k, distance or score to be set

Just offline discussed with @jmazanec15, the default k only support in neural search. In k-NN plugin, k was one required field before radial search introduced and there wasn't default k value. So the existing validation logic here is correct.

Ref to 2.12 KNNQueryBuilder code before radial search introduced: https://github.com/opensearch-project/k-NN/blob/2.12.0.0/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java

heemin32 · 2024-07-15T16:58:15Z

micro-benchmarks/src/main/java/org/opensearch/knn/QueryParsingBenchmarks.java

@@ -0,0 +1,114 @@
+/*


We are using this copyright

/* * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 */

Will update. I think this is broken: https://github.com/opensearch-project/k-NN/blob/main/.idea/copyright/SPDX_ALv2.xml.

heemin32 · 2024-07-15T17:46:55Z

micro-benchmarks/src/main/java/org/opensearch/knn/QueryParsingBenchmarks.java

+@Measurement(iterations = 3, time = 10)
+@Fork(3)
+@State(Scope.Benchmark)
+public class QueryParsingBenchmarks {


Do we really want to have this file in knn repo?
To make this code useful, I think we should force it to run for every change in query parser and fail if there is a degradation. Otherwise, I don't think it is worth to add it here?

I think it should be kept. I think its fine to have this for testing in the future if we want to make some other change around this functionality.

I think we should force it to run for every change in query parser and fail if there is a degradation.

Sure, this is something we could consider in future, but out of scope for this PR. Want to try to get this change in without having to change multiple other things.

micro-benchmarks/src/main/java/org/opensearch/knn/QueryParsingBenchmarks.java

heemin32 · 2024-07-15T17:48:08Z

micro-benchmarks/src/main/java/org/opensearch/knn/QueryParsingBenchmarks.java

+        );
+    }
+
+    protected static NamedXContentRegistry xContentRegistry() {


Suggested change

protected static NamedXContentRegistry xContentRegistry() {

private NamedXContentRegistry xContentRegistry() {

Not sure this makes sense. We need it for initializing NAMED_X_CONTENT_REGISTRY at the top

heemin32 · 2024-07-15T17:54:01Z

src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java

@@ -138,7 +122,7 @@ public static class Builder {
        private String queryName;
        private float boost = DEFAULT_BOOST;

-        private Builder() {}
+        Builder() {}


Just my preference but I think using either private/public make it more readable unless there is specific necessity to make it package access level.

Suggested change

Builder() {}

public Builder() {}

Will keep for now because I dont want to expose publically, but need it in the other class.

Actually updated to public

heemin32 · 2024-07-15T18:01:51Z

src/main/java/org/opensearch/knn/index/query/KNNQueryXContentHelper.java

+ * Helper class responsible for converting XContent to/from KNNQueryBuilders
+ */
+@Log4j2
+public final class KNNQueryXContentHelper {


I am worried about the pattern where KNNQueryBuilder use KNNQueryXContentHelper and KNNQueryXContentHelper use KNNQueryBuilder which could make cyclic dependency by mistake.

Im not worried about cyclic dependency because helper contains static stateless dependencies. Can you give an example?

Something like KNNBuilder calls one of the method in Helper and helper create a new KNNBuilder which call the method in Helper again. It might not happen but this cyclic dependency always leave a room for such case. It is not comfortable to work on such code base as you have to check that the cyclic calls is not happening when you are modifying the code.

Got it, I dont think this will happen.

Yes. It might not happen but it does look like an anti pattern. Isn't it possible to break it?

KNNQueryBuilder will need to depend on KNNQueryBuilderParser just because it needs to delegate the doXContent method. KNNQueryBuilderParser will need KNNQueryBuilder in other methods in order to convert from KNNQueryBuilder to some stream. Given it is stateless static class, Im not sure its going to break.

heemin32 · 2024-07-15T19:04:41Z

src/main/java/org/opensearch/knn/index/query/KNNQueryXContentHelper.java

+    public static final ParseField EF_SEARCH_FIELD = new ParseField(METHOD_PARAMETER_EF_SEARCH);
+    public static final ParseField METHOD_PARAMS_FIELD = new ParseField(METHOD_PARAMETER);
+
+    private static final ObjectParser<KNNQueryBuilder.Builder, Void> INTERNAL_PARSER = createInternalObjectParser();


Can we avoid this static pattern? I feel this pattern is from legacy OpenSearch code where we don't have dependency injection concept and want to have singleton instance. However, this makes writing a unit test hard and not a best practice in OOP.

I think it would be a little difficult to avoid static pattern in this case. Changing this would then lead to a much larger number of touched files and potential blast radius. I think its something that can be taken in the future.

jmazanec15 · 2024-07-22T17:06:24Z

Waiting till after 2.16 to finish this one up because there are a lot of PRs going on.

jmazanec15 · 2024-07-28T22:52:06Z

For method parameter parser, I ended up moving stream in and stream out logic in the parser as well since it was in a way parsing the bytes stream. Wanted to know your thoughts on moving the logic here

@shatejas Im not 100% sure on this. I leaned toward leaving it in the class just because its a constructor and a method. But I think its something we could move out. However, Ill leave it for another PR just to keep this scope limited.

jmazanec15 · 2024-07-29T15:19:19Z

.idea/copyright/SPDX_ALv2.xml

@@ -1,6 +1,6 @@
 <component name="CopyrightManager">
  <copyright>
-    <option name="notice" value="SPDX-License-Identifier: Apache-2.0&#10;&#10;The OpenSearch Contributors require contributions made to&#10;this file be licensed under the Apache-2.0 license or a&#10;compatible open source license.&#10;&#10;Modifications Copyright OpenSearch Contributors. See&#10;GitHub history for details." />
+    <option name="notice" value="Copyright OpenSearch Contributors&#10;SPDX-License-Identifier: Apache-2.0&#10;" />


Updated this to be in line with https://github.com/opensearch-project/k-NN/pull/1824/files#r1678141021

jmazanec15 · 2024-07-29T15:53:40Z

@heemin32 @shatejas I rebased and addressed some of the feedback. I think a lot of the points are good - but I do want to keep scope limited just so that this PR doesnt die on the vine.

Also, wanted to get opinion on one thing. From #1860 I see that this change will break neural because of the introduction of the Helper class. That being said, I think I should move the logic back into the KNNQueryBuilder class. I was luke-warm on moving out in first place and figure that many plugins out there may be using this as an extension. WDYT?

shatejas · 2024-07-29T18:28:48Z

For method parameter parser, I ended up moving stream in and stream out logic in the parser as well since it was in a way parsing the bytes stream. Wanted to know your thoughts on moving the logic here

@shatejas Im not 100% sure on this. I leaned toward leaving it in the class just because its a constructor and a method. But I think its something we could move out. However, Ill leave it for another PR just to keep this scope limited.

Would like to know regarding the doubts. Having stream in and stream out will help reusing it in neural search and decoupling the changes required on the plugin related to knn.

With respect to scope, we are trying to address a tech debt here, would encourage to not push it to a later especially as it doesn't seem to be a heavy lift. Let me know what you think

jmazanec15 · 2024-07-29T18:55:54Z

Would like to know regarding the doubts. Having stream in and stream out will help reusing it in neural search and decoupling the changes required on the plugin related to knn.

I see. So in neural, we cannot directly re-use KNNQueryBuilder IO because there are different required parameters. But, I will follow convention of MethodParameterParser for KNNQueryBuilder and move streaming to the class. Also, why did you define NPROBE in KNNQueryBuilder then? https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/query/KNNQueryBuilder.java#L78-L79. I think Ill do the same not to break neural (https://github.com/search?q=repo%3Aopensearch-project%2Fneural-search%20KNNQueryBuilder&type=code)

With respect to scope, we are trying to address a tech debt here, would encourage to not push it to a later especially as it doesn't seem to be a heavy lift. Let me know what you think

Want to just be clear - I want to get this PR merged, so I am focusing on limited scope, incremental change. The problem is that as scope gets bigger and PR gets bigger, it ends up becoming more difficult to keep up with and eventually gets abandoned due to other priorities. That being said, Im trying to be a bit defensive in adding on refactors to make sure that that doesnt happen to this PR. For example, not going to change to parametrized tests in this PR.

Refactors parsing of KNNQueryBuilder. First, it moves parsing logic to a separate class. Next, it uses ObjectParser instead of parsing manually by hand. To test, it contains a simple jmh benchmark for testing. Signed-off-by: John Mazanec <jmazane@amazon.com>

heemin32 · 2024-07-29T19:38:50Z

Also, wanted to get opinion on one thing. From #1860 I see that this change will break neural because of the introduction of the Helper class. That being said, I think I should move the logic back into the KNNQueryBuilder class. I was luke-warm on moving out in first place and figure that many plugins out there may be using this as an extension. WDYT?

Not sure for this. KNN is not built to be used by other plugins right? We don't know which class/method is being used by third party and which is not. Shouldn't we ask neural to move again from KNN dependency? Or do we have a set of classes that need to be backward compatible?

jmazanec15 · 2024-07-29T19:41:13Z

Not sure for this. KNN is not built to be used by other plugins right? We don't know which class/method is being used by third party and which is not. Shouldn't we ask neural to move again from KNN dependency? Or do we have a set of classes that need to be backward compatible?

In general, I think keeping KNNQueryBuilder bwc is probably a good thing to do - just because its directly required for neural. Hence, I moved the FIELD declarations back into that class as @shatejas had done for MethodParameters

Also, I just updated to handle streams from the KNNQueryBuilderParser class as well to be consistent with MethodParametersParser.

shatejas · 2024-07-29T20:53:01Z

So in neural, we cannot directly re-use KNNQueryBuilder IO because there are different required parameters

I see, might be possible if we club all knn related fields together, but again we will have to check if that can be done in backward compatible way

Also, why did you define NPROBE in KNNQueryBuilder then?

All fields were defined there, didn't think much back then

opensearch-trigger-bot · 2024-07-29T20:59:51Z

The backport to 2.x failed:

The process '/usr/bin/git' failed with exit code 1

To backport manually, run these commands in your terminal:

# Fetch latest updates from GitHub
git fetch
# Create a new working tree
git worktree add .worktrees/backport-2.x 2.x
# Navigate to the new working tree
cd .worktrees/backport-2.x
# Create a new branch
git switch --create backport/backport-1824-to-2.x
# Cherry-pick the merged commit of this pull request and resolve the conflicts
git cherry-pick -x --mainline 1 52636c47c03ca16da3b2c46860d782b46f192118
# Push it to GitHub
git push --set-upstream origin backport/backport-1824-to-2.x
# Go back to the original working tree
cd ../..
# Delete the working tree
git worktree remove .worktrees/backport-2.x

Then, create a pull request where the base branch is 2.x and the compare/head branch is backport/backport-1824-to-2.x.

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 <jmazane@amazon.com> (cherry picked from commit 52636c4)

jmazanec15 requested review from heemin32, navneet1v, VijayanB, vamshin, naveentatikonda, junqiu-lei, martin-gaievski, ryanbogan and luyuncheng as code owners July 13, 2024 00:19

jmazanec15 added backport 2.x Refactoring Improve the design, structure, and implementation while preserving its functionality labels Jul 13, 2024

jmazanec15 mentioned this pull request Jul 13, 2024

Refactor parsing logic of the Query Builder #1793

Closed

1 task

shatejas reviewed Jul 15, 2024

View reviewed changes

heemin32 reviewed Jul 15, 2024

View reviewed changes

jmazanec15 force-pushed the query-builder-ref branch from 33b7438 to 69941f7 Compare July 29, 2024 15:18

jmazanec15 commented Jul 29, 2024

View reviewed changes

jmazanec15 requested review from heemin32 and shatejas July 29, 2024 18:07

jmazanec15 force-pushed the query-builder-ref branch 2 times, most recently from 0db4ae5 to 6ba54e9 Compare July 29, 2024 19:34

jmazanec15 force-pushed the query-builder-ref branch from 6ba54e9 to dc28549 Compare July 29, 2024 19:37

heemin32 approved these changes Jul 29, 2024

View reviewed changes

shatejas approved these changes Jul 29, 2024

View reviewed changes

junqiu-lei approved these changes Jul 29, 2024

View reviewed changes

jmazanec15 merged commit 52636c4 into opensearch-project:main Jul 29, 2024
52 of 53 checks passed

jmazanec15 mentioned this pull request Jul 29, 2024

[Backport 2.x] Refactor parsing of KNNQueryBuilder #1895

Merged

	protected static NamedXContentRegistry xContentRegistry() {
	private NamedXContentRegistry xContentRegistry() {

Refactor parsing logic of the Query Builder (second try) #1824

Refactor parsing logic of the Query Builder (second try) #1824

Conversation

jmazanec15 commented Jul 13, 2024 • edited Loading

Description

Issues Resolved

Check List

shatejas left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

junqiu-lei Jul 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jmazanec15 commented Jul 22, 2024

jmazanec15 commented Jul 28, 2024

Choose a reason for hiding this comment

jmazanec15 commented Jul 29, 2024

shatejas commented Jul 29, 2024

jmazanec15 commented Jul 29, 2024

heemin32 commented Jul 29, 2024

jmazanec15 commented Jul 29, 2024

shatejas commented Jul 29, 2024

opensearch-trigger-bot bot commented Jul 29, 2024

jmazanec15 commented Jul 13, 2024 •

edited

Loading

junqiu-lei Jul 29, 2024 •

edited

Loading