Skip to content

Commit 29ad892

Browse files
SiyaoIsHidingSiyaoIsHiding
committed
orderBy(annOf("c1", CqlVector.newInstance(0.1, 0.2, 0.3))))
1 parent a54ede2 commit 29ad892

File tree

6 files changed

+154
-8
lines changed

6 files changed

+154
-8
lines changed

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import com.datastax.oss.driver.api.core.CqlIdentifier;
2121
import com.datastax.oss.driver.api.core.context.DriverContext;
22+
import com.datastax.oss.driver.api.core.data.CqlVector;
2223
import com.datastax.oss.driver.api.core.metadata.token.Token;
2324
import com.datastax.oss.driver.api.core.type.DataType;
2425
import com.datastax.oss.driver.api.core.type.DataTypes;
@@ -29,6 +30,7 @@
2930
import com.datastax.oss.driver.api.querybuilder.delete.DeleteSelection;
3031
import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
3132
import com.datastax.oss.driver.api.querybuilder.relation.Relation;
33+
import com.datastax.oss.driver.api.querybuilder.select.Ann;
3234
import com.datastax.oss.driver.api.querybuilder.select.SelectFrom;
3335
import com.datastax.oss.driver.api.querybuilder.select.Selector;
3436
import com.datastax.oss.driver.api.querybuilder.term.Term;
@@ -43,6 +45,7 @@
4345
import com.datastax.oss.driver.internal.querybuilder.DefaultRaw;
4446
import com.datastax.oss.driver.internal.querybuilder.delete.DefaultDelete;
4547
import com.datastax.oss.driver.internal.querybuilder.insert.DefaultInsert;
48+
import com.datastax.oss.driver.internal.querybuilder.select.DefaultAnn;
4649
import com.datastax.oss.driver.internal.querybuilder.select.DefaultBindMarker;
4750
import com.datastax.oss.driver.internal.querybuilder.select.DefaultSelect;
4851
import com.datastax.oss.driver.internal.querybuilder.term.BinaryArithmeticTerm;
@@ -538,4 +541,12 @@ public static Truncate truncate(@Nullable String keyspace, @NonNull String table
538541
return truncate(
539542
keyspace == null ? null : CqlIdentifier.fromCql(keyspace), CqlIdentifier.fromCql(table));
540543
}
544+
545+
public static Ann annOf(@NonNull CqlIdentifier cqlIdentifier, @NonNull CqlVector<Number> vector) {
546+
return new DefaultAnn(cqlIdentifier, vector);
547+
}
548+
549+
public static Ann annOf(@NonNull String cqlIdentifier, @NonNull CqlVector<Number> vector) {
550+
return new DefaultAnn(CqlIdentifier.fromCql(cqlIdentifier), vector);
551+
}
541552
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
package com.datastax.oss.driver.api.querybuilder.select;
2+
3+
import com.datastax.oss.driver.api.core.data.CqlVector;
4+
import com.datastax.oss.driver.api.querybuilder.BuildableQuery;
5+
6+
public interface Ann extends BuildableQuery {
7+
8+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,8 @@ default Select orderBy(@NonNull String columnName, @NonNull ClusteringOrder orde
146146
return orderBy(CqlIdentifier.fromCql(columnName), order);
147147
}
148148

149+
@NonNull
150+
Select orderBy(@NonNull Ann ann);
149151
/**
150152
* Adds a LIMIT clause to this query with a literal value.
151153
*

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package com.datastax.oss.driver.internal.querybuilder.select;
2+
3+
import com.datastax.oss.driver.api.core.CqlIdentifier;
4+
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
5+
import com.datastax.oss.driver.api.core.cql.SimpleStatementBuilder;
6+
import com.datastax.oss.driver.api.core.data.CqlVector;
7+
import com.datastax.oss.driver.api.querybuilder.select.Ann;
8+
import edu.umd.cs.findbugs.annotations.NonNull;
9+
10+
import java.util.Map;
11+
12+
import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.literal;
13+
14+
public class DefaultAnn implements Ann {
15+
private final CqlIdentifier cqlIdentifier;
16+
private final CqlVector<Number> vector;
17+
18+
public DefaultAnn(@NonNull CqlIdentifier cqlIdentifier, @NonNull CqlVector<Number> vector) {
19+
this.cqlIdentifier = cqlIdentifier;
20+
this.vector = vector;
21+
}
22+
23+
@NonNull
24+
@Override
25+
public String asCql() {
26+
StringBuilder builder = new StringBuilder();
27+
builder.append("ORDER BY ");
28+
builder.append(this.cqlIdentifier.asCql(true));
29+
builder.append(" ANN OF ");
30+
literal(vector).appendTo(builder);
31+
return builder.toString();
32+
}
33+
34+
@NonNull
35+
@Override
36+
public SimpleStatement build() {
37+
return Ann.super.build();
38+
}
39+
40+
@NonNull
41+
@Override
42+
public SimpleStatement build(@NonNull Object... values) {
43+
return Ann.super.build(values);
44+
}
45+
46+
@NonNull
47+
@Override
48+
public SimpleStatement build(@NonNull Map<String, Object> namedValues) {
49+
return Ann.super.build(namedValues);
50+
}
51+
52+
@NonNull
53+
@Override
54+
public SimpleStatementBuilder builder() {
55+
return Ann.super.builder();
56+
}
57+
58+
}

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

Lines changed: 54 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
2424
import com.datastax.oss.driver.api.querybuilder.BindMarker;
2525
import com.datastax.oss.driver.api.querybuilder.relation.Relation;
26+
import com.datastax.oss.driver.api.querybuilder.select.Ann;
2627
import com.datastax.oss.driver.api.querybuilder.select.Select;
2728
import com.datastax.oss.driver.api.querybuilder.select.SelectFrom;
2829
import com.datastax.oss.driver.api.querybuilder.select.Selector;
@@ -49,6 +50,7 @@ public class DefaultSelect implements SelectFrom, Select {
4950
private final ImmutableList<Relation> relations;
5051
private final ImmutableList<Selector> groupByClauses;
5152
private final ImmutableMap<CqlIdentifier, ClusteringOrder> orderings;
53+
private final Ann ann;
5254
private final Object limit;
5355
private final Object perPartitionLimit;
5456
private final boolean allowsFiltering;
@@ -65,6 +67,7 @@ public DefaultSelect(@Nullable CqlIdentifier keyspace, @NonNull CqlIdentifier ta
6567
ImmutableMap.of(),
6668
null,
6769
null,
70+
null,
6871
false);
6972
}
7073

@@ -74,6 +77,7 @@ public DefaultSelect(@Nullable CqlIdentifier keyspace, @NonNull CqlIdentifier ta
7477
* @param selectors if it contains {@link AllSelector#INSTANCE}, that must be the only element.
7578
* This isn't re-checked because methods that call this constructor internally already do it,
7679
* make sure you do it yourself.
80+
* @param ann Approximate nearest neighbor. ANN ordering does not support secondary ordering or ASC order.
7781
*/
7882
public DefaultSelect(
7983
@Nullable CqlIdentifier keyspace,
@@ -84,6 +88,7 @@ public DefaultSelect(
8488
@NonNull ImmutableList<Relation> relations,
8589
@NonNull ImmutableList<Selector> groupByClauses,
8690
@NonNull ImmutableMap<CqlIdentifier, ClusteringOrder> orderings,
91+
@Nullable Ann ann,
8792
@Nullable Object limit,
8893
@Nullable Object perPartitionLimit,
8994
boolean allowsFiltering) {
@@ -94,6 +99,10 @@ public DefaultSelect(
9499
|| (limit instanceof Integer && (Integer) limit > 0)
95100
|| limit instanceof BindMarker,
96101
"limit must be a strictly positive integer or a bind marker");
102+
Preconditions.checkArgument(
103+
orderings.isEmpty() || ann == null,
104+
"ANN ordering does not support secondary ordering");
105+
this.ann = ann;
97106
this.keyspace = keyspace;
98107
this.table = table;
99108
this.isJson = isJson;
@@ -117,6 +126,7 @@ public SelectFrom json() {
117126
relations,
118127
groupByClauses,
119128
orderings,
129+
null,
120130
limit,
121131
perPartitionLimit,
122132
allowsFiltering);
@@ -134,6 +144,7 @@ public SelectFrom distinct() {
134144
relations,
135145
groupByClauses,
136146
orderings,
147+
null,
137148
limit,
138149
perPartitionLimit,
139150
allowsFiltering);
@@ -193,6 +204,7 @@ public Select withSelectors(@NonNull ImmutableList<Selector> newSelectors) {
193204
relations,
194205
groupByClauses,
195206
orderings,
207+
null,
196208
limit,
197209
perPartitionLimit,
198210
allowsFiltering);
@@ -221,6 +233,7 @@ public Select withRelations(@NonNull ImmutableList<Relation> newRelations) {
221233
newRelations,
222234
groupByClauses,
223235
orderings,
236+
null,
224237
limit,
225238
perPartitionLimit,
226239
allowsFiltering);
@@ -249,6 +262,7 @@ public Select withGroupByClauses(@NonNull ImmutableList<Selector> newGroupByClau
249262
relations,
250263
newGroupByClauses,
251264
orderings,
265+
null,
252266
limit,
253267
perPartitionLimit,
254268
allowsFiltering);
@@ -266,6 +280,12 @@ public Select orderByIds(@NonNull Map<CqlIdentifier, ClusteringOrder> newOrderin
266280
return withOrderings(ImmutableCollections.concat(orderings, newOrderings));
267281
}
268282

283+
@NonNull
284+
@Override
285+
public Select orderBy(@NonNull Ann ann){
286+
return withAnn(ann);
287+
}
288+
269289
@NonNull
270290
public Select withOrderings(@NonNull ImmutableMap<CqlIdentifier, ClusteringOrder> newOrderings) {
271291
return new DefaultSelect(
@@ -277,6 +297,23 @@ public Select withOrderings(@NonNull ImmutableMap<CqlIdentifier, ClusteringOrder
277297
relations,
278298
groupByClauses,
279299
newOrderings,
300+
null,
301+
limit,
302+
perPartitionLimit,
303+
allowsFiltering);
304+
}
305+
306+
@NonNull Select withAnn(@NonNull Ann ann){
307+
return new DefaultSelect(
308+
keyspace,
309+
table,
310+
isJson,
311+
isDistinct,
312+
selectors,
313+
relations,
314+
groupByClauses,
315+
orderings,
316+
ann,
280317
limit,
281318
perPartitionLimit,
282319
allowsFiltering);
@@ -295,6 +332,7 @@ public Select limit(int limit) {
295332
relations,
296333
groupByClauses,
297334
orderings,
335+
null,
298336
limit,
299337
perPartitionLimit,
300338
allowsFiltering);
@@ -312,6 +350,7 @@ public Select limit(@Nullable BindMarker bindMarker) {
312350
relations,
313351
groupByClauses,
314352
orderings,
353+
null,
315354
bindMarker,
316355
perPartitionLimit,
317356
allowsFiltering);
@@ -331,6 +370,7 @@ public Select perPartitionLimit(int perPartitionLimit) {
331370
relations,
332371
groupByClauses,
333372
orderings,
373+
null,
334374
limit,
335375
perPartitionLimit,
336376
allowsFiltering);
@@ -348,6 +388,7 @@ public Select perPartitionLimit(@Nullable BindMarker bindMarker) {
348388
relations,
349389
groupByClauses,
350390
orderings,
391+
null,
351392
limit,
352393
bindMarker,
353394
allowsFiltering);
@@ -365,6 +406,7 @@ public Select allowFiltering() {
365406
relations,
366407
groupByClauses,
367408
orderings,
409+
null,
368410
limit,
369411
perPartitionLimit,
370412
true);
@@ -391,15 +433,19 @@ public String asCql() {
391433
CqlHelper.append(relations, builder, " WHERE ", " AND ", null);
392434
CqlHelper.append(groupByClauses, builder, " GROUP BY ", ",", null);
393435

394-
boolean first = true;
395-
for (Map.Entry<CqlIdentifier, ClusteringOrder> entry : orderings.entrySet()) {
396-
if (first) {
397-
builder.append(" ORDER BY ");
398-
first = false;
399-
} else {
400-
builder.append(",");
436+
if (ann != null){
437+
builder.append(" ").append(this.ann.asCql());
438+
}else{
439+
boolean first = true;
440+
for (Map.Entry<CqlIdentifier, ClusteringOrder> entry : orderings.entrySet()) {
441+
if (first) {
442+
builder.append(" ORDER BY ");
443+
first = false;
444+
} else {
445+
builder.append(",");
446+
}
447+
builder.append(entry.getKey().asCql(true)).append(" ").append(entry.getValue().name());
401448
}
402-
builder.append(entry.getKey().asCql(true)).append(" ").append(entry.getValue().name());
403449
}
404450

405451
if (limit != null) {

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,11 @@
2020
import static com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder.ASC;
2121
import static com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder.DESC;
2222
import static com.datastax.oss.driver.api.querybuilder.Assertions.assertThat;
23+
import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.annOf;
2324
import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.literal;
2425
import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.selectFrom;
2526

27+
import com.datastax.oss.driver.api.core.data.CqlVector;
2628
import com.datastax.oss.driver.api.querybuilder.relation.Relation;
2729
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
2830
import org.junit.Test;
@@ -74,4 +76,23 @@ public void should_replace_previous_ordering() {
7476
.orderBy(ImmutableMap.of("c1", DESC, "c2", ASC)))
7577
.hasCql("SELECT * FROM foo WHERE k=1 ORDER BY c3 ASC,c1 DESC,c2 ASC");
7678
}
79+
80+
@Test
81+
public void should_generate_ann_clause() {
82+
assertThat(
83+
selectFrom("foo")
84+
.all()
85+
.where(Relation.column("k").isEqualTo(literal(1)))
86+
.orderBy(annOf("c1", CqlVector.newInstance(0.1, 0.2, 0.3))))
87+
.hasCql("SELECT * FROM foo WHERE k=1 ORDER BY c1 ANN OF [0.1, 0.2, 0.3]");
88+
}
89+
90+
@Test(expected = IllegalArgumentException.class)
91+
public void should_fail_when_provided_ann_with_other_orderings(){
92+
selectFrom("foo")
93+
.all()
94+
.where(Relation.column("k").isEqualTo(literal(1)))
95+
.orderBy("c1", ASC)
96+
.orderBy(annOf("c1", CqlVector.newInstance(0.1, 0.2, 0.3)));
97+
}
7798
}

0 commit comments

Comments
 (0)