Ensure casting is done using floats so we get the appropriate blocks and folding done

Author: carlosdelest

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -1613,14 +1613,22 @@ private static Expression castStringLiteral(Expression from, DataType target) {
             }
         }
 
+        @SuppressWarnings("unchecked")
         private static Expression processVectorFunction(org.elasticsearch.xpack.esql.core.expression.function.Function vectorFunction) {
             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) {
-                        Literal denseVector = new Literal(arg.source(), folded, DataType.DENSE_VECTOR);
+                        // 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;
                     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/EvalMapper.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
-import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
 import org.elasticsearch.xpack.esql.evaluator.mapper.ExpressionMapper;
 import org.elasticsearch.xpack.esql.expression.predicate.logical.BinaryLogic;
@@ -249,11 +248,7 @@ private static Block block(Literal lit, BlockFactory blockFactory, int positions
                 if (multiValue.isEmpty()) {
                     return blockFactory.newConstantNullBlock(positions);
                 }
-                // dense_vector create internally float values, even if they are specified as doubles
-                ElementType elementType = lit.dataType() == DataType.DENSE_VECTOR
-                    ? ElementType.FLOAT
-                    : ElementType.fromJava(multiValue.get(0).getClass());
-                var wrapper = BlockUtils.wrapperFor(blockFactory, elementType, positions);
+                var wrapper = BlockUtils.wrapperFor(blockFactory, ElementType.fromJava(multiValue.get(0).getClass()), positions);
                 for (int i = 0; i < positions; i++) {
                     wrapper.accept(multiValue);
                 }

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

@@ -2387,16 +2387,17 @@ public void testDenseVectorImplicitCastingKnn() {
         var field = knn.field();
         var queryVector = as(knn.query(), Literal.class);
         assertEquals(DataType.DENSE_VECTOR, queryVector.dataType());
-        assertThat(queryVector.value(), equalTo(List.of(0.342, 0.164, 0.234)));
+        assertThat(queryVector.value(), equalTo(List.of(0.342f, 0.164f, 0.234f)));
     }
 
     public void testDenseVectorImplicitCastingSimilarityFunctions() {
         if (EsqlCapabilities.Cap.COSINE_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
-            checkDenseVectorImplicitCastingSimilarityFunction("v_cosine(vector, [0.342, 0.164, 0.234])");
+            checkDenseVectorImplicitCastingSimilarityFunction("v_cosine(vector, [0.342, 0.164, 0.234])", List.of(0.342f, 0.164f, 0.234f));
+            checkDenseVectorImplicitCastingSimilarityFunction("v_cosine(vector, [1, 2, 3])", List.of(1f, 2f, 3f));
         }
     }
 
-    private void checkDenseVectorImplicitCastingSimilarityFunction(String similarityFunction) {
+    private void checkDenseVectorImplicitCastingSimilarityFunction(String similarityFunction, List<Number> expectedElems) {
         var plan = analyze(String.format(Locale.ROOT, """
             from test | eval similarity = %s
             """, similarityFunction), "mapping-dense_vector.json");
@@ -2410,8 +2411,7 @@ private void checkDenseVectorImplicitCastingSimilarityFunction(String similarity
         assertEquals("vector", left.name());
         var right = as(similarity.right(), Literal.class);
         assertThat(right.dataType(), is(DENSE_VECTOR));
-        assertThat(right.value(), equalTo(List.of(0.342, 0.164, 0.234)));
-        ;
+        assertThat(right.value(), equalTo(expectedElems));
     }
 
     public void testNoDenseVectorFailsSimilarityFunction() {