diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java index e5b6b77515ad3..4c9fd5bd62ef0 100644 --- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java @@ -350,32 +350,6 @@ public void testScoringKeepAndSort() { } } - public void testRrf() { - assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled()); - - var query = """ - FROM test METADATA _score, _id, _index - | WHERE id > 2 - | FORK - ( WHERE content:"fox" | SORT _score, _id DESC ) - ( WHERE content:"dog" | SORT _score, _id DESC ) - | RRF - | EVAL _score = round(_score, 4) - | KEEP id, content, _score, _fork - """; - try (var resp = run(query)) { - assertColumnNames(resp.columns(), List.of("id", "content", "_score", "_fork")); - assertColumnTypes(resp.columns(), List.of("integer", "keyword", "double", "keyword")); - assertThat(getValuesList(resp.values()).size(), equalTo(3)); - Iterable> expectedValues = List.of( - List.of(6, "The quick brown fox jumps over the lazy dog", 0.0325, List.of("fork1", "fork2")), - List.of(4, "The dog is brown but this document is very very long", 0.0164, "fork2"), - List.of(3, "This dog is really brown", 0.0159, "fork2") - ); - assertValues(resp.values(), expectedValues); - } - } - public void testThreeSubQueries() { var query = """ FROM test diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfIT.java new file mode 100644 index 0000000000000..3ffd235882886 --- /dev/null +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfIT.java @@ -0,0 +1,77 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.action; + +import org.elasticsearch.action.index.IndexRequest; +import org.elasticsearch.action.support.WriteRequest; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.plugins.Plugin; +import org.junit.Before; + +import java.util.Collection; +import java.util.List; + +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; +import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList; +import static org.hamcrest.Matchers.equalTo; + +public class RrfIT extends AbstractEsqlIntegTestCase { + @Override + protected Collection> nodePlugins() { + return List.of(EsqlPluginWithEnterpriseOrTrialLicense.class); + } + + @Before + public void setupIndex() { + assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled()); + createAndPopulateIndex(); + } + + public void testRrf() { + var query = """ + FROM test METADATA _score, _id, _index + | WHERE id > 2 + | FORK + ( WHERE content:"fox" | SORT _score, _id DESC ) + ( WHERE content:"dog" | SORT _score, _id DESC ) + | RRF + | EVAL _score = round(_score, 4) + | KEEP id, content, _score, _fork + """; + try (var resp = run(query)) { + assertColumnNames(resp.columns(), List.of("id", "content", "_score", "_fork")); + assertColumnTypes(resp.columns(), List.of("integer", "keyword", "double", "keyword")); + assertThat(getValuesList(resp.values()).size(), equalTo(3)); + Iterable> expectedValues = List.of( + List.of(6, "The quick brown fox jumps over the lazy dog", 0.0325, List.of("fork1", "fork2")), + List.of(4, "The dog is brown but this document is very very long", 0.0164, "fork2"), + List.of(3, "This dog is really brown", 0.0159, "fork2") + ); + assertValues(resp.values(), expectedValues); + } + } + + private void createAndPopulateIndex() { + var indexName = "test"; + var client = client().admin().indices(); + var CreateRequest = client.prepareCreate(indexName) + .setSettings(Settings.builder().put("index.number_of_shards", 1)) + .setMapping("id", "type=integer", "content", "type=text"); + assertAcked(CreateRequest); + client().prepareBulk() + .add(new IndexRequest(indexName).id("1").source("id", 1, "content", "This is a brown fox")) + .add(new IndexRequest(indexName).id("2").source("id", 2, "content", "This is a brown dog")) + .add(new IndexRequest(indexName).id("3").source("id", 3, "content", "This dog is really brown")) + .add(new IndexRequest(indexName).id("4").source("id", 4, "content", "The dog is brown but this document is very very long")) + .add(new IndexRequest(indexName).id("5").source("id", 5, "content", "There is also a white cat")) + .add(new IndexRequest(indexName).id("6").source("id", 6, "content", "The quick brown fox jumps over the lazy dog")) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .get(); + ensureYellow(indexName); + } +} diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfWithInvalidLicenseIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfWithInvalidLicenseIT.java new file mode 100644 index 0000000000000..335a738b0a512 --- /dev/null +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfWithInvalidLicenseIT.java @@ -0,0 +1,62 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.action; + +import org.elasticsearch.ElasticsearchException; +import org.elasticsearch.action.index.IndexRequestBuilder; +import org.elasticsearch.action.support.WriteRequest; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.xpack.esql.VerificationException; +import org.junit.Before; + +import java.util.Collection; +import java.util.List; + +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; +import static org.hamcrest.Matchers.containsString; + +public class RrfWithInvalidLicenseIT extends AbstractEsqlIntegTestCase { + private static final String LICENSE_ERROR_MESSAGE = "current license is non-compliant for [RRF]"; + + @Override + protected Collection> nodePlugins() { + return List.of(EsqlPluginWithNonEnterpriseOrExpiredLicense.class); + } + + @Before + public void setupIndex() { + assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled()); + var indexName = "test"; + var client = client().admin().indices(); + var CreateRequest = client.prepareCreate(indexName) + .setSettings(Settings.builder().put("index.number_of_shards", 1)) + .setMapping("id", "type=integer", "content", "type=text"); + assertAcked(CreateRequest); + client().prepareBulk() + .add(new IndexRequestBuilder(client, indexName).setId("1").setSource("id", 1, "content", "This is a brown fox")) + .add(new IndexRequestBuilder(client, indexName).setId("2").setSource("id", 2, "content", "This is a brown dog")) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .get(); + ensureYellow(indexName); + } + + public void testRrf() { + var query = """ + FROM test METADATA _score, _id, _index + | FORK + ( WHERE content:"fox" ) + ( WHERE content:"dog" ) + | RRF + """; + + ElasticsearchException e = expectThrows(VerificationException.class, () -> run(query)); + + assertThat(e.getMessage(), containsString(LICENSE_ERROR_MESSAGE)); + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RrfScoreEval.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RrfScoreEval.java index 9647e4c5885e1..afd7a69b5123b 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RrfScoreEval.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/RrfScoreEval.java @@ -8,6 +8,9 @@ package org.elasticsearch.xpack.esql.plan.logical; import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.license.License; +import org.elasticsearch.license.XPackLicenseState; +import org.elasticsearch.xpack.esql.LicenseAware; import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware; import org.elasticsearch.xpack.esql.common.Failures; import org.elasticsearch.xpack.esql.core.expression.Attribute; @@ -20,7 +23,7 @@ import static org.elasticsearch.xpack.esql.common.Failure.fail; -public class RrfScoreEval extends UnaryPlan implements PostAnalysisVerificationAware { +public class RrfScoreEval extends UnaryPlan implements PostAnalysisVerificationAware, LicenseAware { private final Attribute forkAttr; private final Attribute scoreAttr; @@ -93,4 +96,9 @@ public boolean equals(Object obj) { RrfScoreEval rrf = (RrfScoreEval) obj; return child().equals(rrf.child()) && scoreAttr.equals(rrf.scoreAttribute()) && forkAttr.equals(forkAttribute()); } + + @Override + public boolean licenseCheck(XPackLicenseState state) { + return state.isAllowedByLicense(License.OperationMode.ENTERPRISE); + } }