Skip to content

Commit e425094

Browse files
mridula-s109mridula-s109
authored
Merge branch 'main' into SEARCH-1006-text-similarity-reranker-does-not-propagate-min-score-correctly
2 parents e5f30a2 + 0702e42 commit e425094

File tree

40 files changed

+427
-148
lines changed

40 files changed

+427
-148
lines changed

‎qa/vector/build.gradle‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
* License v3.0 only", or the "Server Side Public License, v 1".
88
*/
99

10+
import org.elasticsearch.gradle.internal.test.TestUtil
11+
1012
apply plugin: 'elasticsearch.java'
1113
apply plugin: 'elasticsearch.build'
1214

@@ -23,6 +25,8 @@ dependencies {
2325
api "org.apache.lucene:lucene-core:${versions.lucene}"
2426
api "org.apache.lucene:lucene-queries:${versions.lucene}"
2527
api "org.apache.lucene:lucene-codecs:${versions.lucene}"
28+
implementation project(':libs:simdvec')
29+
implementation project(':libs:native')
2630
implementation project(':libs:logging')
2731
implementation project(':server')
2832
}
@@ -37,6 +41,7 @@ tasks.register("checkVec", JavaExec) {
3741
// Configure logging to console
3842
systemProperty "es.logger.out", "console"
3943
systemProperty "es.logger.level", "INFO" // Change to DEBUG if needed
44+
systemProperty 'es.nativelibs.path', TestUtil.getTestLibraryPath(file("../../libs/native/libraries/build/platform/").toString())
4045

4146
if (buildParams.getRuntimeJavaVersion().map { it.majorVersion.toInteger() }.get() >= 21) {
4247
jvmArgs '-Xms4g', '-Xmx4g', '--add-modules=jdk.incubator.vector', '--enable-native-access=ALL-UNNAMED', '-Djava.util.concurrent.ForkJoinPool.common.parallelism=8', '-XX:+UnlockDiagnosticVMOptions', '-XX:+DebugNonSafepoints', '-XX:+HeapDumpOnOutOfMemoryError'

‎qa/vector/src/main/java/module-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
requires org.elasticsearch.base;
1212
requires org.elasticsearch.server;
1313
requires org.elasticsearch.xcontent;
14+
requires org.elasticsearch.cli;
1415
requires org.apache.lucene.core;
1516
requires org.apache.lucene.codecs;
1617
requires org.apache.lucene.queries;

‎qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java‎

Lines changed: 17 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,19 @@
1515
import org.apache.lucene.codecs.KnnVectorsFormat;
1616
import org.apache.lucene.codecs.lucene101.Lucene101Codec;
1717
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
18+
import org.elasticsearch.cli.ProcessInfo;
1819
import org.elasticsearch.common.Strings;
1920
import org.elasticsearch.common.logging.LogConfigurator;
21+
import org.elasticsearch.common.settings.Settings;
2022
import org.elasticsearch.core.PathUtils;
2123
import org.elasticsearch.index.codec.vectors.ES813Int8FlatVectorFormat;
2224
import org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat;
2325
import org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
2426
import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
2527
import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
2628
import org.elasticsearch.logging.Level;
29+
import org.elasticsearch.logging.LogManager;
30+
import org.elasticsearch.logging.Logger;
2731
import org.elasticsearch.xcontent.XContentParser;
2832
import org.elasticsearch.xcontent.XContentParserConfiguration;
2933
import org.elasticsearch.xcontent.XContentType;
@@ -35,19 +39,26 @@
3539
import java.util.ArrayList;
3640
import java.util.List;
3741
import java.util.Locale;
42+
import java.util.Map;
3843

3944
/**
4045
* A utility class to create and test KNN indices using Lucene.
4146
* It supports various index types (HNSW, FLAT, IVF) and configurations.
4247
*/
4348
public class KnnIndexTester {
44-
static final Level LOG_LEVEL = Level.DEBUG;
45-
46-
static final SysOutLogger logger = new SysOutLogger();
49+
static final Logger logger;
4750

4851
static {
4952
LogConfigurator.loadLog4jPlugins();
50-
LogConfigurator.configureESLogging(); // native access requires logging to be initialized
53+
54+
// necessary otherwise the es.logger.level system configuration in build.gradle is ignored
55+
ProcessInfo pinfo = ProcessInfo.fromSystem();
56+
Map<String, String> sysprops = pinfo.sysprops();
57+
String loggerLevel = sysprops.getOrDefault("es.logger.level", Level.INFO.name());
58+
Settings settings = Settings.builder().put("logger.level", loggerLevel).build();
59+
LogConfigurator.configureWithoutConfig(settings);
60+
61+
logger = LogManager.getLogger(KnnIndexTester.class);
5162
}
5263

5364
static final String INDEX_DIR = "target/knn_index";
@@ -163,7 +174,7 @@ public static void main(String[] args) throws Exception {
163174
FormattedResults formattedResults = new FormattedResults();
164175
for (CmdLineArgs cmdLineArgs : cmdLineArgsList) {
165176
Results result = new Results(cmdLineArgs.indexType().name().toLowerCase(Locale.ROOT), cmdLineArgs.numDocs());
166-
System.out.println("Running KNN index tester with arguments: " + cmdLineArgs);
177+
logger.info("Running KNN index tester with arguments: " + cmdLineArgs);
167178
Codec codec = createCodec(cmdLineArgs);
168179
Path indexPath = PathUtils.get(formatIndexPath(cmdLineArgs));
169180
if (cmdLineArgs.reindex() || cmdLineArgs.forceMerge()) {
@@ -195,8 +206,7 @@ public static void main(String[] args) throws Exception {
195206
}
196207
formattedResults.results.add(result);
197208
}
198-
System.out.println("Results:");
199-
System.out.println(formattedResults);
209+
logger.info("Results: \n" + formattedResults);
200210
}
201211

202212
static class FormattedResults {
@@ -326,57 +336,6 @@ static class Results {
326336
}
327337
}
328338

329-
static final class SysOutLogger {
330-
331-
void warn(String message) {
332-
if (LOG_LEVEL.ordinal() >= Level.WARN.ordinal()) {
333-
System.out.println(message);
334-
}
335-
}
336-
337-
void warn(String message, Object... params) {
338-
if (LOG_LEVEL.ordinal() >= Level.WARN.ordinal()) {
339-
System.out.println(String.format(Locale.ROOT, message, params));
340-
}
341-
}
342-
343-
void info(String message) {
344-
if (LOG_LEVEL.ordinal() >= Level.INFO.ordinal()) {
345-
System.out.println(message);
346-
}
347-
}
348-
349-
void info(String message, Object... params) {
350-
if (LOG_LEVEL.ordinal() >= Level.INFO.ordinal()) {
351-
System.out.println(String.format(Locale.ROOT, message, params));
352-
}
353-
}
354-
355-
void debug(String message) {
356-
if (LOG_LEVEL.ordinal() >= Level.DEBUG.ordinal()) {
357-
System.out.println(message);
358-
}
359-
}
360-
361-
void debug(String message, Object... params) {
362-
if (LOG_LEVEL.ordinal() >= Level.DEBUG.ordinal()) {
363-
System.out.println(String.format(Locale.ROOT, message, params));
364-
}
365-
}
366-
367-
void trace(String message) {
368-
if (LOG_LEVEL == Level.TRACE) {
369-
System.out.println(message);
370-
}
371-
}
372-
373-
void trace(String message, Object... params) {
374-
if (LOG_LEVEL == Level.TRACE) {
375-
System.out.println(String.format(Locale.ROOT, message, params));
376-
}
377-
}
378-
}
379-
380339
static final class ThreadDetails {
381340
private static final ThreadMXBean threadBean = (ThreadMXBean) java.lang.management.ManagementFactory.getThreadMXBean();
382341
public final long[] threadIDs;

‎qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexer.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -117,15 +117,15 @@ public boolean isEnabled(String component) {
117117
}
118118
});
119119
logger.debug(
120-
"KnnIndexer: using codec=%s, vectorEncoding=%s, dim=%d, similarityFunction=%s",
120+
"KnnIndexer: using codec={}, vectorEncoding={}, dim={}, similarityFunction={}",
121121
codec.getName(),
122122
vectorEncoding,
123123
dim,
124124
similarityFunction
125125
);
126126

127127
if (Files.exists(indexPath)) {
128-
logger.debug("KnnIndexer: existing index at %s", indexPath);
128+
logger.debug("KnnIndexer: existing index at {}", indexPath);
129129
} else {
130130
Files.createDirectories(indexPath);
131131
}
@@ -143,7 +143,7 @@ public boolean isEnabled(String component) {
143143
);
144144
}
145145
logger.info(
146-
"docsPathSizeInBytes=%d, dim=%d, vectorEncoding=%s, byteSize=%d",
146+
"docsPathSizeInBytes={}, dim={}, vectorEncoding={}, byteSize={}",
147147
docsPathSizeInBytes,
148148
dim,
149149
vectorEncoding,
@@ -170,7 +170,7 @@ public boolean isEnabled(String component) {
170170
}
171171

172172
long elapsed = System.nanoTime() - start;
173-
logger.debug("Indexing took %d ms for %d docs", TimeUnit.NANOSECONDS.toMillis(elapsed), numDocs);
173+
logger.debug("Indexing took {} ms for {} docs", TimeUnit.NANOSECONDS.toMillis(elapsed), numDocs);
174174
result.indexTimeMS = TimeUnit.NANOSECONDS.toMillis(elapsed);
175175
}
176176

@@ -183,14 +183,14 @@ public boolean isEnabled(String component) {
183183
}
184184
});
185185
iwc.setCodec(codec);
186-
logger.debug("KnnIndexer: forceMerge in %s", indexPath);
186+
logger.debug("KnnIndexer: forceMerge in {}", indexPath);
187187
long startNS = System.nanoTime();
188188
try (IndexWriter iw = new IndexWriter(FSDirectory.open(indexPath), iwc)) {
189189
iw.forceMerge(1);
190190
}
191191
long endNS = System.nanoTime();
192192
long elapsedNSec = (endNS - startNS);
193-
logger.info("forceMerge took %d ms", TimeUnit.NANOSECONDS.toMillis(elapsedNSec));
193+
logger.info("forceMerge took {} ms", TimeUnit.NANOSECONDS.toMillis(elapsedNSec));
194194
results.forceMergeTimeMS = TimeUnit.NANOSECONDS.toMillis(elapsedNSec);
195195
}
196196

‎qa/vector/src/main/java/org/elasticsearch/test/knn/KnnSearcher.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ void runSearch(KnnIndexTester.Results finalResults) throws IOException {
181181
resultIds[i] = getResultIds(results[i], storedFields);
182182
}
183183
logger.info(
184-
"completed %d searches in %d ms: %d QPS CPU time=%dms",
184+
"completed {} searches in {} ms: {} QPS CPU time={}ms",
185185
numQueryVectors,
186186
elapsed,
187187
(1000L * numQueryVectors) / elapsed,

‎server/src/internalClusterTest/java/org/elasticsearch/index/shard/IndexShardIT.java‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import org.apache.lucene.index.DirectoryReader;
1212
import org.apache.lucene.store.LockObtainFailedException;
13+
import org.apache.lucene.util.SetOnce;
1314
import org.elasticsearch.ExceptionsHelper;
1415
import org.elasticsearch.action.ActionListener;
1516
import org.elasticsearch.action.index.IndexRequest;
@@ -19,6 +20,8 @@
1920
import org.elasticsearch.cluster.ClusterInfoServiceUtils;
2021
import org.elasticsearch.cluster.ClusterState;
2122
import org.elasticsearch.cluster.InternalClusterInfoService;
23+
import org.elasticsearch.cluster.ShardHeapUsage;
24+
import org.elasticsearch.cluster.ShardHeapUsageCollector;
2225
import org.elasticsearch.cluster.metadata.IndexMetadata;
2326
import org.elasticsearch.cluster.node.DiscoveryNode;
2427
import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
@@ -62,6 +65,7 @@
6265
import org.elasticsearch.indices.IndicesService;
6366
import org.elasticsearch.indices.breaker.CircuitBreakerService;
6467
import org.elasticsearch.indices.recovery.RecoveryState;
68+
import org.elasticsearch.plugins.ClusterPlugin;
6569
import org.elasticsearch.plugins.Plugin;
6670
import org.elasticsearch.search.builder.SearchSourceBuilder;
6771
import org.elasticsearch.test.DummyShardLock;
@@ -82,6 +86,7 @@
8286
import java.util.Collections;
8387
import java.util.List;
8488
import java.util.Locale;
89+
import java.util.Map;
8590
import java.util.Optional;
8691
import java.util.concurrent.BrokenBarrierException;
8792
import java.util.concurrent.CountDownLatch;
@@ -90,6 +95,7 @@
9095
import java.util.concurrent.atomic.AtomicBoolean;
9196
import java.util.concurrent.atomic.AtomicReference;
9297
import java.util.function.Predicate;
98+
import java.util.stream.Collectors;
9399
import java.util.stream.Stream;
94100

95101
import static com.carrotsearch.randomizedtesting.RandomizedTest.randomAsciiLettersOfLength;
@@ -111,12 +117,13 @@
111117
import static org.hamcrest.Matchers.equalTo;
112118
import static org.hamcrest.Matchers.greaterThan;
113119
import static org.hamcrest.Matchers.instanceOf;
120+
import static org.hamcrest.Matchers.lessThanOrEqualTo;
114121

115122
public class IndexShardIT extends ESSingleNodeTestCase {
116123

117124
@Override
118125
protected Collection<Class<? extends Plugin>> getPlugins() {
119-
return pluginList(InternalSettingsPlugin.class);
126+
return pluginList(InternalSettingsPlugin.class, BogusShardHeapUsagePlugin.class);
120127
}
121128

122129
public void testLockTryingToDelete() throws Exception {
@@ -254,6 +261,20 @@ public void testExpectedShardSizeIsPresent() throws InterruptedException {
254261
assertThat(dataSetSize.get(), greaterThan(0L));
255262
}
256263

264+
public void testHeapUsageEstimateIsPresent() {
265+
InternalClusterInfoService clusterInfoService = (InternalClusterInfoService) getInstanceFromNode(ClusterInfoService.class);
266+
ClusterInfoServiceUtils.refresh(clusterInfoService);
267+
ClusterState state = getInstanceFromNode(ClusterService.class).state();
268+
Map<String, ShardHeapUsage> shardHeapUsages = clusterInfoService.getClusterInfo().getShardHeapUsages();
269+
assertNotNull(shardHeapUsages);
270+
assertEquals(state.nodes().size(), shardHeapUsages.size());
271+
for (DiscoveryNode node : state.nodes()) {
272+
assertTrue(shardHeapUsages.containsKey(node.getId()));
273+
ShardHeapUsage shardHeapUsage = shardHeapUsages.get(node.getId());
274+
assertThat(shardHeapUsage.estimatedFreeBytes(), lessThanOrEqualTo(shardHeapUsage.totalBytes()));
275+
}
276+
}
277+
257278
public void testIndexCanChangeCustomDataPath() throws Exception {
258279
final String index = "test-custom-data-path";
259280
final Path sharedDataPath = getInstanceFromNode(Environment.class).sharedDataDir().resolve(randomAsciiLettersOfLength(10));
@@ -797,4 +818,40 @@ private static void assertAllIndicesRemovedAndDeletionCompleted(Iterable<Indices
797818
assertBusy(() -> assertFalse(indicesService.hasUncompletedPendingDeletes()), 1, TimeUnit.MINUTES);
798819
}
799820
}
821+
822+
public static class BogusShardShardHeapUsageCollector implements ShardHeapUsageCollector {
823+
824+
private final BogusShardHeapUsagePlugin plugin;
825+
826+
public BogusShardShardHeapUsageCollector(BogusShardHeapUsagePlugin plugin) {
827+
this.plugin = plugin;
828+
}
829+
830+
@Override
831+
public void collectClusterHeapUsage(ActionListener<Map<String, Long>> listener) {
832+
ActionListener.completeWith(
833+
listener,
834+
() -> plugin.getClusterService()
835+
.state()
836+
.nodes()
837+
.stream()
838+
.collect(Collectors.toUnmodifiableMap(DiscoveryNode::getId, node -> randomNonNegativeLong()))
839+
);
840+
}
841+
}
842+
843+
public static class BogusShardHeapUsagePlugin extends Plugin implements ClusterPlugin {
844+
845+
private final SetOnce<ClusterService> clusterService = new SetOnce<>();
846+
847+
@Override
848+
public Collection<?> createComponents(PluginServices services) {
849+
clusterService.set(services.clusterService());
850+
return List.of();
851+
}
852+
853+
public ClusterService getClusterService() {
854+
return clusterService.get();
855+
}
856+
}
800857
}

‎server/src/internalClusterTest/resources/META-INF/services/org.elasticsearch.cluster.ShardHeapUsageCollector‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
#
2+
# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
# or more contributor license agreements. Licensed under the "Elastic License
4+
# 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
# Public License v 1"; you may not use this file except in compliance with, at
6+
# your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
# License v3.0 only", or the "Server Side Public License, v 1".
8+
#
9+
10+
org.elasticsearch.index.shard.IndexShardIT$BogusShardShardHeapUsageCollector

‎server/src/main/java/org/elasticsearch/TransportVersions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,7 @@ static TransportVersion def(int id) {
293293
public static final TransportVersion SNAPSHOT_INDEX_SHARD_STATUS_MISSING_STATS = def(9_093_0_00);
294294
public static final TransportVersion ML_INFERENCE_ELASTIC_RERANK = def(9_094_0_00);
295295
public static final TransportVersion SEARCH_LOAD_PER_INDEX_STATS = def(9_095_0_00);
296+
public static final TransportVersion HEAP_USAGE_IN_CLUSTER_INFO = def(9_096_0_00);
296297

297298
/*
298299
* STOP! READ THIS FIRST! No, really,

0 commit comments

Comments
 (0)