Implement v_magnitude function

docs/changelog/132765.yaml

@@ -0,0 +1,5 @@
+pr: 132765
+summary: Implement `v_magnitude` function
+area: ES|QL
+type: feature
+issues: []

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/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1356,6 +1356,11 @@ public enum Cap {
          */
         CORRECT_SKIPPED_SHARDS_COUNT,
 
+        /*
+         * Support for calculating the scalar vector magnitude.
+         */
+        MAGNITUDE_SCALAR_VECTOR_FUNCTION(Build.current().isSnapshot()),
+
         /**
          * Byte elements dense vector field type support.
          */

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,179 @@
+/*
+ * 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.compute.data.Block;
+import org.elasticsearch.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.DriverContext;
+import org.elasticsearch.compute.operator.EvalOperator;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
+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.core.type.DataType;
+import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
+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;
+
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isNotNull;
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
+import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
+
+public class Magnitude extends UnaryScalarFunction implements EvaluatorMapper, VectorFunction {
+
+    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 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));
+    }
+
+    @Override
+    public DataType dataType() {
+        return DataType.DOUBLE;
+    }
+
+    @Override
+    protected TypeResolution resolveType() {
+        if (childrenResolved() == false) {
+            return new TypeResolution("Unresolved children");
+        }
+
+        return isNotNull(field(), sourceText(), TypeResolutions.ParamOrdinal.FIRST).and(
+            isType(field(), dt -> dt == DENSE_VECTOR, sourceText(), TypeResolutions.ParamOrdinal.FIRST, "dense_vector")
+        );
+    }
+
+    /**
+     * Functional interface for evaluating the scalar value of the underlying float array.
+     */
+    @FunctionalInterface
+    public interface ScalarEvaluatorFunction {
+        float calculateScalar(float[] scratch);
+    }
+
+    @Override
+    public Object fold(FoldContext ctx) {
+        return EvaluatorMapper.super.fold(source(), ctx);
+    }
+
+    @Override
+    public final EvalOperator.ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator) {
+        return new ScalarEvaluatorFactory(toEvaluator.apply(field()), SCALAR_FUNCTION, getClass().getSimpleName() + "Evaluator");
+    }
+
+    private record ScalarEvaluatorFactory(
+        EvalOperator.ExpressionEvaluator.Factory child,
+        ScalarEvaluatorFunction scalarFunction,
+        String evaluatorName
+    ) implements EvalOperator.ExpressionEvaluator.Factory {
+
+        @Override
+        public EvalOperator.ExpressionEvaluator get(DriverContext context) {
+            // TODO check whether to use this custom evaluator or reuse / define an existing one
+            return new EvalOperator.ExpressionEvaluator() {
+                @Override
+                public Block eval(Page page) {
+                    try (FloatBlock block = (FloatBlock) child.get(context).eval(page);) {
+                        int positionCount = page.getPositionCount();
+                        int dimensions = 0;
+                        // Get the first non-empty vector to calculate the dimension
+                        for (int p = 0; p < positionCount; p++) {
+                            if (block.getValueCount(p) != 0) {
+                                dimensions = block.getValueCount(p);
+                                break;
+                            }
+                        }
+                        if (dimensions == 0) {
+                            return context.blockFactory().newConstantFloatBlockWith(0F, 0);
+                        }
+
+                        float[] scratch = new float[dimensions];
+                        try (var builder = context.blockFactory().newDoubleBlockBuilder(positionCount * dimensions)) {
+                            for (int p = 0; p < positionCount; p++) {
+                                int dims = block.getValueCount(p);
+                                if (dims == 0) {
+                                    // A null value for the vector, by default append 0 as result.
+                                    builder.appendNull();
+                                    continue;
+                                }
+                                readFloatArray(block, block.getFirstValueIndex(p), dimensions, scratch);
+                                float result = scalarFunction.calculateScalar(scratch);
+                                builder.appendDouble(result);
+                            }
+                            return builder.build();
+                        }
+                    }
+                }
+
+                @Override
+                public String toString() {
+                    return evaluatorName() + "[child=" + child + "]";
+                }
+
+                @Override
+                public void close() {}
+            };
+        }
+
+        private static void readFloatArray(FloatBlock block, int position, int dimensions, float[] scratch) {
+            for (int i = 0; i < dimensions; i++) {
+                scratch[i] = block.getFloat(position + i);
+            }
+        }
+
+        @Override
+        public String toString() {
+            return evaluatorName() + "[child=" + child + "]";
+        }
+    }
+}

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

@@ -42,6 +42,9 @@ public static List<NamedWriteableRegistry.Entry> getNamedWritables() {
         if (EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
             entries.add(L2Norm.ENTRY);
         }
+        if (EsqlCapabilities.Cap.MAGNITUDE_SCALAR_VECTOR_FUNCTION.isEnabled()) {
+            entries.add(Magnitude.ENTRY);
+        }
 
         return Collections.unmodifiableList(entries);
     }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java

@@ -59,6 +59,7 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.Concat;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.Substring;
 import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
+import org.elasticsearch.xpack.esql.expression.function.vector.Magnitude;
 import org.elasticsearch.xpack.esql.expression.function.vector.VectorSimilarityFunction;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
@@ -2432,6 +2433,30 @@ private void checkNoDenseVectorFailsSimilarityFunction(String similarityFunction
         );
     }
 
+    public void testMagnitudePlanWithDenseVectorImplicitCasting() {
+        var plan = analyze(String.format(Locale.ROOT, """
+            from test | eval scalar = v_magnitude([1, 2, 3])
+            """), "mapping-dense_vector.json");
+
+        var limit = as(plan, Limit.class);
+        var eval = as(limit.child(), Eval.class);
+        var alias = as(eval.fields().get(0), Alias.class);
+        assertEquals("scalar", alias.name());
+        var scalar = as(alias.child(), Magnitude.class);
+        var child = as(scalar.field(), Literal.class);
+        assertThat(child.dataType(), is(DENSE_VECTOR));
+        assertThat(child.value(), equalTo(List.of(1.0f, 2.0f, 3.0f)));
+    }
+
+    public void testNoDenseVectorFailsForMagnitude() {
+        var query = String.format(Locale.ROOT, "row a = 1 |  eval scalar = v_magnitude(0.342)");
+        VerificationException error = expectThrows(VerificationException.class, () -> analyze(query));
+        assertThat(
+            error.getMessage(),
+            containsString("first argument of [v_magnitude(0.342)] must be [dense_vector], found value [0.342] type [double]")
+        );
+    }
+
     public void testRateRequiresCounterTypes() {
         assumeTrue("rate requires snapshot builds", Build.current().isSnapshot());
         Analyzer analyzer = analyzer(tsdbIndexResolution());