CASSANDRA-19837: Support ORDER BY ANN in query builder

query-builder/revapi.json

@@ -2772,6 +2772,11 @@
         "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::orderBy(com.datastax.oss.driver.api.core.CqlIdentifier, com.datastax.oss.driver.api.querybuilder.term.Term, com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder)",
+        "justification": "CASSANDRA-19837: New method added to support ORDER BY ANN clause in query builder"
       }
     ]
   }

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

@@ -23,6 +23,7 @@
 import com.datastax.oss.driver.api.querybuilder.BuildableQuery;
 import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
 import com.datastax.oss.driver.api.querybuilder.relation.OngoingWhereClause;
+import com.datastax.oss.driver.api.querybuilder.term.Term;
 import com.datastax.oss.driver.internal.core.CqlIdentifiers;
 import com.datastax.oss.driver.shaded.guava.common.collect.Iterables;
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -137,6 +138,26 @@ default Select orderBy(@NonNull Map<String, ClusteringOrder> orderings) {
   @NonNull
   Select orderBy(@NonNull CqlIdentifier columnId, @NonNull ClusteringOrder order);
 
+  /**
+   * Adds the provided ORDER BY ANN clause to the query.
+   *
+   * <p>If an ordering was already defined for this identifier, it will be removed and the new
+   * clause will be appended at the end of the current list for this query.
+   */
+  @NonNull
+  Select orderBy(
+      @NonNull CqlIdentifier columnId, @NonNull Term vector, @NonNull ClusteringOrder order);
+
+  /**
+   * Shortcut for {@link #orderBy(CqlIdentifier, Term, ClusteringOrder)
+   * orderBy(CqlIdentifier.fromCql(columnName), order)}.
+   */
+  @NonNull
+  default Select orderBy(
+      @NonNull String columnName, @NonNull Term vector, @NonNull ClusteringOrder order) {
+    return orderBy(CqlIdentifier.fromCql(columnName), vector, order);
+  }
+
   /**
    * Shortcut for {@link #orderBy(CqlIdentifier, ClusteringOrder)
    * orderBy(CqlIdentifier.fromCql(columnName), order)}.

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

@@ -26,6 +26,7 @@
 import com.datastax.oss.driver.api.querybuilder.select.Select;
 import com.datastax.oss.driver.api.querybuilder.select.SelectFrom;
 import com.datastax.oss.driver.api.querybuilder.select.Selector;
+import com.datastax.oss.driver.api.querybuilder.term.Term;
 import com.datastax.oss.driver.internal.querybuilder.CqlHelper;
 import com.datastax.oss.driver.internal.querybuilder.ImmutableCollections;
 import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
@@ -48,7 +49,7 @@ public class DefaultSelect implements SelectFrom, Select {
   private final ImmutableList<Selector> selectors;
   private final ImmutableList<Relation> relations;
   private final ImmutableList<Selector> groupByClauses;
-  private final ImmutableMap<CqlIdentifier, ClusteringOrder> orderings;
+  private final ImmutableMap<CqlIdentifier, OrderDirection> orderings;
   private final Object limit;
   private final Object perPartitionLimit;
   private final boolean allowsFiltering;
@@ -83,7 +84,7 @@ public DefaultSelect(
       @NonNull ImmutableList<Selector> selectors,
       @NonNull ImmutableList<Relation> relations,
       @NonNull ImmutableList<Selector> groupByClauses,
-      @NonNull ImmutableMap<CqlIdentifier, ClusteringOrder> orderings,
+      @NonNull ImmutableMap<CqlIdentifier, OrderDirection> orderings,
       @Nullable Object limit,
       @Nullable Object perPartitionLimit,
       boolean allowsFiltering) {
@@ -257,17 +258,28 @@ public Select withGroupByClauses(@NonNull ImmutableList<Selector> newGroupByClau
   @NonNull
   @Override
   public Select orderBy(@NonNull CqlIdentifier columnId, @NonNull ClusteringOrder order) {
-    return withOrderings(ImmutableCollections.append(orderings, columnId, order));
+    return withOrderings(
+        ImmutableCollections.append(orderings, columnId, new OrderDirection(order)));
+  }
+
+  @NonNull
+  @Override
+  public Select orderBy(
+      @NonNull CqlIdentifier columnId, @NonNull Term vector, @NonNull ClusteringOrder order) {
+    return withOrderings(
+        ImmutableCollections.append(orderings, columnId, new OrderDirection(order, vector)));
   }
 
   @NonNull
   @Override
   public Select orderByIds(@NonNull Map<CqlIdentifier, ClusteringOrder> newOrderings) {
-    return withOrderings(ImmutableCollections.concat(orderings, newOrderings));
+    ImmutableMap.Builder<CqlIdentifier, OrderDirection> builder = ImmutableMap.builder();
+    newOrderings.forEach((key, order) -> builder.put(key, new OrderDirection(order)));
+    return withOrderings(ImmutableCollections.concat(orderings, builder.build()));
   }
 
   @NonNull
-  public Select withOrderings(@NonNull ImmutableMap<CqlIdentifier, ClusteringOrder> newOrderings) {
+  public Select withOrderings(@NonNull ImmutableMap<CqlIdentifier, OrderDirection> newOrderings) {
     return new DefaultSelect(
         keyspace,
         table,
@@ -392,14 +404,15 @@ public String asCql() {
     CqlHelper.append(groupByClauses, builder, " GROUP BY ", ",", null);
 
     boolean first = true;
-    for (Map.Entry<CqlIdentifier, ClusteringOrder> entry : orderings.entrySet()) {
+    for (Map.Entry<CqlIdentifier, OrderDirection> 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));
+      entry.getValue().appendTo(builder);
     }
 
     if (limit != null) {
@@ -490,7 +503,7 @@ public ImmutableList<Selector> getGroupByClauses() {
   }
 
   @NonNull
-  public ImmutableMap<CqlIdentifier, ClusteringOrder> getOrderings() {
+  public ImmutableMap<CqlIdentifier, OrderDirection> getOrderings() {
     return orderings;
   }
 

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

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.datastax.oss.driver.internal.querybuilder.select;
+
+import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
+import com.datastax.oss.driver.api.querybuilder.term.Term;
+import edu.umd.cs.findbugs.annotations.NonNull;
+
+public class OrderDirection {
+  private final ClusteringOrder clusteringOrder;
+  private final Term clusteringVector;
+
+  public OrderDirection(ClusteringOrder clusteringOrder) {
+    this(clusteringOrder, null);
+  }
+
+  public OrderDirection(ClusteringOrder clusteringOrder, Term clusteringVector) {
+    this.clusteringOrder = clusteringOrder;
+    this.clusteringVector = clusteringVector;
+  }
+
+  public void appendTo(@NonNull StringBuilder builder) {
+    if (clusteringVector != null) {
+      builder.append(" ANN OF ");
+      clusteringVector.appendTo(builder);
+    }
+    if (clusteringOrder != null) {
+      builder.append(" ").append(clusteringOrder.name());
+    }
+  }
+}

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/select/SelectOrderingTest.java

@@ -23,6 +23,11 @@
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.literal;
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.selectFrom;
 
+import com.datastax.oss.driver.api.core.data.CqlVector;
+import com.datastax.oss.driver.api.core.type.DataTypes;
+import com.datastax.oss.driver.api.core.type.VectorType;
+import com.datastax.oss.driver.api.core.type.codec.TypeCodec;
+import com.datastax.oss.driver.api.core.type.codec.TypeCodecs;
 import com.datastax.oss.driver.api.querybuilder.relation.Relation;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import org.junit.Test;
@@ -46,6 +51,18 @@ public void should_generate_ordering_clauses() {
         .hasCql("SELECT * FROM foo WHERE k=1 ORDER BY c1 ASC,c2 DESC");
   }
 
+  @Test
+  public void should_generate_vector_ordering_clauses() {
+    VectorType vectorType = DataTypes.vectorOf(DataTypes.FLOAT, 3);
+    TypeCodec<CqlVector<Float>> codec = TypeCodecs.vectorOf(vectorType, TypeCodecs.FLOAT);
+    assertThat(
+            selectFrom("foo")
+                .all()
+                .where(Relation.column("k").isEqualTo(literal(1)))
+                .orderBy("c1", literal(CqlVector.newInstance(1.01f, 2.5f, -3.0f), codec), ASC))
+        .hasCql("SELECT * FROM foo WHERE k=1 ORDER BY c1 ANN OF [1.01, 2.5, -3.0] ASC");
+  }
+
   @Test(expected = IllegalArgumentException.class)
   public void should_fail_when_provided_names_resolve_to_the_same_id() {
     selectFrom("foo")