From cb83b979cbd4eb865e46710a6d2ac606a8c88dfa Mon Sep 17 00:00:00 2001
From: Aaron Morton <aaron.morton@datastax.com>
Date: Thu, 14 Nov 2024 09:52:23 +1300
Subject: [PATCH] fixes #1700 default limit for ANN find now 1k (#1706)

---
 .../jsonapi/service/operation/OperationAttempt.java |  1 -
 .../service/operation/query/OrderByCqlClause.java   | 13 +++++++++++++
 .../operation/tables/TableOrderByANNCqlClause.java  | 11 +++++++++--
 .../service/resolver/ReadCommandResolver.java       |  3 ++-
 .../resolver/sort/TableCqlSortClauseResolver.java   |  5 ++++-
 5 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/OperationAttempt.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/OperationAttempt.java
index 32d954e87..ca60f0761 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/OperationAttempt.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/OperationAttempt.java
@@ -580,7 +580,6 @@ protected void logStatement(Logger logger, String prefix, SimpleStatement statem
         var floatPos = start;
         for (int i = 0; i < 5; i++) {
           var nextPos = cql.indexOf(",", floatPos + 1, end);
-          System.out.println(nextPos);
           if (nextPos > -1) {
             floatPos = nextPos;
           } else {
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/query/OrderByCqlClause.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/query/OrderByCqlClause.java
index bdca1a07a..a5934f244 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/query/OrderByCqlClause.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/query/OrderByCqlClause.java
@@ -36,4 +36,17 @@ public interface OrderByCqlClause extends Function<Select, Select>, CQLClause {
   default boolean fullyCoversCommand() {
     return false;
   }
+
+  /**
+   * Called to get the default limit for the query when using this order by.
+   *
+   * <p>ANN order by has a different default that regular queries to limit the number of rows in
+   * consideration.
+   *
+   * @return Default returns {@link Integer#MAX_VALUE} so queries can read all the rows in the
+   *     table.
+   */
+  default Integer getDefaultLimit() {
+    return Integer.MAX_VALUE;
+  }
 }
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableOrderByANNCqlClause.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableOrderByANNCqlClause.java
index 72c78e085..9ab68d3a4 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableOrderByANNCqlClause.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/TableOrderByANNCqlClause.java
@@ -17,11 +17,13 @@ public class TableOrderByANNCqlClause implements OrderByCqlClause {
 
   private final ApiColumnDef apiColumnDef;
   private final CqlVector<Float> vector;
+  private final Integer defaultLimit;
 
-  public TableOrderByANNCqlClause(ApiColumnDef apiColumnDef, CqlVector<Float> vector) {
+  public TableOrderByANNCqlClause(
+      ApiColumnDef apiColumnDef, CqlVector<Float> vector, Integer defaultLimit) {
     this.apiColumnDef = Objects.requireNonNull(apiColumnDef, "apiColumnDef must not be null");
     this.vector = Objects.requireNonNull(vector, "vector must not be null");
-
+    this.defaultLimit = Objects.requireNonNull(defaultLimit, "defaultLimit must not be null");
     // sanity check
     if (apiColumnDef.type().typeName() != ApiTypeName.VECTOR) {
       throw new IllegalArgumentException(
@@ -38,4 +40,9 @@ public Select apply(Select select) {
   public boolean fullyCoversCommand() {
     return true;
   }
+
+  @Override
+  public Integer getDefaultLimit() {
+    return defaultLimit;
+  }
 }
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/ReadCommandResolver.java b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/ReadCommandResolver.java
index 1f5b22965..044d0c56a 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/ReadCommandResolver.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/ReadCommandResolver.java
@@ -74,7 +74,8 @@ protected GenericOperation<TableSchemaObject, ReadAttempt<TableSchemaObject>> bu
 
     // if the user did not provide a limit,we read all the possible rows. Paging is then handled
     // by the driver pagination
-    int commandLimit = command.limit().orElseGet(() -> Integer.MAX_VALUE);
+    int commandLimit =
+        command.limit().orElseGet(() -> orderByWithWarnings.target().getDefaultLimit());
 
     int commandSkip = command.skip().orElse(0);
 
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/sort/TableCqlSortClauseResolver.java b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/sort/TableCqlSortClauseResolver.java
index b070bebb8..4dceb7b66 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/sort/TableCqlSortClauseResolver.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/sort/TableCqlSortClauseResolver.java
@@ -251,7 +251,10 @@ private WithWarnings<OrderByCqlClause> resolveVectorSort(
         "Vector sorting on column {}", cqlIdentifierToMessageString(vectorSortColumn.name()));
     var cqlVector = CqlVectorUtil.floatsToCqlVector(vectorSortExpression.vector());
     return WithWarnings.of(
-        new TableOrderByANNCqlClause(vectorSortColumn, cqlVector),
+        new TableOrderByANNCqlClause(
+            vectorSortColumn,
+            cqlVector,
+            commandContext.getConfig(OperationsConfig.class).maxVectorSearchLimit()),
         List.of(WarningException.Code.ZERO_FILTER_OPERATIONS));
   }