From fa402d4b607ca95bde62b636512c7366f2a0fef5 Mon Sep 17 00:00:00 2001 From: Lukasz Antoniak Date: Mon, 19 Aug 2024 08:56:42 +0200 Subject: [PATCH] CASSANDRA-19837: Support ORDER BY ANN in query builder --- query-builder/revapi.json | 5 ++ .../api/querybuilder/select/Select.java | 21 +++++++++ .../querybuilder/select/DefaultSelect.java | 29 ++++++++---- .../querybuilder/select/OrderDirection.java | 46 +++++++++++++++++++ .../select/SelectOrderingTest.java | 17 +++++++ 5 files changed, 110 insertions(+), 8 deletions(-) create mode 100644 query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/OrderDirection.java diff --git a/query-builder/revapi.json b/query-builder/revapi.json index 9d0163b487e..3dc6f834ed7 100644 --- a/query-builder/revapi.json +++ b/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" } ] } diff --git a/query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/select/Select.java b/query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/select/Select.java index a22b45c35bd..31a92c46312 100644 --- a/query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/select/Select.java +++ b/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 orderings) { @NonNull Select orderBy(@NonNull CqlIdentifier columnId, @NonNull ClusteringOrder order); + /** + * Adds the provided ORDER BY ANN clause to the query. + * + *

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)}. diff --git a/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/DefaultSelect.java b/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/DefaultSelect.java index 86a2a07a3f2..7f1defdadca 100644 --- a/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/DefaultSelect.java +++ b/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 selectors; private final ImmutableList relations; private final ImmutableList groupByClauses; - private final ImmutableMap orderings; + private final ImmutableMap orderings; private final Object limit; private final Object perPartitionLimit; private final boolean allowsFiltering; @@ -83,7 +84,7 @@ public DefaultSelect( @NonNull ImmutableList selectors, @NonNull ImmutableList relations, @NonNull ImmutableList groupByClauses, - @NonNull ImmutableMap orderings, + @NonNull ImmutableMap orderings, @Nullable Object limit, @Nullable Object perPartitionLimit, boolean allowsFiltering) { @@ -257,17 +258,28 @@ public Select withGroupByClauses(@NonNull ImmutableList 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 newOrderings) { - return withOrderings(ImmutableCollections.concat(orderings, newOrderings)); + ImmutableMap.Builder 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 newOrderings) { + public Select withOrderings(@NonNull ImmutableMap 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 entry : orderings.entrySet()) { + for (Map.Entry 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 getGroupByClauses() { } @NonNull - public ImmutableMap getOrderings() { + public ImmutableMap getOrderings() { return orderings; } diff --git a/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/OrderDirection.java b/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/select/OrderDirection.java new file mode 100644 index 00000000000..f4cefdb9edc --- /dev/null +++ b/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()); + } + } +} diff --git a/query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/select/SelectOrderingTest.java b/query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/select/SelectOrderingTest.java index ff27fde4f8f..9ae2cc85bfb 100644 --- a/query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/select/SelectOrderingTest.java +++ b/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> 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")