Skip to content

Commit a6c0bb3

Browse files
neetikasinghalneetikasinghal
authored
Fix sort related ITs for concurrent search (opensearch-project#9466)
Signed-off-by: Neetika Singhal <neetiks@amazon.com>
1 parent f5a6e6d commit a6c0bb3

File tree

6 files changed

+81
-16
lines changed

6 files changed

+81
-16
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
152152
- Add support for wrapping CollectorManager with profiling during concurrent execution ([#9129](https://github.com/opensearch-project/OpenSearch/pull/9129))
153153
- Rethrow OpenSearch exception for non-concurrent path while using concurrent search ([#9177](https://github.com/opensearch-project/OpenSearch/pull/9177))
154154
- Improve performance of encoding composite keys in multi-term aggregations ([#9412](https://github.com/opensearch-project/OpenSearch/pull/9412))
155+
- Fix sort related ITs for concurrent search ([#9177](https://github.com/opensearch-project/OpenSearch/pull/9466)
155156

156157
### Deprecated
157158

‎server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232

3333
package org.opensearch.search.sort;
3434

35+
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
36+
3537
import org.apache.lucene.tests.util.TestUtil;
3638
import org.apache.lucene.util.BytesRef;
3739
import org.apache.lucene.util.UnicodeUtil;
@@ -45,6 +47,7 @@
4547
import org.opensearch.cluster.metadata.IndexMetadata;
4648
import org.opensearch.common.Numbers;
4749
import org.opensearch.common.settings.Settings;
50+
import org.opensearch.common.util.FeatureFlags;
4851
import org.opensearch.common.xcontent.XContentFactory;
4952
import org.opensearch.core.rest.RestStatus;
5053
import org.opensearch.core.xcontent.MediaTypeRegistry;
@@ -60,7 +63,7 @@
6063
import org.opensearch.search.SearchHit;
6164
import org.opensearch.search.SearchHits;
6265
import org.opensearch.test.InternalSettingsPlugin;
63-
import org.opensearch.test.OpenSearchIntegTestCase;
66+
import org.opensearch.test.ParameterizedOpenSearchIntegTestCase;
6467
import org.hamcrest.Matchers;
6568

6669
import java.io.IOException;
@@ -86,6 +89,7 @@
8689
import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
8790
import static org.opensearch.index.query.functionscore.ScoreFunctionBuilders.fieldValueFactorFunction;
8891
import static org.opensearch.script.MockScriptPlugin.NAME;
92+
import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
8993
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
9094
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertFirstHit;
9195
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
@@ -105,7 +109,24 @@
105109
import static org.hamcrest.Matchers.nullValue;
106110
import static org.hamcrest.Matchers.oneOf;
107111

108-
public class FieldSortIT extends OpenSearchIntegTestCase {
112+
public class FieldSortIT extends ParameterizedOpenSearchIntegTestCase {
113+
public FieldSortIT(Settings dynamicSettings) {
114+
super(dynamicSettings);
115+
}
116+
117+
@ParametersFactory
118+
public static Collection<Object[]> parameters() {
119+
return Arrays.asList(
120+
new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), false).build() },
121+
new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), true).build() }
122+
);
123+
}
124+
125+
@Override
126+
protected Settings featureFlagSettings() {
127+
return Settings.builder().put(super.featureFlagSettings()).put(FeatureFlags.CONCURRENT_SEGMENT_SEARCH, "true").build();
128+
}
129+
109130
public static class CustomScriptPlugin extends MockScriptPlugin {
110131
@Override
111132
protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {

‎server/src/internalClusterTest/java/org/opensearch/search/sort/SimpleSortIT.java

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,15 @@
3232

3333
package org.opensearch.search.sort;
3434

35+
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
36+
3537
import org.opensearch.action.index.IndexRequestBuilder;
3638
import org.opensearch.action.search.SearchResponse;
3739
import org.opensearch.action.search.ShardSearchFailure;
3840
import org.opensearch.common.geo.GeoPoint;
3941
import org.opensearch.common.geo.GeoUtils;
42+
import org.opensearch.common.settings.Settings;
43+
import org.opensearch.common.util.FeatureFlags;
4044
import org.opensearch.index.fielddata.ScriptDocValues;
4145
import org.opensearch.plugins.Plugin;
4246
import org.opensearch.script.MockScriptPlugin;
@@ -45,7 +49,7 @@
4549
import org.opensearch.search.SearchHit;
4650
import org.opensearch.search.sort.ScriptSortBuilder.ScriptSortType;
4751
import org.opensearch.test.InternalSettingsPlugin;
48-
import org.opensearch.test.OpenSearchIntegTestCase;
52+
import org.opensearch.test.ParameterizedOpenSearchIntegTestCase;
4953

5054
import java.io.IOException;
5155
import java.util.ArrayList;
@@ -61,6 +65,7 @@
6165
import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
6266
import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
6367
import static org.opensearch.index.query.QueryBuilders.termQuery;
68+
import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
6469
import static org.opensearch.search.sort.SortBuilders.scriptSort;
6570
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
6671
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
@@ -70,10 +75,27 @@
7075
import static org.hamcrest.Matchers.equalTo;
7176
import static org.hamcrest.Matchers.not;
7277

73-
public class SimpleSortIT extends OpenSearchIntegTestCase {
78+
public class SimpleSortIT extends ParameterizedOpenSearchIntegTestCase {
7479

7580
private static final String DOUBLE_APOSTROPHE = "\u0027\u0027";
7681

82+
public SimpleSortIT(Settings dynamicSettings) {
83+
super(dynamicSettings);
84+
}
85+
86+
@ParametersFactory
87+
public static Collection<Object[]> parameters() {
88+
return Arrays.asList(
89+
new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), false).build() },
90+
new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), true).build() }
91+
);
92+
}
93+
94+
@Override
95+
protected Settings featureFlagSettings() {
96+
return Settings.builder().put(super.featureFlagSettings()).put(FeatureFlags.CONCURRENT_SEGMENT_SEARCH, "true").build();
97+
}
98+
7799
@Override
78100
protected Collection<Class<? extends Plugin>> nodePlugins() {
79101
return Arrays.asList(CustomScriptPlugin.class, InternalSettingsPlugin.class);

‎server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/BytesRefFieldComparatorSource.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOEx
9191
return indexFieldData.load(context).getBytesValues();
9292
}
9393

94-
protected void setScorer(Scorable scorer) {}
94+
protected void setScorer(Scorable scorer, LeafReaderContext context) {}
9595

9696
@Override
9797
public FieldComparator<?> newComparator(String fieldname, int numHits, boolean enableSkipping, boolean reversed) {
@@ -134,9 +134,11 @@ protected SortedDocValues getSortedDocValues(LeafReaderContext context, String f
134134
}
135135

136136
return new FieldComparator.TermValComparator(numHits, null, sortMissingLast) {
137+
LeafReaderContext leafReaderContext;
137138

138139
@Override
139140
protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String field) throws IOException {
141+
leafReaderContext = context;
140142
final SortedBinaryDocValues values = getValues(context);
141143
final BinaryDocValues selectedValues;
142144
if (nested == null) {
@@ -152,7 +154,7 @@ protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String f
152154

153155
@Override
154156
public void setScorer(Scorable scorer) {
155-
BytesRefFieldComparatorSource.this.setScorer(scorer);
157+
BytesRefFieldComparatorSource.this.setScorer(scorer, leafReaderContext);
156158
}
157159

158160
};

‎server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/DoubleValuesComparatorSource.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ private NumericDoubleValues getNumericDocValues(LeafReaderContext context, doubl
9595
}
9696
}
9797

98-
protected void setScorer(Scorable scorer) {}
98+
protected void setScorer(Scorable scorer, LeafReaderContext context) {}
9999

100100
@Override
101101
public FieldComparator<?> newComparator(String fieldname, int numHits, boolean enableSkipping, boolean reversed) {
@@ -115,7 +115,7 @@ protected NumericDocValues getNumericDocValues(LeafReaderContext context, String
115115

116116
@Override
117117
public void setScorer(Scorable scorer) {
118-
DoubleValuesComparatorSource.this.setScorer(scorer);
118+
DoubleValuesComparatorSource.this.setScorer(scorer, context);
119119
}
120120
};
121121
}

‎server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,11 @@
6969
import org.opensearch.search.MultiValueMode;
7070

7171
import java.io.IOException;
72+
import java.io.UncheckedIOException;
7273
import java.util.Locale;
74+
import java.util.Map;
7375
import java.util.Objects;
76+
import java.util.concurrent.ConcurrentHashMap;
7477

7578
import static org.opensearch.core.xcontent.ConstructingObjectParser.constructorArg;
7679
import static org.opensearch.search.sort.FieldSortBuilder.validateMaxChildrenExistOnlyInTopLevelNestedSort;
@@ -355,11 +358,19 @@ private IndexFieldData.XFieldComparatorSource fieldComparatorSource(QueryShardCo
355358
final StringSortScript.Factory factory = context.compile(script, StringSortScript.CONTEXT);
356359
final StringSortScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
357360
return new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
358-
StringSortScript leafScript;
361+
// introducing a map to keep a mapping between the leaf reader context and leaf script
362+
// such that the functions of the class are thread safe in case of concurrent search
363+
final Map<LeafReaderContext, StringSortScript> leafContextSortScriptMap = new ConcurrentHashMap<>();
359364

360365
@Override
361366
protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
362-
leafScript = searchScript.newInstance(context);
367+
final StringSortScript leafScript = leafContextSortScriptMap.computeIfAbsent(context, ctx -> {
368+
try {
369+
return searchScript.newInstance(ctx);
370+
} catch (IOException e) {
371+
throw new UncheckedIOException(e);
372+
}
373+
});
363374
final BinaryDocValues values = new AbstractBinaryDocValues() {
364375
final BytesRefBuilder spare = new BytesRefBuilder();
365376

@@ -379,8 +390,8 @@ public BytesRef binaryValue() {
379390
}
380391

381392
@Override
382-
protected void setScorer(Scorable scorer) {
383-
leafScript.setScorer(scorer);
393+
protected void setScorer(Scorable scorer, LeafReaderContext context) {
394+
leafContextSortScriptMap.get(context).setScorer(scorer);
384395
}
385396

386397
@Override
@@ -403,11 +414,19 @@ public BucketedSort newBucketedSort(
403414
final NumberSortScript.Factory numberSortFactory = context.compile(script, NumberSortScript.CONTEXT);
404415
final NumberSortScript.LeafFactory numberSortScript = numberSortFactory.newFactory(script.getParams(), context.lookup());
405416
return new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) {
406-
NumberSortScript leafScript;
417+
// introducing a map to keep a mapping between the leaf reader context and leaf script
418+
// such that the functions of the class are thread safe in case of concurrent search
419+
final Map<LeafReaderContext, NumberSortScript> leafContextSortScriptMap = new ConcurrentHashMap<>();
407420

408421
@Override
409422
protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws IOException {
410-
leafScript = numberSortScript.newInstance(context);
423+
final NumberSortScript leafScript = leafContextSortScriptMap.computeIfAbsent(context, ctx -> {
424+
try {
425+
return numberSortScript.newInstance(ctx);
426+
} catch (IOException e) {
427+
throw new UncheckedIOException(e);
428+
}
429+
});
411430
final NumericDoubleValues values = new NumericDoubleValues() {
412431
@Override
413432
public boolean advanceExact(int doc) throws IOException {
@@ -424,8 +443,8 @@ public double doubleValue() {
424443
}
425444

426445
@Override
427-
protected void setScorer(Scorable scorer) {
428-
leafScript.setScorer(scorer);
446+
protected void setScorer(Scorable scorer, LeafReaderContext context) {
447+
leafContextSortScriptMap.get(context).setScorer(scorer);
429448
}
430449
};
431450
default:

0 commit comments

Comments
 (0)