Implement v_magnitude function

x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-magnitude.csv-spec

@@ -0,0 +1,75 @@
+ # Tests for v_magnitude scalar function
+
+ magnitudeWithVectorField
+ required_capability: magnitude_scalar_vector_function
+
+// tag::vector-magnitude[]
+ from colors
+ | eval magnitude = v_magnitude(rgb_vector) 
+ | sort magnitude desc, color asc 
+// end::vector-magnitude[]
+ | limit 10
+ | keep color, magnitude
+ ;
+
+// tag::vector-magnitude-result[]
+color:text | magnitude:double
+white      | 441.6729431152344
+snow       | 435.9185791015625
+azure      | 433.1858825683594
+ivory      | 433.1858825683594
+mint cream | 433.0704345703125
+sea shell  | 426.25579833984375
+honeydew   | 424.5291442871094
+old lace   | 420.6352233886719
+corn silk  | 418.2451477050781
+linen      | 415.93267822265625
+// end::vector-magnitude-result[] 
+;
+
+ magnitudeAsPartOfExpression
+ required_capability: magnitude_scalar_vector_function
+
+ from colors
+ | eval score = round((1 + v_magnitude(rgb_vector) / 2), 3) 
+ | sort score desc, color asc 
+ | limit 10
+ | keep color, score
+ ;
+
+color:text | score:double
+white      | 221.836
+snow       | 218.959
+azure      | 217.593
+ivory      | 217.593
+mint cream | 217.535
+sea shell  | 214.128
+honeydew   | 213.265
+old lace   | 211.318
+corn silk  | 210.123
+linen      | 208.966
+;
+
+magnitudeWithLiteralVectors
+required_capability: magnitude_scalar_vector_function
+
+row a = 1
+| eval magnitude = round(v_magnitude([1, 2, 3]), 3) 
+| keep magnitude
+;
+
+magnitude:double
+3.742
+;
+
+ magnitudeWithStats
+ required_capability: magnitude_scalar_vector_function
+
+ from colors
+ | eval magnitude = round(v_magnitude(rgb_vector), 3) 
+ | stats avg = round(avg(magnitude), 3), min = min(magnitude), max = max(magnitude)
+ ;
+
+avg:double | min:double | max:double
+313.692    | 0.0        | 441.673
+;

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/vector/VectorScalarFunctionsIT.java

@@ -0,0 +1,192 @@
+/*
+ * 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.vector;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.elasticsearch.action.index.IndexRequestBuilder;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.compute.operator.exchange.ExchangeService;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xpack.esql.EsqlTestUtils;
+import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+import org.elasticsearch.xpack.esql.expression.function.vector.Magnitude;
+import org.elasticsearch.xpack.esql.expression.function.vector.VectorScalarFunction;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.Locale;
+
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+
+public class VectorScalarFunctionsIT extends AbstractEsqlIntegTestCase {
+
+    @ParametersFactory
+    public static Iterable<Object[]> parameters() throws Exception {
+        List<Object[]> params = new ArrayList<>();
+        if (EsqlCapabilities.Cap.MAGNITUDE_SCALAR_VECTOR_FUNCTION.isEnabled()) {
+            params.add(new Object[] { "v_magnitude", (VectorScalarFunction.ScalarEvaluatorFunction) Magnitude::calculateScalar });
+        }
+        return params;
+    }
+
+    @Override
+    protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
+        return Settings.builder()
+            .put(super.nodeSettings(nodeOrdinal, otherSettings))
+            // testDifferentDimensions fails the final driver on the coordinator, leading to cancellation of the entire request.
+            // If the exchange sink is opened on a remote node but the compute request hasn't been sent, we cannot close the exchange
+            // sink (for now).Here, we reduce the inactive sinks interval to ensure those inactive sinks are removed quickly.
+            .put(ExchangeService.INACTIVE_SINKS_INTERVAL_SETTING, TimeValue.timeValueMillis(between(3000, 4000)))
+            .build();
+    }
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return CollectionUtils.appendToCopy(super.nodePlugins(), InternalExchangePlugin.class);
+    }
+
+    private final String functionName;
+    private final VectorScalarFunction.ScalarEvaluatorFunction scalarFunction;
+    private int numDims;
+
+    public VectorScalarFunctionsIT(
+        @Name("functionName") String functionName,
+        @Name("scalarFunction") VectorScalarFunction.ScalarEvaluatorFunction scalarFunction
+    ) {
+        this.functionName = functionName;
+        this.scalarFunction = scalarFunction;
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testEvalOverVector() {
+        var query = String.format(Locale.ROOT, """
+                FROM test
+                | EVAL result = %s(vector)
+                | KEEP vector, result
+            """, functionName);
+
+        try (var resp = run(query)) {
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            valuesList.forEach(values -> {
+                float[] v = readVector((List<Float>) values.get(0));
+                Double result = (Double) values.get(1);
+
+                assertNotNull(result);
+                float expected = scalarFunction.calculateScalar(v);
+                assertEquals(expected, result, 0.0001);
+            });
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testEvalOverConstant() {
+        var randomVector = randomVectorArray();
+        var query = String.format(Locale.ROOT, """
+                FROM test
+                | EVAL result = %s(%s)
+                | KEEP vector, result
+            """, functionName, Arrays.toString(randomVector));
+
+        try (var resp = run(query)) {
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            valuesList.forEach(values -> {
+                float[] v = readVector((List<Float>) values.get(0));
+                Double result = (Double) values.get(1);
+
+                assertNotNull(result);
+                float expected = scalarFunction.calculateScalar(randomVector);
+                assertEquals(expected, result, 0.0001);
+            });
+        }
+    }
+
+    private static float[] readVector(List<Float> leftVector) {
+        float[] leftScratch = new float[leftVector.size()];
+        for (int i = 0; i < leftVector.size(); i++) {
+            leftScratch[i] = leftVector.get(i);
+        }
+        return leftScratch;
+    }
+
+    @Before
+    public void setup() throws IOException {
+        assumeTrue("Dense vector type is disabled", EsqlCapabilities.Cap.DENSE_VECTOR_FIELD_TYPE.isEnabled());
+
+        createIndexWithDenseVector("test");
+
+        numDims = randomIntBetween(32, 64) * 2; // min 64, even number
+        int numDocs = randomIntBetween(10, 100);
+        IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
+        for (int i = 0; i < numDocs; i++) {
+            List<Float> v = randomVector();
+            docs[i] = prepareIndex("test").setId("" + i).setSource("id", String.valueOf(i), "vector", v);
+        }
+
+        indexRandom(true, docs);
+    }
+
+    private List<Float> randomVector() {
+        assert numDims != 0 : "numDims must be set before calling randomVector()";
+        List<Float> vector = new ArrayList<>(numDims);
+        for (int j = 0; j < numDims; j++) {
+            vector.add(randomFloat());
+        }
+        return vector;
+    }
+
+    private float[] randomVectorArray() {
+        assert numDims != 0 : "numDims must be set before calling randomVectorArray()";
+        return randomVectorArray(numDims);
+    }
+
+    private static float[] randomVectorArray(int dimensions) {
+        float[] vector = new float[dimensions];
+        for (int j = 0; j < dimensions; j++) {
+            vector[j] = randomFloat();
+        }
+        return vector;
+    }
+
+    private void createIndexWithDenseVector(String indexName) throws IOException {
+        var client = client().admin().indices();
+        XContentBuilder mapping = XContentFactory.jsonBuilder()
+            .startObject()
+            .startObject("properties")
+            .startObject("id")
+            .field("type", "integer")
+            .endObject();
+        createDenseVectorField(mapping, "vector");
+        mapping.endObject().endObject();
+        Settings.Builder settingsBuilder = Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, randomIntBetween(1, 5));
+
+        var CreateRequest = client.prepareCreate(indexName)
+            .setSettings(Settings.builder().put("index.number_of_shards", 1))
+            .setMapping(mapping)
+            .setSettings(settingsBuilder.build());
+        assertAcked(CreateRequest);
+    }
+
+    private void createDenseVectorField(XContentBuilder mapping, String fieldName) throws IOException {
+        mapping.startObject(fieldName).field("type", "dense_vector").field("similarity", "cosine");
+        mapping.endObject();
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1345,7 +1345,12 @@ public enum Cap {
         /**
          * Support correct counting of skipped shards.
          */
-        CORRECT_SKIPPED_SHARDS_COUNT;
+        CORRECT_SKIPPED_SHARDS_COUNT,
+
+        /*
+         * Support for calculating the scalar vector magnitude.
+         */
+        MAGNITUDE_SCALAR_VECTOR_FUNCTION(Build.current().isSnapshot());
 
         private final boolean enabled;
 

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

@@ -188,6 +188,7 @@
 import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
 import org.elasticsearch.xpack.esql.expression.function.vector.L1Norm;
 import org.elasticsearch.xpack.esql.expression.function.vector.L2Norm;
+import org.elasticsearch.xpack.esql.expression.function.vector.Magnitude;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 import org.elasticsearch.xpack.esql.session.Configuration;
 
@@ -503,7 +504,8 @@ private static FunctionDefinition[][] snapshotFunctions() {
                 def(CosineSimilarity.class, CosineSimilarity::new, "v_cosine"),
                 def(DotProduct.class, DotProduct::new, "v_dot_product"),
                 def(L1Norm.class, L1Norm::new, "v_l1_norm"),
-                def(L2Norm.class, L2Norm::new, "v_l2_norm") } };
+                def(L2Norm.class, L2Norm::new, "v_l2_norm"),
+                def(Magnitude.class, Magnitude::new, "v_magnitude") } };
     }
 
     public EsqlFunctionRegistry snapshotRegistry() {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Magnitude.java

@@ -0,0 +1,71 @@
+/*
+ * 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.expression.function.vector;
+
+import org.apache.lucene.util.VectorUtil;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.function.scalar.UnaryScalarFunction;
+import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.expression.function.Example;
+import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesTo;
+import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesToLifecycle;
+import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
+import org.elasticsearch.xpack.esql.expression.function.Param;
+
+import java.io.IOException;
+
+public class Magnitude extends VectorScalarFunction {
+
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Hamming", Magnitude::new);
+    static final ScalarEvaluatorFunction SCALAR_FUNCTION = Magnitude::calculateScalar;
+
+    @FunctionInfo(
+        returnType = "double",
+        preview = true,
+        description = "Calculates the magnitude of a dense_vector.",
+        examples = { @Example(file = "vector-magnitude", tag = "vector-magnitude") },
+        appliesTo = { @FunctionAppliesTo(lifeCycle = FunctionAppliesToLifecycle.DEVELOPMENT) }
+    )
+    public Magnitude(
+        Source source,
+        @Param(name = "input", type = { "dense_vector" }, description = "dense_vector for which to compute the magnitude") Expression input
+    ) {
+        super(source, input);
+    }
+
+    private Magnitude(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    protected UnaryScalarFunction replaceChild(Expression newChild) {
+        return new Magnitude(source(), newChild);
+    }
+
+    @Override
+    protected ScalarEvaluatorFunction getScalarFunction() {
+        return SCALAR_FUNCTION;
+    }
+
+    @Override
+    protected NodeInfo<? extends Expression> info() {
+        return NodeInfo.create(this, Magnitude::new, field());
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
+    public static float calculateScalar(float[] scratch) {
+        return (float) Math.sqrt(VectorUtil.dotProduct(scratch, scratch));
+    }
+}