Skip to content

Commit

Permalink
Integ tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Varun Jain <[email protected]>
  • Loading branch information
@vibrantvarun
vibrantvarun committed Jun 11, 2024
1 parent 291294a commit 1b34723
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 48 deletions.
14 changes: 8 additions & 6 deletions src/test/java/org/opensearch/neuralsearch/query/HybridQuerySortIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,6 @@
import org.junit.Before;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryBuilders;
import org.opensearch.index.query.TermQueryBuilder;
import org.opensearch.index.query.MatchQueryBuilder;
import org.opensearch.index.query.RangeQueryBuilder;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryBuilder;
Expand All @@ -40,7 +38,7 @@ public class HybridQuerySortIT extends BaseNeuralSearchIT {
private static final String TEXT_FIELD_VALUE_4_MI_2 = "Mission Impossible 2";
private static final String TEXT_FIELD_VALUE_5_TERMINAL = "The Terminal";
private static final String TEXT_FIELD_VALUE_6_AVENGERS = "Avengers";
private static final int INTEGER_FIELD_STOCK_1_25 = 24;
private static final int INTEGER_FIELD_STOCK_1_25 = 25;
private static final int INTEGER_FIELD_STOCK_2_22 = 22;
private static final int INTEGER_FIELD_STOCK_3_256 = 256;
private static final int INTEGER_FIELD_STOCK_4_25 = 25;
Expand Down Expand Up @@ -74,6 +72,7 @@ protected boolean preserveClusterUponCompletion() {
@SneakyThrows
public void testSingleFieldSort_whenMultipleSubQueriesOnIndexWithSingleShard_thenSuccessful() {
try {
updateClusterSettings("search.concurrent_segment_search.enabled", false);
prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER);
HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery(
"mission",
Expand All @@ -82,6 +81,9 @@ public void testSingleFieldSort_whenMultipleSubQueriesOnIndexWithSingleShard_the
GTE_OF_RANGE_IN_HYBRID_QUERY
);

// QueryBuilder queryBuilder = QueryBuilders.matchAllQuery();
// .gte(GTE_OF_RANGE_IN_HYBRID_QUERY)
// .lte(LTE_OF_RANGE_IN_HYBRID_QUERY);
Map<String, SortOrder> fieldSortOrderMap = new HashMap<>();
fieldSortOrderMap.put("stock", SortOrder.DESC);

Expand Down Expand Up @@ -273,11 +275,11 @@ public void testSearchAfter_whenAfterFieldIsNotPassed_thenFail() {
}

private HybridQueryBuilder createHybridQueryBuilderWithMatchTermAndRangeQuery(String text, String value, int lte, int gte) {
MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(TEXT_FIELD_1_NAME, text);
TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEXT_FIELD_1_NAME, value);
// MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(TEXT_FIELD_1_NAME, text);
// TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEXT_FIELD_1_NAME, value);
RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(INTEGER_FIELD_1_STOCK).gte(gte).lte(lte);
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
hybridQueryBuilder.add(matchQueryBuilder).add(termQueryBuilder).add(rangeQueryBuilder);
hybridQueryBuilder.add(rangeQueryBuilder);
return hybridQueryBuilder;
}

Expand Down
86 changes: 44 additions & 42 deletions src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -626,23 +626,23 @@ protected void addKnnDoc(
) {
Request request = new Request("POST", "/" + index + "/_doc/" + docId + "?refresh=true");
XContentBuilder builder = XContentFactory.jsonBuilder().startObject();
// for (int i = 0; i < vectorFieldNames.size(); i++) {
// builder.field(vectorFieldNames.get(i), vectors.get(i));
// }
for (int i = 0; i < vectorFieldNames.size(); i++) {
builder.field(vectorFieldNames.get(i), vectors.get(i));
}

for (int i = 0; i < textFieldNames.size(); i++) {
builder.field(textFieldNames.get(i), texts.get(i));
}

// for (int i = 0; i < nestedFieldNames.size(); i++) {
// builder.field(nestedFieldNames.get(i));
// builder.startObject();
// Map<String, String> nestedValues = nestedFields.get(i);
// for (Map.Entry<String, String> entry : nestedValues.entrySet()) {
// builder.field(entry.getKey(), entry.getValue());
// }
// builder.endObject();
// }
for (int i = 0; i < nestedFieldNames.size(); i++) {
builder.field(nestedFieldNames.get(i));
builder.startObject();
Map<String, String> nestedValues = nestedFields.get(i);
for (Map.Entry<String, String> entry : nestedValues.entrySet()) {
builder.field(entry.getKey(), entry.getValue());
}
builder.endObject();
}

for (int i = 0; i < integerFieldNames.size(); i++) {
builder.field(integerFieldNames.get(i), integerFieldValues.get(i));
Expand All @@ -652,9 +652,9 @@ protected void addKnnDoc(
builder.field(keywordFieldNames.get(i), keywordFieldValues.get(i));
}

// for (int i = 0; i < dateFieldNames.size(); i++) {
// builder.field(dateFieldNames.get(i), dateFieldValues.get(i));
// }
for (int i = 0; i < dateFieldNames.size(); i++) {
builder.field(dateFieldNames.get(i), dateFieldValues.get(i));
}
builder.endObject();

request.setJsonEntity(builder.toString());
Expand Down Expand Up @@ -845,37 +845,39 @@ protected String buildIndexConfiguration(
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()
.startObject()
.startObject("settings")
.startObject("index")
.field("number_of_shards", numberOfShards)
// .field("index.knn", true)
.endObject()
.endObject()
.startObject("mappings")
.startObject("properties");

// for (KNNFieldConfig knnFieldConfig : knnFieldConfigs) {
// xContentBuilder.startObject(knnFieldConfig.getName())
// .field("type", "knn_vector")
// .field("dimension", Integer.toString(knnFieldConfig.getDimension()))
// .startObject("method")
// .field("engine", "lucene")
// .field("space_type", knnFieldConfig.getSpaceType().getValue())
// .field("name", "hnsw")
// .endObject()
// .endObject();
// }
for (KNNFieldConfig knnFieldConfig : knnFieldConfigs) {
xContentBuilder.startObject(knnFieldConfig.getName())
.field("type", "knn_vector")
.field("dimension", Integer.toString(knnFieldConfig.getDimension()))
.startObject("method")
.field("engine", "lucene")
.field("space_type", knnFieldConfig.getSpaceType().getValue())
.field("name", "hnsw")
.endObject()
.endObject();
}
// treat the list in a manner that first element is always the type name and all others are keywords
// if (!nestedFields.isEmpty()) {
// String nestedFieldName = nestedFields.get(0);
// xContentBuilder.startObject(nestedFieldName).field("type", "nested");
// if (nestedFields.size() > 1) {
// xContentBuilder.startObject("properties");
// for (int i = 1; i < nestedFields.size(); i++) {
// String innerNestedTypeField = nestedFields.get(i);
// xContentBuilder.startObject(innerNestedTypeField).field("type", "keyword").endObject();
// }
// xContentBuilder.endObject();
// }
// xContentBuilder.endObject();
// }
if (!nestedFields.isEmpty()) {
String nestedFieldName = nestedFields.get(0);
xContentBuilder.startObject(nestedFieldName).field("type", "nested");
if (nestedFields.size() > 1) {
xContentBuilder.startObject("properties");
for (int i = 1; i < nestedFields.size(); i++) {
String innerNestedTypeField = nestedFields.get(i);
xContentBuilder.startObject(innerNestedTypeField).field("type", "keyword").endObject();
}
xContentBuilder.endObject();
}
xContentBuilder.endObject();
}

for (String intField : intFields) {
xContentBuilder.startObject(intField).field("type", "integer").endObject();
Expand All @@ -885,9 +887,9 @@ protected String buildIndexConfiguration(
xContentBuilder.startObject(keywordField).field("type", "keyword").endObject();
}

// for (String dateField : dateFields) {
// xContentBuilder.startObject(dateField).field("type", "date").field("format", "MM/dd/yyyy").endObject();
// }
for (String dateField : dateFields) {
xContentBuilder.startObject(dateField).field("type", "date").field("format", "MM/dd/yyyy").endObject();
}

for (String textField : textFields) {
xContentBuilder.startObject(textField).field("type", "text").endObject();
Expand Down

0 comments on commit 1b34723

Please sign in to comment.