refactor and pass tests

Author: SiyaoIsHiding

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

@@ -19,7 +19,6 @@
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.context.DriverContext;
-import com.datastax.oss.driver.api.core.data.CqlVector;
 import com.datastax.oss.driver.api.core.metadata.token.Token;
 import com.datastax.oss.driver.api.core.type.DataType;
 import com.datastax.oss.driver.api.core.type.DataTypes;
@@ -30,7 +29,6 @@
 import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection;
 import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
 import com.datastax.oss.driver.api.querybuilder.relation.Relation;
-import com.datastax.oss.driver.api.querybuilder.select.Ann;
 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;
@@ -45,7 +43,6 @@
 import com.datastax.oss.driver.internal.querybuilder.DefaultRaw;
 import com.datastax.oss.driver.internal.querybuilder.delete.DefaultDelete;
 import com.datastax.oss.driver.internal.querybuilder.insert.DefaultInsert;
-import com.datastax.oss.driver.internal.querybuilder.select.DefaultAnn;
 import com.datastax.oss.driver.internal.querybuilder.select.DefaultBindMarker;
 import com.datastax.oss.driver.internal.querybuilder.select.DefaultSelect;
 import com.datastax.oss.driver.internal.querybuilder.term.BinaryArithmeticTerm;
@@ -541,12 +538,4 @@ public static Truncate truncate(@Nullable String keyspace, @NonNull String table
     return truncate(
         keyspace == null ? null : CqlIdentifier.fromCql(keyspace), CqlIdentifier.fromCql(table));
   }
-
-  public static Ann annOf(@NonNull CqlIdentifier cqlIdentifier, @NonNull CqlVector<Number> vector) {
-    return new DefaultAnn(cqlIdentifier, vector);
-  }
-
-  public static Ann annOf(@NonNull String cqlIdentifier, @NonNull CqlVector<Number> vector) {
-    return new DefaultAnn(CqlIdentifier.fromCql(cqlIdentifier), vector);
-  }
 }

query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/select/Ann.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,8 +147,16 @@ default Select orderBy(@NonNull String columnName, @NonNull ClusteringOrder orde
     return orderBy(CqlIdentifier.fromCql(columnName), order);
   }
 
+  /**
+   * Shortcut for {@link #orderBy(CqlIdentifier, CqlVector)}, adding an ORDER BY ... ANN OF ...
+   * clause
+   */
+  @NonNull
+  Select orderBy(@NonNull String columnName, @NonNull CqlVector<? extends Number> ann);
+
+  /** Adds the ORDER BY ... ANN OF ... clause */
   @NonNull
-  Select orderBy(@NonNull Ann ann);
+  Select orderBy(@NonNull CqlIdentifier columnId, @NonNull CqlVector<? extends Number> ann);
   /**
    * Adds a LIMIT clause to this query with a literal value.
    *

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

@@ -20,10 +20,11 @@
 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.Ann;
 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;
@@ -276,14 +277,20 @@ public Select orderBy(@NonNull CqlIdentifier columnId, @NonNull ClusteringOrder
 
   @NonNull
   @Override
-  public Select orderByIds(@NonNull Map<CqlIdentifier, ClusteringOrder> newOrderings) {
-    return withOrderings(ImmutableCollections.concat(orderings, newOrderings));
+  public Select orderBy(@NonNull String columnName, @NonNull CqlVector<? extends Number> ann) {
+    return withAnn(new Ann(CqlIdentifier.fromCql(columnName), ann));
+  }
+
+  @NonNull
+  @Override
+  public Select orderBy(@NonNull CqlIdentifier columnId, @NonNull CqlVector<? extends Number> ann) {
+    return withAnn(new Ann(columnId, ann));
   }
 
   @NonNull
   @Override
-  public Select orderBy(@NonNull Ann ann) {
-    return withAnn(ann);
+  public Select orderByIds(@NonNull Map<CqlIdentifier, ClusteringOrder> newOrderings) {
+    return withOrderings(ImmutableCollections.concat(orderings, newOrderings));
   }
 
   @NonNull
@@ -435,7 +442,8 @@ public String asCql() {
     CqlHelper.append(groupByClauses, builder, " GROUP BY ", ",", null);
 
     if (ann != null) {
-      builder.append(" ").append(this.ann.asCql());
+      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()) {
@@ -559,4 +567,14 @@ public boolean allowsFiltering() {
   public String toString() {
     return asCql();
   }
+
+  public static class Ann {
+    private final CqlVector<? extends Number> vector;
+    private final CqlIdentifier columnId;
+
+    private Ann(CqlIdentifier columnId, CqlVector<? extends Number> vector) {
+      this.vector = vector;
+      this.columnId = columnId;
+    }
+  }
 }

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

@@ -20,7 +20,6 @@
 import static com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder.ASC;
 import static com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder.DESC;
 import static com.datastax.oss.driver.api.querybuilder.Assertions.assertThat;
-import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.annOf;
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.literal;
 import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.selectFrom;
 
@@ -83,7 +82,7 @@ public void should_generate_ann_clause() {
             selectFrom("foo")
                 .all()
                 .where(Relation.column("k").isEqualTo(literal(1)))
-                .orderBy(annOf("c1", CqlVector.newInstance(0.1, 0.2, 0.3))))
+                .orderBy("c1", CqlVector.newInstance(0.1, 0.2, 0.3)))
         .hasCql("SELECT * FROM foo WHERE k=1 ORDER BY c1 ANN OF [0.1, 0.2, 0.3]");
   }
 
@@ -93,6 +92,6 @@ public void should_fail_when_provided_ann_with_other_orderings() {
         .all()
         .where(Relation.column("k").isEqualTo(literal(1)))
         .orderBy("c1", ASC)
-        .orderBy(annOf("c1", CqlVector.newInstance(0.1, 0.2, 0.3)));
+        .orderBy("c2", CqlVector.newInstance(0.1, 0.2, 0.3));
   }
 }