JAVA-3118: Add support for vector data type in Schema Builder, QueryBuilder

patch by Jane He; reviewed by Mick Semb Wever and Bret McGuire for JAVA-3118
reference: #1931

Author: SiyaoIsHiding

manual/query_builder/select/README.md

@@ -387,6 +387,29 @@ selectFrom("sensor_data")
 // SELECT reading FROM sensor_data WHERE id=? ORDER BY date DESC
 ```
 
+Vector Search:
+
+```java
+
+import com.datastax.oss.driver.api.core.data.CqlVector;
+
+selectFrom("foo")
+    .all()
+    .where(Relation.column("k").isEqualTo(literal(1)))
+    .orderByAnnOf("c1", CqlVector.newInstance(0.1, 0.2, 0.3));
+// SELECT * FROM foo WHERE k=1 ORDER BY c1 ANN OF [0.1, 0.2, 0.3]
+
+selectFrom("cycling", "comments_vs")
+    .column("comment")
+    .function(
+        "similarity_cosine",
+        Selector.column("comment_vector"),
+        literal(CqlVector.newInstance(0.2, 0.15, 0.3, 0.2, 0.05)))
+    .orderByAnnOf("comment_vector", CqlVector.newInstance(0.1, 0.15, 0.3, 0.12, 0.05))
+    .limit(1);
+// SELECT comment,similarity_cosine(comment_vector,[0.2, 0.15, 0.3, 0.2, 0.05]) FROM cycling.comments_vs ORDER BY comment_vector ANN OF [0.1, 0.15, 0.3, 0.12, 0.05] LIMIT 1
+```
+
 Limits:
 
 ```java

query-builder/revapi.json

@@ -2772,6 +2772,16 @@
         "code": "java.method.addedToInterface",
         "new": "method com.datastax.oss.driver.api.querybuilder.update.UpdateStart com.datastax.oss.driver.api.querybuilder.update.UpdateStart::usingTtl(int)",
         "justification": "JAVA-2210: Add ability to set TTL for modification queries"
+      },
+      {
+        "code": "java.method.addedToInterface",
+        "new": "method com.datastax.oss.driver.api.querybuilder.select.Select com.datastax.oss.driver.api.querybuilder.select.Select::orderByAnnOf(java.lang.String, com.datastax.oss.driver.api.core.data.CqlVector<?>)",
+        "justification": "JAVA-3118: Add support for vector data type in Schema Builder, QueryBuilder"
+      },
+      {
+        "code": "java.method.addedToInterface",
+        "new": "method com.datastax.oss.driver.api.querybuilder.select.Select com.datastax.oss.driver.api.querybuilder.select.Select::orderByAnnOf(com.datastax.oss.driver.api.core.CqlIdentifier, com.datastax.oss.driver.api.core.data.CqlVector<?>)",
+        "justification": "JAVA-3118: Add support for vector data type in Schema Builder, QueryBuilder"
       }
     ]
   }

query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/select/Select.java

@@ -18,6 +18,7 @@
 package com.datastax.oss.driver.api.querybuilder.select;
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
+import com.datastax.oss.driver.api.core.data.CqlVector;
 import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
 import com.datastax.oss.driver.api.querybuilder.BindMarker;
 import com.datastax.oss.driver.api.querybuilder.BuildableQuery;
@@ -146,6 +147,16 @@ default Select orderBy(@NonNull String columnName, @NonNull ClusteringOrder orde
     return orderBy(CqlIdentifier.fromCql(columnName), order);
   }
 
+  /**
+   * Shortcut for {@link #orderByAnnOf(CqlIdentifier, CqlVector)}, adding an ORDER BY ... ANN OF ...
+   * clause
+   */
+  @NonNull
+  Select orderByAnnOf(@NonNull String columnName, @NonNull CqlVector<?> ann);
+
+  /** Adds the ORDER BY ... ANN OF ... clause, usually used for vector search */
+  @NonNull
+  Select orderByAnnOf(@NonNull CqlIdentifier columnId, @NonNull CqlVector<?> ann);
   /**
    * Adds a LIMIT clause to this query with a literal value.
    *

query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/DefaultSelect.java

@@ -20,8 +20,10 @@
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
 import com.datastax.oss.driver.api.core.cql.SimpleStatementBuilder;
+import com.datastax.oss.driver.api.core.data.CqlVector;
 import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
 import com.datastax.oss.driver.api.querybuilder.BindMarker;
+import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
 import com.datastax.oss.driver.api.querybuilder.relation.Relation;
 import com.datastax.oss.driver.api.querybuilder.select.Select;
 import com.datastax.oss.driver.api.querybuilder.select.SelectFrom;
@@ -49,6 +51,7 @@ public class DefaultSelect implements SelectFrom, Select {
   private final ImmutableList<Relation> relations;
   private final ImmutableList<Selector> groupByClauses;
   private final ImmutableMap<CqlIdentifier, ClusteringOrder> orderings;
+  private final Ann ann;
   private final Object limit;
   private final Object perPartitionLimit;
   private final boolean allowsFiltering;
@@ -65,6 +68,7 @@ public DefaultSelect(@Nullable CqlIdentifier keyspace, @NonNull CqlIdentifier ta
         ImmutableMap.of(),
         null,
         null,
+        null,
         false);
   }
 
@@ -74,6 +78,8 @@ public DefaultSelect(@Nullable CqlIdentifier keyspace, @NonNull CqlIdentifier ta
    * @param selectors if it contains {@link AllSelector#INSTANCE}, that must be the only element.
    *     This isn't re-checked because methods that call this constructor internally already do it,
    *     make sure you do it yourself.
+   * @param ann Approximate nearest neighbor. ANN ordering does not support secondary ordering or
+   *     ASC order.
    */
   public DefaultSelect(
       @Nullable CqlIdentifier keyspace,
@@ -84,6 +90,7 @@ public DefaultSelect(
       @NonNull ImmutableList<Relation> relations,
       @NonNull ImmutableList<Selector> groupByClauses,
       @NonNull ImmutableMap<CqlIdentifier, ClusteringOrder> orderings,
+      @Nullable Ann ann,
       @Nullable Object limit,
       @Nullable Object perPartitionLimit,
       boolean allowsFiltering) {
@@ -94,6 +101,9 @@ public DefaultSelect(
             || (limit instanceof Integer && (Integer) limit > 0)
             || limit instanceof BindMarker,
         "limit must be a strictly positive integer or a bind marker");
+    Preconditions.checkArgument(
+        orderings.isEmpty() || ann == null, "ANN ordering does not support secondary ordering");
+    this.ann = ann;
     this.keyspace = keyspace;
     this.table = table;
     this.isJson = isJson;
@@ -117,6 +127,7 @@ public SelectFrom json() {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -134,6 +145,7 @@ public SelectFrom distinct() {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -193,6 +205,7 @@ public Select withSelectors(@NonNull ImmutableList<Selector> newSelectors) {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -221,6 +234,7 @@ public Select withRelations(@NonNull ImmutableList<Relation> newRelations) {
         newRelations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -249,6 +263,7 @@ public Select withGroupByClauses(@NonNull ImmutableList<Selector> newGroupByClau
         relations,
         newGroupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -260,6 +275,18 @@ public Select orderBy(@NonNull CqlIdentifier columnId, @NonNull ClusteringOrder
     return withOrderings(ImmutableCollections.append(orderings, columnId, order));
   }
 
+  @NonNull
+  @Override
+  public Select orderByAnnOf(@NonNull String columnName, @NonNull CqlVector<?> ann) {
+    return withAnn(new Ann(CqlIdentifier.fromCql(columnName), ann));
+  }
+
+  @NonNull
+  @Override
+  public Select orderByAnnOf(@NonNull CqlIdentifier columnId, @NonNull CqlVector<?> ann) {
+    return withAnn(new Ann(columnId, ann));
+  }
+
   @NonNull
   @Override
   public Select orderByIds(@NonNull Map<CqlIdentifier, ClusteringOrder> newOrderings) {
@@ -277,6 +304,24 @@ public Select withOrderings(@NonNull ImmutableMap<CqlIdentifier, ClusteringOrder
         relations,
         groupByClauses,
         newOrderings,
+        ann,
+        limit,
+        perPartitionLimit,
+        allowsFiltering);
+  }
+
+  @NonNull
+  Select withAnn(@NonNull Ann ann) {
+    return new DefaultSelect(
+        keyspace,
+        table,
+        isJson,
+        isDistinct,
+        selectors,
+        relations,
+        groupByClauses,
+        orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -295,6 +340,7 @@ public Select limit(int limit) {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -312,6 +358,7 @@ public Select limit(@Nullable BindMarker bindMarker) {
         relations,
         groupByClauses,
         orderings,
+        ann,
         bindMarker,
         perPartitionLimit,
         allowsFiltering);
@@ -331,6 +378,7 @@ public Select perPartitionLimit(int perPartitionLimit) {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         allowsFiltering);
@@ -348,6 +396,7 @@ public Select perPartitionLimit(@Nullable BindMarker bindMarker) {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         bindMarker,
         allowsFiltering);
@@ -365,6 +414,7 @@ public Select allowFiltering() {
         relations,
         groupByClauses,
         orderings,
+        ann,
         limit,
         perPartitionLimit,
         true);
@@ -391,15 +441,20 @@ public String asCql() {
     CqlHelper.append(relations, builder, " WHERE ", " AND ", null);
     CqlHelper.append(groupByClauses, builder, " GROUP BY ", ",", null);
 
-    boolean first = true;
-    for (Map.Entry<CqlIdentifier, ClusteringOrder> entry : orderings.entrySet()) {
-      if (first) {
-        builder.append(" ORDER BY ");
-        first = false;
-      } else {
-        builder.append(",");
+    if (ann != null) {
+      builder.append(" ORDER BY ").append(this.ann.columnId.asCql(true)).append(" ANN OF ");
+      QueryBuilder.literal(ann.vector).appendTo(builder);
+    } else {
+      boolean first = true;
+      for (Map.Entry<CqlIdentifier, ClusteringOrder> entry : orderings.entrySet()) {
+        if (first) {
+          builder.append(" ORDER BY ");
+          first = false;
+        } else {
+          builder.append(",");
+        }
+        builder.append(entry.getKey().asCql(true)).append(" ").append(entry.getValue().name());
       }
-      builder.append(entry.getKey().asCql(true)).append(" ").append(entry.getValue().name());
     }
 
     if (limit != null) {
@@ -499,6 +554,11 @@ public Object getLimit() {
     return limit;
   }
 
+  @Nullable
+  public Ann getAnn() {
+    return ann;
+  }
+
   @Nullable
   public Object getPerPartitionLimit() {
     return perPartitionLimit;
@@ -512,4 +572,14 @@ public boolean allowsFiltering() {
   public String toString() {
     return asCql();
   }
+
+  public static class Ann {
+    private final CqlVector<?> vector;
+    private final CqlIdentifier columnId;
+
+    private Ann(CqlIdentifier columnId, CqlVector<?> vector) {
+      this.vector = vector;
+      this.columnId = columnId;
+    }
+  }
 }

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/delete/DeleteSelectorTest.java

@@ -22,6 +22,7 @@
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.deleteFrom;
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.literal;
 
+import com.datastax.oss.driver.api.core.data.CqlVector;
 import org.junit.Test;
 
 public class DeleteSelectorTest {
@@ -34,6 +35,16 @@ public void should_generate_column_deletion() {
         .hasCql("DELETE v FROM ks.foo WHERE k=?");
   }
 
+  @Test
+  public void should_generate_vector_deletion() {
+    assertThat(
+            deleteFrom("foo")
+                .column("v")
+                .whereColumn("k")
+                .isEqualTo(literal(CqlVector.newInstance(0.1, 0.2))))
+        .hasCql("DELETE v FROM foo WHERE k=[0.1, 0.2]");
+  }
+
   @Test
   public void should_generate_field_deletion() {
     assertThat(

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/insert/RegularInsertTest.java

@@ -23,6 +23,7 @@
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.literal;
 import static org.assertj.core.api.Assertions.catchThrowable;
 
+import com.datastax.oss.driver.api.core.data.CqlVector;
 import com.datastax.oss.driver.api.querybuilder.term.Term;
 import com.datastax.oss.driver.internal.querybuilder.insert.DefaultInsert;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
@@ -41,6 +42,12 @@ public void should_generate_column_assignments() {
         .hasCql("INSERT INTO foo (a,b) VALUES (?,?)");
   }
 
+  @Test
+  public void should_generate_vector_literals() {
+    assertThat(insertInto("foo").value("a", literal(CqlVector.newInstance(0.1, 0.2, 0.3))))
+        .hasCql("INSERT INTO foo (a) VALUES ([0.1, 0.2, 0.3])");
+  }
+
   @Test
   public void should_keep_last_assignment_if_column_listed_twice() {
     assertThat(

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/AlterTableTest.java

@@ -108,4 +108,10 @@ public void should_generate_alter_table_with_no_compression() {
     assertThat(alterTable("bar").withNoCompression())
         .hasCql("ALTER TABLE bar WITH compression={'sstable_compression':''}");
   }
+
+  @Test
+  public void should_generate_alter_table_with_vector() {
+    assertThat(alterTable("bar").alterColumn("v", DataTypes.vectorOf(DataTypes.FLOAT, 3)))
+        .hasCql("ALTER TABLE bar ALTER v TYPE vector<float, 3>");
+  }
 }

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/AlterTypeTest.java

@@ -53,4 +53,10 @@ public void should_generate_alter_table_with_rename_three_columns() {
     assertThat(alterType("bar").renameField("x", "y").renameField("u", "v").renameField("b", "a"))
         .hasCql("ALTER TYPE bar RENAME x TO y AND u TO v AND b TO a");
   }
+
+  @Test
+  public void should_generate_alter_type_with_vector() {
+    assertThat(alterType("foo", "bar").alterField("vec", DataTypes.vectorOf(DataTypes.FLOAT, 3)))
+        .hasCql("ALTER TYPE foo.bar ALTER vec TYPE vector<float, 3>");
+  }
 }

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableTest.java

@@ -314,4 +314,13 @@ public void should_generate_create_table_time_window_compaction() {
         .hasCql(
             "CREATE TABLE bar (k int PRIMARY KEY,v text) WITH compaction={'class':'TimeWindowCompactionStrategy','compaction_window_size':10,'compaction_window_unit':'DAYS','timestamp_resolution':'MICROSECONDS','unsafe_aggressive_sstable_expiration':false}");
   }
+
+  @Test
+  public void should_generate_vector_column() {
+    assertThat(
+            createTable("foo")
+                .withPartitionKey("k", DataTypes.INT)
+                .withColumn("v", DataTypes.vectorOf(DataTypes.FLOAT, 3)))
+        .hasCql("CREATE TABLE foo (k int PRIMARY KEY,v vector<float, 3>)");
+  }
 }