Small refactoring.

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/inference/InferenceFunction.java

@@ -17,18 +17,17 @@
 
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 
 public abstract class InferenceFunction extends Function {
-
-    private Expression inferenceId;
-
+    private final Expression inferenceId;
     private final Expression options;
 
     @SuppressWarnings("this-escape")
     protected InferenceFunction(Source source, List<Expression> children, Expression options) {
         super(source, Stream.concat(children.stream(), Stream.of(options)).toList());
-        this.inferenceId = parseInferenceId(options);
+        this.inferenceId = parseInferenceId(options, this::defaultInferenceId);
         this.options = options;
     }
 
@@ -40,22 +39,6 @@ public Expression options() {
         return options;
     }
 
-    protected abstract Expression parseInferenceId(Expression options);
-
-    public abstract List<Attribute> temporaryAttributes();
-
-    protected Expression readOption(String optionName, Expression options) {
-        return readOption(optionName, options, Literal.NULL);
-    }
-
-    protected Expression readOption(String optionName, Expression options, Expression defaultValue) {
-        if (options != null && options.dataType() != DataType.NULL && options instanceof MapExpression mapOptions) {
-            return mapOptions.getOrDefault(optionName, defaultValue);
-        }
-
-        return defaultValue;
-    }
-
     @Override
     protected TypeResolution resolveType() {
         if (childrenResolved() == false) {
@@ -65,6 +48,10 @@ protected TypeResolution resolveType() {
         return resolveParams().and(resolveOptions());
     }
 
+    protected abstract Expression defaultInferenceId();
+
+    public abstract List<Attribute> temporaryAttributes();
+
     protected abstract TypeResolution resolveParams();
 
     protected abstract TypeResolution resolveOptions();
@@ -81,4 +68,16 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(super.hashCode(), inferenceId, options);
     }
+
+    private static Expression parseInferenceId(Expression options, Supplier<Expression> defautlInferenceIdSupplier) {
+        return readOption("inference_id", options, defautlInferenceIdSupplier);
+    }
+
+    private static Expression readOption(String optionName, Expression options, Supplier<Expression> defaultValueSupplier) {
+        if (options != null && options.dataType() != DataType.NULL && options instanceof MapExpression mapOptions) {
+            return mapOptions.getOrDefault(optionName, defaultValueSupplier.get());
+        }
+
+        return defaultValueSupplier.get();
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/inference/RerankFunction.java

@@ -114,11 +114,7 @@ public String getWriteableName() {
     }
 
     @Override
-    protected Expression parseInferenceId(Expression options) {
-        return readOption("inference_id", options, defaultInferenceId());
-    }
-
-    private Literal defaultInferenceId() {
+    protected Literal defaultInferenceId() {
         return Literal.keyword(Source.EMPTY, org.elasticsearch.xpack.esql.plan.logical.inference.Rerank.DEFAULT_INFERENCE_ID);
     }
 
@@ -129,7 +125,7 @@ public List<Attribute> temporaryAttributes() {
 
     @Override
     protected TypeResolution resolveParams() {
-        return resolveField().and(resoolveQuery());
+        return resolveField().and(resolveQuery());
     }
 
     @Override
@@ -143,7 +139,7 @@ private TypeResolution resolveField() {
         );
     }
 
-    private TypeResolution resoolveQuery() {
+    private TypeResolution resolveQuery() {
         return isString(query, sourceText(), TypeResolutions.ParamOrdinal.SECOND).and(
             isNotNull(query, sourceText(), TypeResolutions.ParamOrdinal.SECOND)
         ).and(isFoldable(query, sourceText(), TypeResolutions.ParamOrdinal.SECOND));

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/RerankFunctionErrorTests.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.expression.function.inference;
 
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -15,6 +16,7 @@
 import org.elasticsearch.xpack.esql.expression.function.ErrorsForCasesWithoutExamplesTestCase;
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
 import org.hamcrest.Matcher;
+import org.junit.Before;
 
 import java.util.List;
 import java.util.Locale;
@@ -25,6 +27,12 @@
 import static org.hamcrest.Matchers.equalTo;
 
 public class RerankFunctionErrorTests extends ErrorsForCasesWithoutExamplesTestCase {
+
+    @Before
+    public void checkCapability() {
+        assumeTrue("RERANK_FUNCTION is not enabled", EsqlCapabilities.Cap.RERANK_FUNCTION.isEnabled());
+    }
+
     @Override
     protected List<TestCaseSupplier> cases() {
         return paramsToSuppliers(RerankFunctionTests.parameters());

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/RerankFunctionSerializationTests.java

@@ -7,15 +7,22 @@
 
 package org.elasticsearch.xpack.esql.expression.function.inference;
 
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.expression.AbstractExpressionSerializationTests;
+import org.junit.Before;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
 public class RerankFunctionSerializationTests extends AbstractExpressionSerializationTests<RerankFunction> {
+    @Before
+    public void checkCapability() {
+        assumeTrue("RERANK_FUNCTION is not enabled", EsqlCapabilities.Cap.RERANK_FUNCTION.isEnabled());
+    }
+
     @Override
     protected RerankFunction createTestInstance() {
         Source source = randomSource();

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/RerankFunctionTests.java

@@ -12,6 +12,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.xpack.core.security.authc.support.mapper.expressiondsl.FieldExpression;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.MapExpression;
@@ -21,6 +22,7 @@
 import org.elasticsearch.xpack.esql.expression.function.FunctionName;
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
 import org.hamcrest.Matchers;
+import org.junit.Before;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -35,6 +37,11 @@
 
 @FunctionName("rerank")
 public class RerankFunctionTests extends AbstractFunctionTestCase {
+    @Before
+    public void checkCapability() {
+        assumeTrue("RERANK_FUNCTION is not enabled", EsqlCapabilities.Cap.RERANK_FUNCTION.isEnabled());
+    }
+
     public RerankFunctionTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
         this.testCase = testCaseSupplier.get();
     }