Add tests for casting using ToDenseVector

carlosdelest · carlosdelest · commit 7383dc98f019 · 2025-09-02T14:11:23.000+02:00
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/knn-function.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/knn-function.csv-spec
@@ -1,7 +1,3 @@
-# TODO Most tests explicitly set k. Until knn function uses LIMIT as k, we need to explicitly set it to all values
-# in the dataset to avoid test failures due to docs allocation in different shards, which can impact results for a 
-# top-n query at the shard level 
-
 knnSearch
 required_capability: knn_function_v3
 
@@ -306,3 +302,28 @@ c: long       | primary: boolean
 41            | false          
 9             | true           
 ;
+
+knnWithCasting
+required_capability: knn_function_v3
+required_capability: to_dense_vector_function
+
+from colors metadata _score
+| eval query = [0, 120, 0]
+| where knn(rgb_vector, query, 10) 
+| sort _score desc, color asc
+| keep color, rgb_vector
+| limit 10
+;
+
+color:text | rgb_vector:dense_vector
+green      | [0.0, 128.0, 0.0]
+black      | [0.0, 0.0, 0.0]
+olive      | [128.0, 128.0, 0.0]
+teal       | [0.0, 128.0, 128.0]
+lime       | [0.0, 255.0, 0.0]
+sienna     | [160.0, 82.0, 45.0]
+maroon     | [128.0, 0.0, 0.0]
+navy       | [0.0, 0.0, 128.0]
+gray       | [128.0, 128.0, 128.0]
+chartreuse | [127.0, 255.0, 0.0]
+;
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-cosine-similarity.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-cosine-similarity.csv-spec
@@ -90,17 +90,40 @@ total_null:long
 59
 ;
 
-# TODO Need to implement a conversion function to convert a non-foldable row to a dense_vector
-similarityWithRow-Ignore
+similarityWithRow
 required_capability: cosine_vector_similarity_function
+required_capability: to_dense_vector_function
  
 row vector = [1, 2, 3] 
 | eval similarity = round(v_cosine(vector, [0, 1, 2]), 3) 
+;
+
+vector: integer      | similarity:double
+[1, 2, 3]      | 0.978  
+;
+
+similarityWithVectorField
+required_capability: cosine_vector_similarity_function
+required_capability: to_dense_vector_function
+
+from colors
+| where color != "black" 
+| eval query = [0, 255, 255]
+| eval similarity = v_cosine(rgb_vector, query) 
 | sort similarity desc, color asc 
 | limit 10
 | keep color, similarity
 ;
 
-similarity:double
-0.978  
+color:text     | similarity:double
+cyan           | 1.0
+teal           | 1.0
+turquoise      | 0.9890533685684204
+aqua marine    | 0.964962363243103
+azure          | 0.916246771812439
+lavender       | 0.9136701822280884
+mint cream     | 0.9122757911682129
+honeydew       | 0.9122424125671387
+gainsboro      | 0.9082483053207397
+gray           | 0.9082483053207397  
 ;
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-dot-product.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-dot-product.csv-spec
@@ -88,17 +88,39 @@ total_null:long
 ;
 
 
-# TODO Need to implement a conversion function to convert a non-foldable row to a dense_vector
-similarityWithRow-Ignore
+similarityWithRow
 required_capability: dot_product_vector_similarity_function
+required_capability: to_dense_vector_function
  
 row vector = [1, 2, 3] 
 | eval similarity = round(v_dot_product(vector, [0, 1, 2]), 3) 
+;
+
+vector: integer | similarity:double
+[1, 2, 3]       | 4.5
+;
+
+similarityWithVectorField
+required_capability: dot_product_vector_similarity_function
+required_capability: to_dense_vector_function
+
+from colors
+| eval query = [0, 255, 255]
+| eval similarity = v_dot_product(rgb_vector, query) 
 | sort similarity desc, color asc 
 | limit 10
 | keep color, similarity
 ;
 
-similarity:double
-0.978  
+color:text | similarity:double
+azure      | 65025.5
+cyan       | 65025.5
+white      | 65025.5
+mint cream | 64388.0
+snow       | 63750.5
+honeydew   | 63113.0
+ivory      | 63113.0
+sea shell  | 61583.0
+lavender   | 61200.5
+old lace   | 60563.0
 ;
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-hamming.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-hamming.csv-spec
@@ -87,17 +87,39 @@ total_null:long
 59
 ;
 
-# TODO Need to implement a conversion function to convert a non-foldable row to a dense_vector
-similarityWithRow-Ignore
+similarityWithRow
 required_capability: hamming_vector_similarity_function
+required_capability: to_dense_vector_function
  
 row vector = [1, 2, 3] 
 | eval similarity = round(v_hamming(vector, [0, 1, 2]), 3) 
+;
+
+vector: integer  | similarity:double
+[1, 2, 3]        | 4.0
+;
+
+similarityWithVectorField
+required_capability: hamming_vector_similarity_function
+required_capability: to_dense_vector_function
+
+from colors
+| eval query = [0, 255, 255]
+| eval similarity = v_hamming(rgb_vector, query) 
 | sort similarity desc, color asc 
 | limit 10
 | keep color, similarity
 ;
-
-similarity:double
-0.978  
+ 
+color:text | similarity:double
+red        | 24.0
+orange     | 20.0
+gold       | 18.0
+indigo     | 18.0
+bisque     | 17.0
+maroon     | 17.0
+pink       | 17.0
+salmon     | 17.0
+black      | 16.0
+firebrick  | 16.0
 ;
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-l1-norm.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-l1-norm.csv-spec
@@ -87,17 +87,39 @@ total_null:long
 59
 ;
 
-# TODO Need to implement a conversion function to convert a non-foldable row to a dense_vector
-similarityWithRow-Ignore
+similarityWithRow
 required_capability: l1_norm_vector_similarity_function
+required_capability: to_dense_vector_function
  
 row vector = [1, 2, 3] 
 | eval similarity = round(v_l1_norm(vector, [0, 1, 2]), 3) 
+;
+
+vector: integer  | similarity:double
+[1, 2, 3]        | 3.0
+;
+
+similarityWithVectorField
+required_capability: l1_norm_vector_similarity_function
+required_capability: to_dense_vector_function
+
+from colors
+| eval query = [0, 255, 255]
+| eval similarity = v_l1_norm(rgb_vector, query) 
 | sort similarity desc, color asc 
 | limit 10
 | keep color, similarity
 ;
-
-similarity:double
-0.978  
+ 
+color:text | similarity:double
+red        | 765.0
+crimson    | 650.0
+maroon     | 638.0
+firebrick  | 620.0
+orange     | 600.0
+tomato     | 595.0
+brown      | 591.0
+chocolate  | 585.0
+coral      | 558.0
+gold       | 550.0
 ;
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-l2-norm.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-l2-norm.csv-spec
@@ -87,17 +87,39 @@ total_null:long
 59
 ;
 
-# TODO Need to implement a conversion function to convert a non-foldable row to a dense_vector
-similarityWithRow-Ignore
+similarityWithRow
 required_capability: l2_norm_vector_similarity_function
+required_capability: to_dense_vector_function
  
 row vector = [1, 2, 3] 
 | eval similarity = round(v_l2_norm(vector, [0, 1, 2]), 3) 
+;
+
+vector: integer  | similarity:double
+[1, 2, 3]        | 1.732
+;
+
+similarityWithVectorField
+required_capability: l2_norm_vector_similarity_function
+required_capability: to_dense_vector_function
+
+from colors
+| eval query = [0, 255, 255]
+| eval similarity = v_l2_norm(rgb_vector, query) 
 | sort similarity desc, color asc 
 | limit 10
 | keep color, similarity
 ;
-
-similarity:double
-0.978  
+ 
+color:text | similarity:double
+red        | 441.6729431152344
+maroon     | 382.6669616699219
+crimson    | 376.36419677734375
+orange     | 371.68536376953125
+gold       | 362.8360595703125
+black      | 360.62445068359375
+magenta    | 360.62445068359375
+yellow     | 360.62445068359375
+firebrick  | 359.67486572265625
+tomato     | 351.0227966308594
 ;
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-magnitude.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-magnitude.csv-spec
@@ -85,3 +85,15 @@ row a = 1
 magnitude:double
 null
 ;
+
+magnitudeWithRow
+required_capability: magnitude_scalar_vector_function
+required_capability: to_dense_vector_function
+ 
+row vector = [1, 2, 3] 
+| eval magnitude = round(v_magnitude(vector), 3) 
+;
+
+vector: integer      | magnitude:double
+[1, 2, 3]      | 3.742
+;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java
@@ -1416,7 +1416,12 @@ public enum Cap {
         /**
          * URL encoding function.
          */
-        URL_ENCODE(Build.current().isSnapshot());
+        URL_ENCODE(Build.current().isSnapshot()),
+
+        /**
+         * TO_DENSE_VECTOR function.
+         */
+        TO_DENSE_VECTOR_FUNCTION(Build.current().isSnapshot());
 
         private final boolean enabled;
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java
@@ -76,6 +76,7 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromAggregateMetricDouble;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToAggregateMetricDouble;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDateNanos;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDenseVector;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDouble;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToInteger;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
@@ -168,6 +169,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.TIME_DURATION;
 import static org.elasticsearch.xpack.esql.core.type.DataType.UNSUPPORTED;
 import static org.elasticsearch.xpack.esql.core.type.DataType.VERSION;
+import static org.elasticsearch.xpack.esql.core.type.DataType.isString;
 import static org.elasticsearch.xpack.esql.core.type.DataType.isTemporalAmount;
 import static org.elasticsearch.xpack.esql.telemetry.FeatureMetric.LIMIT;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.maybeParseTemporalAmount;
@@ -1668,18 +1670,24 @@ private static Expression processVectorFunction(org.elasticsearch.xpack.esql.cor
             List<Expression> args = vectorFunction.arguments();
             List<Expression> newArgs = new ArrayList<>();
             for (Expression arg : args) {
-                if (arg.resolved() && arg.dataType().isNumeric() && arg.foldable()) {
-                    Object folded = arg.fold(FoldContext.small() /* TODO remove me */);
-                    if (folded instanceof List) {
-                        // Convert to floats so blocks are created accordingly
-                        List<Float> floatVector;
-                        if (arg.dataType() == FLOAT) {
-                            floatVector = (List<Float>) folded;
-                        } else {
-                            floatVector = ((List<Number>) folded).stream().map(Number::floatValue).collect(Collectors.toList());
+                if (arg.resolved()) {
+                    if (arg.foldable() && arg.dataType().isNumeric()) {
+                        Object folded = arg.fold(FoldContext.small() /* TODO remove me */);
+                        if (folded instanceof List) {
+                            // Convert to floats so blocks are created accordingly
+                            List<Float> floatVector;
+                            if (arg.dataType() == FLOAT) {
+                                floatVector = (List<Float>) folded;
+                            } else {
+                                floatVector = ((List<Number>) folded).stream().map(Number::floatValue).collect(Collectors.toList());
+                            }
+                            Literal denseVector = new Literal(arg.source(), floatVector, DataType.DENSE_VECTOR);
+                            newArgs.add(denseVector);
+                            continue;
                         }
-                        Literal denseVector = new Literal(arg.source(), floatVector, DataType.DENSE_VECTOR);
-                        newArgs.add(denseVector);
+                    } else if (arg.dataType().isNumeric() || isString(arg.dataType())) {
+                        // add casting function
+                        newArgs.add(new ToDenseVector(arg.source(), arg));
                         continue;
                     }
                 }
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java
