Add WithFieldName implementation to QueryBuilders (opensearch-project#1285)

Signed-off-by: Owais <[email protected]>
Signed-off-by: yeonghyeonKo <[email protected]>

Author: owaiskazi19

src/main/java/org/opensearch/neuralsearch/query/NeuralKNNQueryBuilder.java

@@ -11,6 +11,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.query.AbstractQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.WithFieldName;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.knn.index.query.KNNQueryBuilder;
@@ -48,7 +49,7 @@
  */
 
 @Getter
-public class NeuralKNNQueryBuilder extends AbstractQueryBuilder<NeuralKNNQueryBuilder> {
+public class NeuralKNNQueryBuilder extends AbstractQueryBuilder<NeuralKNNQueryBuilder> implements WithFieldName {
     /**
      * The name of the query
      */
@@ -83,6 +84,7 @@ public static Builder builder() {
      *
      * @return The field name used in the KNN query
      */
+    @Override
     public String fieldName() {
         return knnQueryBuilder.fieldName();
     }

src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java

@@ -38,6 +38,7 @@
 import org.opensearch.index.mapper.NumberFieldMapper;
 import org.opensearch.index.query.AbstractQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.WithFieldName;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.knn.index.query.parser.MethodParametersParser;
@@ -69,7 +70,7 @@
 @Accessors(chain = true, fluent = true)
 @NoArgsConstructor(access = AccessLevel.PRIVATE)
 @AllArgsConstructor(access = AccessLevel.PRIVATE)
-public class NeuralQueryBuilder extends AbstractQueryBuilder<NeuralQueryBuilder> implements ModelInferenceQueryBuilder {
+public class NeuralQueryBuilder extends AbstractQueryBuilder<NeuralQueryBuilder> implements ModelInferenceQueryBuilder, WithFieldName {
 
     public static final String NAME = "neural";
 
@@ -610,4 +611,14 @@ private static boolean validateKNNQueryType(Integer k, Float maxDistance, Float
         }
         return queryCount == 1;
     }
+
+    /**
+     * Gets the field name that this query is searching against.
+     *
+     * @return The field name used in the Neural query
+     */
+    @Override
+    public String fieldName() {
+        return this.fieldName;
+    }
 }

src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java

@@ -41,6 +41,7 @@
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.query.AbstractQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.WithFieldName;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor;
@@ -69,7 +70,10 @@
 @Accessors(chain = true, fluent = true)
 @NoArgsConstructor
 @AllArgsConstructor
-public class NeuralSparseQueryBuilder extends AbstractQueryBuilder<NeuralSparseQueryBuilder> implements ModelInferenceQueryBuilder {
+public class NeuralSparseQueryBuilder extends AbstractQueryBuilder<NeuralSparseQueryBuilder>
+    implements
+        ModelInferenceQueryBuilder,
+        WithFieldName {
     public static final String NAME = "neural_sparse";
     @VisibleForTesting
     static final ParseField QUERY_TEXT_FIELD = new ParseField("query_text");
@@ -540,4 +544,14 @@ private boolean shouldUseModelId() {
     public static float bytesToFloat(byte[] bytes) {
         return ByteBuffer.wrap(bytes).getFloat();
     }
+
+    /**
+     * Gets the field name that this query is searching against.
+     *
+     * @return The field name used in the Neural Sparse query
+     */
+    @Override
+    public String fieldName() {
+        return this.fieldName;
+    }
 }