Merge pull request #1725 from marklogic/feature/17146-top-ann-k

MLE-17146 Added support for annTopK

Author: rjrudin

marklogic-client-api/src/main/java/com/marklogic/client/expression/PlanBuilder.java

@@ -1498,6 +1498,21 @@ public interface ModifyPlan extends PreparePlan, PlanBuilderBase.ModifyPlanBase
   * @return  a ModifyPlan object
   */
   public abstract ModifyPlan bindAs(PlanColumn column, ServerExpression expression);
+
+	/**
+	 * Facilitates Approximate Nearest Neighbor (ann) vector search. Given a query vector, it searches for K nearest
+	 * neighbor vector embeddings that are stored in the database.
+	 *
+	 * @param k This positive integer k is the top-K rows to return as a result of the index lookup.
+	 * @param vectorColumn The column representing the vector ann-indexed column to perform the index lookup against.
+	 * @param queryVector Specifies the query vector to perform the index lookup with.
+	 * @param distanceColumn Optional output column that captures the values of the distance metric of the vectors retrieved from the index associated with vectorColumn and the queryVector.
+	 * @param queryTolerance Specifies the query tolerance to help balance recall and search time. The value is between 0.0 and 1.0. At 0.0, the recall will be highest. At 1.0 the recall will likely see a large degradation, but queries will be quick. The default value is 0.0.
+	 * @return
+	 * @since 7.1.0
+	 */
+  ModifyPlan annTopK(int k, PlanColumn vectorColumn, ServerExpression queryVector, PlanColumn distanceColumn, float queryTolerance);
+
 /**
   * This method restricts the left row set to rows where a row with the same columns and values doesn't exist in the right row set.
   * @param right  The row set from the right view.

marklogic-client-api/src/main/java/com/marklogic/client/impl/BaseTypeImpl.java

@@ -409,7 +409,11 @@ static <T extends BaseArgImpl> T[] convertList(Object[] items, Class<T> as) {
       Arrays.stream(items)
         .map(item -> {
           if (item != null && !as.isInstance(item)) {
-            throw new IllegalArgumentException("expected "+as.getName()+" argument instead of "+item.getClass().getName());
+			  // Prior to 7.1.0, this threw an exception, as it was requiring every item to be an instance of the given
+			  // class. This meant that a primitive value could never be passed. But that forces the server to support
+			  // both a primitive value and a "wrapped" value (e.g. with ns=xs, fn=float, args=value) for every
+			  // argument. This instead assumes that it can just write the item as-is and the server will accept it.
+			  return (BaseArgImpl) serializedPlanBuilder -> serializedPlanBuilder.append(item);
           }
           return (T) item;
         })

marklogic-client-api/src/main/java/com/marklogic/client/impl/PlanBuilderSubImpl.java

@@ -986,6 +986,13 @@ static class ModifyPlanSubImpl
       super(prior, fnPrefix, fnName, fnArgs);
     }
 
+	  @Override
+	  public ModifyPlan annTopK(int k, PlanColumn vectorColumn, ServerExpression queryVector, PlanColumn distanceColumn, float queryTolerance) {
+		  return new PlanBuilderSubImpl.ModifyPlanSubImpl(this, "op", "annTopK", new Object[]{
+			  k, vectorColumn, queryVector, distanceColumn, queryTolerance
+		  });
+	  }
+
 	  @Override
 	  public ModifyPlan patch(String docColumn, PatchBuilder patchDef) {
 		  return new PlanBuilderSubImpl.ModifyPlanSubImpl(this, "op", "patch", new Object[]{ this.col(docColumn), patchDef });

marklogic-client-api/src/test/java/com/marklogic/client/test/rows/VectorTest.java

@@ -38,15 +38,15 @@ void vectorFunctionsHappyPath() {
 		PlanBuilder.ModifyPlan plan =
 			op.fromView("vectors", "persons")
 				.bind(op.as("sampleVector", op.vec.vector(sampleVector)))
-				.bind(op.as("cosineSimilarity", op.vec.cosineSimilarity(op.col("embedding"),op.col("sampleVector"))))
-				.bind(op.as("dotProduct", op.vec.dotProduct(op.col("embedding"),op.col("sampleVector"))))
-				.bind(op.as("euclideanDistance", op.vec.euclideanDistance(op.col("embedding"),op.col("sampleVector"))))
+				.bind(op.as("cosineSimilarity", op.vec.cosineSimilarity(op.col("embedding"), op.col("sampleVector"))))
+				.bind(op.as("dotProduct", op.vec.dotProduct(op.col("embedding"), op.col("sampleVector"))))
+				.bind(op.as("euclideanDistance", op.vec.euclideanDistance(op.col("embedding"), op.col("sampleVector"))))
 				.bind(op.as("dimension", op.vec.dimension(op.col("sampleVector"))))
 				.bind(op.as("normalize", op.vec.normalize(op.col("sampleVector"))))
 				.bind(op.as("magnitude", op.vec.magnitude(op.col("sampleVector"))))
 				.bind(op.as("get", op.vec.get(op.col("sampleVector"), op.xs.integer(2))))
-				.bind(op.as("add", op.vec.add(op.col("embedding"),op.col("sampleVector"))))
-				.bind(op.as("subtract", op.vec.subtract(op.col("embedding"),op.col("sampleVector"))))
+				.bind(op.as("add", op.vec.add(op.col("embedding"), op.col("sampleVector"))))
+				.bind(op.as("subtract", op.vec.subtract(op.col("embedding"), op.col("sampleVector"))))
 				.bind(op.as("base64Encode", op.vec.base64Encode(op.col("sampleVector"))))
 				.bind(op.as("base64Decode", op.vec.base64Decode(op.col("base64Encode"))))
 				.bind(op.as("subVector", op.vec.subvector(op.col("sampleVector"), op.xs.integer(1), op.xs.integer(1))))
@@ -64,15 +64,15 @@ void vectorFunctionsHappyPath() {
 		rows.forEach(row -> {
 //			 Simple a sanity checks to verify that the functions ran. Very little concern about the actual return values.
 			double cosineSimilarity = row.getDouble("cosineSimilarity");
-			assertTrue((cosineSimilarity > 0) && (cosineSimilarity < 1),"Unexpected value: " + cosineSimilarity);
+			assertTrue((cosineSimilarity > 0) && (cosineSimilarity < 1), "Unexpected value: " + cosineSimilarity);
 			double dotProduct = row.getDouble("dotProduct");
 			Assertions.assertTrue(dotProduct > 0, "Unexpected value: " + dotProduct);
 			double euclideanDistance = row.getDouble("euclideanDistance");
 			Assertions.assertTrue(euclideanDistance > 0, "Unexpected value: " + euclideanDistance);
 			assertEquals(3, row.getInt("dimension"));
 			assertEquals(3, ((ArrayNode) row.get("normalize")).size());
 			double magnitude = row.getDouble("magnitude");
-			assertTrue( magnitude > 0, "Unexpected value: " + magnitude);
+			assertTrue(magnitude > 0, "Unexpected value: " + magnitude);
 			assertEquals(3, ((ArrayNode) row.get("add")).size());
 			assertEquals(3, ((ArrayNode) row.get("subtract")).size());
 			assertFalse(row.getString("base64Encode").isEmpty());
@@ -89,7 +89,7 @@ void cosineSimilarity_DimensionMismatch() {
 		PlanBuilder.ModifyPlan plan =
 			op.fromView("vectors", "persons")
 				.bind(op.as("sampleVector", op.vec.vector(twoDimensionalVector)))
-				.bind(op.as("cosineSimilarity", op.vec.cosineSimilarity(op.col("embedding"),op.col("sampleVector"))))
+				.bind(op.as("cosineSimilarity", op.vec.cosineSimilarity(op.col("embedding"), op.col("sampleVector"))))
 				.select(op.col("name"), op.col("summary"), op.col("cosineSimilarity"));
 		Exception exception = assertThrows(FailedRequestException.class, () -> resultRows(plan));
 		String actualMessage = exception.getMessage();
@@ -102,7 +102,7 @@ void cosineSimilarity_InvalidVector() {
 		PlanBuilder.ModifyPlan plan =
 			op.fromView("vectors", "persons")
 				.bind(op.as("sampleVector", invalidVector))
-				.bind(op.as("cosineSimilarity", op.vec.cosineSimilarity(op.col("embedding"),op.col("sampleVector"))))
+				.bind(op.as("cosineSimilarity", op.vec.cosineSimilarity(op.col("embedding"), op.col("sampleVector"))))
 				.select(op.col("name"), op.col("summary"), op.col("cosineSimilarity"));
 		Exception exception = assertThrows(FailedRequestException.class, () -> resultRows(plan));
 		String actualMessage = exception.getMessage();
@@ -111,20 +111,20 @@ void cosineSimilarity_InvalidVector() {
 	}
 
 	@Test
-	// As of 07/26/24, this test will fail with the ML12 develop branch.
-	// However, it will succeed with the 12ea1 build.
-	// See https://progresssoftware.atlassian.net/browse/MLE-15707
+		// As of 07/26/24, this test will fail with the ML12 develop branch.
+		// However, it will succeed with the 12ea1 build.
+		// See https://progresssoftware.atlassian.net/browse/MLE-15707
 	void bindVectorFromDocs() {
 		PlanBuilder.ModifyPlan plan =
 			op.fromSearchDocs(
-				op.cts.andQuery(
-					op.cts.documentQuery("/optic/vectors/alice.json"),
-					op.cts.elementQuery(
-						"person",
-						op.cts.trueQuery()
-    				)
-  				))
-			.bind(op.as("embedding", op.vec.vector(op.xpath("doc", "/person/embedding"))));
+					op.cts.andQuery(
+						op.cts.documentQuery("/optic/vectors/alice.json"),
+						op.cts.elementQuery(
+							"person",
+							op.cts.trueQuery()
+						)
+					))
+				.bind(op.as("embedding", op.vec.vector(op.xpath("doc", "/person/embedding"))));
 		List<RowRecord> rows = resultRows(plan);
 		assertEquals(1, rows.size());
 	}
@@ -138,4 +138,30 @@ void vecVectorWithCol() {
 		List<RowRecord> rows = resultRows(plan);
 		assertEquals(2, rows.size());
 	}
+
+	@Test
+	void annTopK() {
+		PlanBuilder.ModifyPlan plan = op.fromView("vectors", "persons")
+			.annTopK(10, op.col("embedding"), op.vec.vector(sampleVector), op.col("distance"), 0.5f);
+
+		List<RowRecord> rows = resultRows(plan);
+		assertEquals(2, rows.size(), "Verifying that annTopK worked and returned both rows from the view.");
+
+		rows.forEach(row -> {
+			float distance = row.getFloat("distance");
+			assertTrue(distance > 0, "Just verifying that annTopK both worked and put a valid value into the 'distance' column.");
+		});
+	}
+
+	@Test
+	void dslAnnTopK() {
+		String query = "const qualityVector = vec.vector([ 1.1, 2.2, 3.3 ]);\n" +
+			"op.fromView('vectors', 'persons')\n" +
+			"  .bind(op.as('myVector', op.vec.vector(op.col('embedding'))))\n" +
+			"  .annTopK(2, op.col('myVector'), qualityVector, op.col('distance'), 0.5)";
+
+		RawQueryDSLPlan plan = rowManager.newRawQueryDSLPlan(new StringHandle(query));
+		List<RowRecord> rows = resultRows(plan);
+		assertEquals(2, rows.size(), "Just verifying that 'annTopK' works via the DSL and v1/rows.");
+	}
 }