Esql resultset dense vector (#1208) (#1210)

* make result set support dense vector

* improve test

Co-authored-by: Laura Trotta <153528055+l-trotta@users.noreply.github.com>

Author: github-actions[bot]

java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/jdbc/EsType.java

@@ -61,7 +61,8 @@ enum EsType implements SQLType {
     GEO_SHAPE(ExtraTypes.GEOMETRY),
     SHAPE(ExtraTypes.GEOMETRY),
     UNSIGNED_LONG(Types.NUMERIC),
-    VERSION(Types.VARCHAR);
+    VERSION(Types.VARCHAR),
+    DENSE_VECTOR(Types.ARRAY);
 
     private final Integer type;
 

java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/jdbc/TypeConverter.java

@@ -211,6 +211,7 @@ static Object convert(Object v, EsType columnType, String typeString) throws SQL
             case BOOLEAN:
             case TEXT:
             case KEYWORD:
+            case DENSE_VECTOR:
                 return v; // These types are already represented correctly in JSON
             case BYTE:
                 return ((Number) v).byteValue(); // Parser might return it as integer or long - need to update to the correct type

java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/jdbc/TypeUtils.java

@@ -31,6 +31,7 @@
 import java.util.EnumSet;
 import java.util.GregorianCalendar;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
@@ -117,6 +118,7 @@ private TypeUtils() {}
         types.put(EsType.GEO_POINT, String.class);
         types.put(EsType.GEO_SHAPE, String.class);
         types.put(EsType.SHAPE, String.class);
+        types.put(EsType.DENSE_VECTOR, List.class);
 
         TYPE_TO_CLASS = unmodifiableMap(types);
 

java-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlAdapterTest.java

@@ -22,17 +22,17 @@
 import co.elastic.clients.elasticsearch.ElasticsearchClient;
 import co.elastic.clients.elasticsearch._helpers.esql.jdbc.ResultSetEsqlAdapter;
 import co.elastic.clients.elasticsearch._helpers.esql.objects.ObjectsEsqlAdapter;
-import co.elastic.clients.elasticsearch.esql.EsqlFormat;
 import co.elastic.clients.json.JsonpMappingException;
 import co.elastic.clients.json.jackson.JacksonJsonpMapper;
 import co.elastic.clients.testkit.MockHttpClient;
-import co.elastic.clients.transport.endpoints.BinaryResponse;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
 import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Types;
+import java.util.List;
 
 public class EsqlAdapterTest extends Assertions {
 
@@ -120,11 +120,6 @@ public void testProfilingInfo() throws IOException {
     @Test
     public void testObjectDeserializer() throws IOException {
 
-        BinaryResponse response = esClient.esql().query(q -> q
-            .query("FROM foo")
-            .format(EsqlFormat.Json)
-        );
-
         Iterable<Data> data = esClient.esql().query(
             new ObjectsEsqlAdapter<>(Data.class),
             "FROM employees | STATS avg_salary = AVG(salary) by country"
@@ -151,4 +146,54 @@ public void testResultSetAdapter() throws Exception {
             System.out.println(resultSet.getDouble("avg_salary") + " " + resultSet.getString(2));
         }
     }
+
+    @Test
+    public void testDenseVector() throws IOException, SQLException {
+
+        String jsonEmbeddings = "{\n" +
+                                "  \"took\": 8,\n" +
+                                "  \"is_partial\": false,\n" +
+                                "  \"completion_time_in_millis\": 1776348683193,\n" +
+                                "  \"documents_found\": 1,\n" +
+                                "  \"values_loaded\": 252725,\n" +
+                                "  \"start_time_in_millis\": 1776348683185,\n" +
+                                "  \"expiration_time_in_millis\": 1776780683134,\n" +
+                                "  \"columns\": [\n" +
+                                "    {\n" +
+                                "      \"name\": \"content\",\n" +
+                                "      \"type\": \"text\"\n" +
+                                "    },\n" +
+                                "    {\n" +
+                                "      \"name\": \"content.keyword\",\n" +
+                                "      \"type\": \"keyword\"\n" +
+                                "    },\n" +
+                                "    {\n" +
+                                "      \"name\": \"embedding\",\n" +
+                                "      \"type\": \"dense_vector\"\n" +
+                                "    }\n" +
+                                "    ],\n" +
+                                "  \"values\": [\n" +
+                                "[\"some text\",\"id\",[0.1,0.2,0.3]]\n" +
+                                "]\n" +
+                                "}";
+
+        ElasticsearchClient esClient = new MockHttpClient()
+            .add("/_query", "application/json", jsonEmbeddings)
+            .client(new JacksonJsonpMapper());
+
+        ResultSet resultSet = esClient.esql().query(
+                ResultSetEsqlAdapter.INSTANCE,
+                "FROM embeddings");
+
+        assertEquals(3, resultSet.getMetaData().getColumnCount());
+        assertEquals(Types.VARCHAR, resultSet.getMetaData().getColumnType(1));
+        assertEquals(Types.VARCHAR, resultSet.getMetaData().getColumnType(2));
+        assertEquals(Types.ARRAY, resultSet.getMetaData().getColumnType(3));
+
+        while (resultSet.next()) {
+            List vec = resultSet.getObject("embedding", List.class);
+            assertEquals(3,vec.size());
+            System.out.println(resultSet.getString("embedding"));
+        }
+    }
 }