Skip to content

Commit 706230a

Browse files
carlosdelestcarlosdelest
committed
Merge remote-tracking branch 'origin/8.19' into backport/8.19/pr-127322
2 parents bd9d066 + beb58e9 commit 706230a

File tree

32 files changed

+810
-155
lines changed

32 files changed

+810
-155
lines changed

‎build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/DependenciesInfoPlugin.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,13 @@
1515
import org.gradle.api.Project;
1616
import org.gradle.api.artifacts.Configuration;
1717
import org.gradle.api.attributes.Category;
18+
import org.gradle.api.attributes.Usage;
1819
import org.gradle.api.plugins.JavaPlugin;
1920

2021
public class DependenciesInfoPlugin implements Plugin<Project> {
22+
23+
public static String USAGE_ATTRIBUTE = "DependenciesInfo";
24+
2125
@Override
2226
public void apply(final Project project) {
2327
project.getPlugins().apply(CompileOnlyResolvePlugin.class);
@@ -43,6 +47,9 @@ public void apply(final Project project) {
4347
)
4448
);
4549

50+
dependenciesInfoFilesConfiguration.attributes(
51+
attributes -> attributes.attribute(Usage.USAGE_ATTRIBUTE, project.getObjects().named(Usage.class, USAGE_ATTRIBUTE))
52+
);
4653
project.getArtifacts().add("dependenciesInfoFiles", depsInfo);
4754

4855
}

‎distribution/build.gradle‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ configurations {
3131
attributes {
3232
attribute(Category.CATEGORY_ATTRIBUTE, project.getObjects().named(Category.class, Category.DOCUMENTATION))
3333
}
34+
attributes {
35+
attribute(Usage.USAGE_ATTRIBUTE, project.getObjects().named(Usage.class, DependenciesInfoPlugin.USAGE_ATTRIBUTE))
36+
}
3437
}
3538
featuresMetadata {
3639
attributes {

‎docs/changelog/128694.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 128694
2+
summary: "Adding Google VertexAI completion integration"
3+
area: Inference
4+
type: enhancement
5+
issues: [ ]

‎muted-tests.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,12 @@ tests:
224224
- class: org.elasticsearch.discovery.ClusterDisruptionIT
225225
method: testAckedIndexing
226226
issue: https://github.com/elastic/elasticsearch/issues/117024
227+
- class: org.elasticsearch.xpack.esql.qa.mixed.MixedClusterEsqlSpecIT
228+
method: test {lookup-join.LookupJoinOnTimeSeriesIndex ASYNC}
229+
issue: https://github.com/elastic/elasticsearch/issues/129078
230+
- class: org.elasticsearch.xpack.esql.qa.mixed.MixedClusterEsqlSpecIT
231+
method: test {lookup-join.LookupJoinOnTimeSeriesIndex SYNC}
232+
issue: https://github.com/elastic/elasticsearch/issues/129078
227233

228234
# Examples:
229235
#

‎test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java‎

Lines changed: 144 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,14 @@
99

1010
package org.elasticsearch.test;
1111

12+
import org.apache.lucene.index.IndexReader;
1213
import org.apache.lucene.search.BoostQuery;
14+
import org.apache.lucene.search.IndexSearcher;
1315
import org.apache.lucene.search.MatchNoDocsQuery;
1416
import org.apache.lucene.search.Query;
1517
import org.apache.lucene.search.TermQuery;
18+
import org.apache.lucene.store.Directory;
19+
import org.apache.lucene.tests.index.RandomIndexWriter;
1620
import org.elasticsearch.ElasticsearchParseException;
1721
import org.elasticsearch.TransportVersion;
1822
import org.elasticsearch.action.support.PlainActionFuture;
@@ -45,6 +49,7 @@
4549
import org.elasticsearch.xcontent.json.JsonStringEncoder;
4650
import org.elasticsearch.xcontent.json.JsonXContent;
4751

52+
import java.io.Closeable;
4853
import java.io.IOException;
4954
import java.time.Instant;
5055
import java.time.ZoneOffset;
@@ -453,82 +458,87 @@ protected boolean builderGeneratesCacheableQueries() {
453458
return true;
454459
}
455460

461+
protected IndexReaderManager getIndexReaderManager() {
462+
return NullIndexReaderManager.INSTANCE;
463+
}
464+
456465
/**
457466
* Test creates the {@link Query} from the {@link QueryBuilder} under test and delegates the
458467
* assertions being made on the result to the implementing subclass.
459468
*/
460469
public void testToQuery() throws IOException {
461470
for (int runs = 0; runs < NUMBER_OF_TESTQUERIES; runs++) {
462-
SearchExecutionContext context = createSearchExecutionContext();
463-
assert context.isCacheable();
464-
context.setAllowUnmappedFields(true);
465-
QB firstQuery = createTestQueryBuilder();
466-
QB controlQuery = copyQuery(firstQuery);
467-
/* we use a private rewrite context here since we want the most realistic way of asserting that we are cacheable or not.
468-
* We do it this way in SearchService where
469-
* we first rewrite the query with a private context, then reset the context and then build the actual lucene query*/
470-
QueryBuilder rewritten = rewriteQuery(firstQuery, createQueryRewriteContext(), new SearchExecutionContext(context));
471-
Query firstLuceneQuery = rewritten.toQuery(context);
472-
assertNotNull("toQuery should not return null", firstLuceneQuery);
473-
assertLuceneQuery(firstQuery, firstLuceneQuery, context);
474-
// remove after assertLuceneQuery since the assertLuceneQuery impl might access the context as well
475-
assertEquals(
476-
"query is not equal to its copy after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
477-
firstQuery,
478-
controlQuery
479-
);
480-
assertEquals(
481-
"equals is not symmetric after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
482-
controlQuery,
483-
firstQuery
484-
);
485-
assertThat(
486-
"query copy's hashcode is different from original hashcode after calling toQuery, firstQuery: "
487-
+ firstQuery
488-
+ ", secondQuery: "
489-
+ controlQuery,
490-
controlQuery.hashCode(),
491-
equalTo(firstQuery.hashCode())
492-
);
493-
494-
QB secondQuery = copyQuery(firstQuery);
495-
// query _name never should affect the result of toQuery, we randomly set it to make sure
496-
if (randomBoolean()) {
497-
secondQuery.queryName(
498-
secondQuery.queryName() == null
499-
? randomAlphaOfLengthBetween(1, 30)
500-
: secondQuery.queryName() + randomAlphaOfLengthBetween(1, 10)
501-
);
502-
}
503-
context = new SearchExecutionContext(context);
504-
Query secondLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context)).toQuery(
505-
context
506-
);
507-
assertNotNull("toQuery should not return null", secondLuceneQuery);
508-
assertLuceneQuery(secondQuery, secondLuceneQuery, context);
509-
510-
if (builderGeneratesCacheableQueries()) {
471+
try (IndexReaderManager irm = getIndexReaderManager()) {
472+
SearchExecutionContext context = createSearchExecutionContext(irm.getIndexSearcher());
473+
assert context.isCacheable();
474+
context.setAllowUnmappedFields(true);
475+
QB firstQuery = createTestQueryBuilder();
476+
QB controlQuery = copyQuery(firstQuery);
477+
/* we use a private rewrite context here since we want the most realistic way of asserting that we are cacheable or not.
478+
* We do it this way in SearchService where
479+
* we first rewrite the query with a private context, then reset the context and then build the actual lucene query*/
480+
QueryBuilder rewritten = rewriteQuery(firstQuery, createQueryRewriteContext(), new SearchExecutionContext(context));
481+
Query firstLuceneQuery = rewritten.toQuery(context);
482+
assertNotNull("toQuery should not return null", firstLuceneQuery);
483+
assertLuceneQuery(firstQuery, firstLuceneQuery, context);
484+
// remove after assertLuceneQuery since the assertLuceneQuery impl might access the context as well
511485
assertEquals(
512-
"two equivalent query builders lead to different lucene queries hashcode",
513-
secondLuceneQuery.hashCode(),
514-
firstLuceneQuery.hashCode()
486+
"query is not equal to its copy after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
487+
firstQuery,
488+
controlQuery
515489
);
516490
assertEquals(
517-
"two equivalent query builders lead to different lucene queries",
518-
rewrite(secondLuceneQuery),
519-
rewrite(firstLuceneQuery)
491+
"equals is not symmetric after calling toQuery, firstQuery: " + firstQuery + ", secondQuery: " + controlQuery,
492+
controlQuery,
493+
firstQuery
494+
);
495+
assertThat(
496+
"query copy's hashcode is different from original hashcode after calling toQuery, firstQuery: "
497+
+ firstQuery
498+
+ ", secondQuery: "
499+
+ controlQuery,
500+
controlQuery.hashCode(),
501+
equalTo(firstQuery.hashCode())
520502
);
521-
}
522503

523-
if (supportsBoost() && firstLuceneQuery instanceof MatchNoDocsQuery == false) {
524-
secondQuery.boost(firstQuery.boost() + 1f + randomFloat());
525-
Query thirdLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
504+
QB secondQuery = copyQuery(firstQuery);
505+
// query _name never should affect the result of toQuery, we randomly set it to make sure
506+
if (randomBoolean()) {
507+
secondQuery.queryName(
508+
secondQuery.queryName() == null
509+
? randomAlphaOfLengthBetween(1, 30)
510+
: secondQuery.queryName() + randomAlphaOfLengthBetween(1, 10)
511+
);
512+
}
513+
context = new SearchExecutionContext(context);
514+
Query secondLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
526515
.toQuery(context);
527-
assertNotEquals(
528-
"modifying the boost doesn't affect the corresponding lucene query",
529-
rewrite(firstLuceneQuery),
530-
rewrite(thirdLuceneQuery)
531-
);
516+
assertNotNull("toQuery should not return null", secondLuceneQuery);
517+
assertLuceneQuery(secondQuery, secondLuceneQuery, context);
518+
519+
if (builderGeneratesCacheableQueries()) {
520+
assertEquals(
521+
"two equivalent query builders lead to different lucene queries hashcode",
522+
secondLuceneQuery.hashCode(),
523+
firstLuceneQuery.hashCode()
524+
);
525+
assertEquals(
526+
"two equivalent query builders lead to different lucene queries",
527+
rewrite(secondLuceneQuery),
528+
rewrite(firstLuceneQuery)
529+
);
530+
}
531+
532+
if (supportsBoost() && firstLuceneQuery instanceof MatchNoDocsQuery == false) {
533+
secondQuery.boost(firstQuery.boost() + 1f + randomFloat());
534+
Query thirdLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
535+
.toQuery(context);
536+
assertNotEquals(
537+
"modifying the boost doesn't affect the corresponding lucene query",
538+
rewrite(firstLuceneQuery),
539+
rewrite(thirdLuceneQuery)
540+
);
541+
}
532542
}
533543
}
534544
}
@@ -938,9 +948,75 @@ public boolean isTextField(String fieldName) {
938948
*/
939949
public void testCacheability() throws IOException {
940950
QB queryBuilder = createTestQueryBuilder();
941-
SearchExecutionContext context = createSearchExecutionContext();
942-
QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, createQueryRewriteContext(), new SearchExecutionContext(context));
943-
assertNotNull(rewriteQuery.toQuery(context));
944-
assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
951+
try (IndexReaderManager irm = getIndexReaderManager()) {
952+
SearchExecutionContext context = createSearchExecutionContext(irm.getIndexSearcher());
953+
QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, createQueryRewriteContext(), new SearchExecutionContext(context));
954+
assertNotNull(rewriteQuery.toQuery(context));
955+
assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
956+
}
957+
}
958+
959+
public static class IndexReaderManager implements Closeable {
960+
private final Directory directory;
961+
private RandomIndexWriter indexWriter;
962+
private IndexReader indexReader;
963+
private IndexSearcher indexSearcher;
964+
965+
public IndexReaderManager() {
966+
this.directory = newDirectory();
967+
}
968+
969+
private IndexReaderManager(Directory directory) {
970+
this.directory = directory;
971+
}
972+
973+
public IndexReader getIndexReader() throws IOException {
974+
if (indexReader == null) {
975+
indexWriter = new RandomIndexWriter(random(), directory);
976+
initIndexWriter(indexWriter);
977+
indexReader = indexWriter.getReader();
978+
}
979+
return indexReader;
980+
}
981+
982+
public IndexSearcher getIndexSearcher() throws IOException {
983+
if (indexSearcher == null) {
984+
indexSearcher = newSearcher(getIndexReader());
985+
}
986+
return indexSearcher;
987+
}
988+
989+
@Override
990+
public void close() throws IOException {
991+
if (indexReader != null) {
992+
indexReader.close();
993+
}
994+
if (indexWriter != null) {
995+
indexWriter.close();
996+
}
997+
if (directory != null) {
998+
directory.close();
999+
}
1000+
}
1001+
1002+
protected void initIndexWriter(RandomIndexWriter indexWriter) {}
1003+
}
1004+
1005+
public static class NullIndexReaderManager extends IndexReaderManager {
1006+
public static final NullIndexReaderManager INSTANCE = new NullIndexReaderManager();
1007+
1008+
public NullIndexReaderManager() {
1009+
super(null);
1010+
}
1011+
1012+
@Override
1013+
public IndexReader getIndexReader() {
1014+
return null;
1015+
}
1016+
1017+
@Override
1018+
public IndexSearcher getIndexSearcher() {
1019+
return null;
1020+
}
9451021
}
9461022
}

‎x-pack/plugin/core/template-resources/src/main/resources/[email protected]‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@
4242
"jobtype": {
4343
"type": "keyword"
4444
},
45+
"scheduled_report_id": {
46+
"type": "keyword"
47+
},
4548
"payload": {
4649
"type": "object",
4750
"enabled": false

‎x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
134134

135135
public void testGetServicesWithCompletionTaskType() throws IOException {
136136
List<Object> services = getServices(TaskType.COMPLETION);
137-
assertThat(services.size(), equalTo(14));
137+
assertThat(services.size(), equalTo(15));
138138

139139
var providers = providers(services);
140140

@@ -155,6 +155,7 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
155155
"streaming_completion_test_service",
156156
"hugging_face",
157157
"amazon_sagemaker",
158+
"googlevertexai",
158159
"mistral"
159160
).toArray()
160161
)

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiModel.java‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ public abstract class GoogleVertexAiModel extends RateLimitGroupingModel {
2424

2525
private final GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings;
2626

27-
protected URI uri;
27+
protected URI nonStreamingUri;
2828

2929
public GoogleVertexAiModel(
3030
ModelConfigurations configurations,
@@ -39,14 +39,14 @@ public GoogleVertexAiModel(
3939
public GoogleVertexAiModel(GoogleVertexAiModel model, ServiceSettings serviceSettings) {
4040
super(model, serviceSettings);
4141

42-
uri = model.uri();
42+
nonStreamingUri = model.nonStreamingUri();
4343
rateLimitServiceSettings = model.rateLimitServiceSettings();
4444
}
4545

4646
public GoogleVertexAiModel(GoogleVertexAiModel model, TaskSettings taskSettings) {
4747
super(model, taskSettings);
4848

49-
uri = model.uri();
49+
nonStreamingUri = model.nonStreamingUri();
5050
rateLimitServiceSettings = model.rateLimitServiceSettings();
5151
}
5252

@@ -56,17 +56,8 @@ public GoogleVertexAiRateLimitServiceSettings rateLimitServiceSettings() {
5656
return rateLimitServiceSettings;
5757
}
5858

59-
public URI uri() {
60-
return uri;
61-
}
62-
63-
@Override
64-
public int rateLimitGroupingHash() {
65-
// In VertexAI rate limiting is scoped to the project, region and model. URI already has this information so we are using that.
66-
// API Key does not affect the quota
67-
// https://ai.google.dev/gemini-api/docs/rate-limits
68-
// https://cloud.google.com/vertex-ai/docs/quotas
69-
return Objects.hash(uri);
59+
public URI nonStreamingUri() {
60+
return nonStreamingUri;
7061
}
7162

7263
@Override

0 commit comments

Comments
 (0)