Merge pull request #907 from marklogic/feature/17147-annTopK

MLE-17147 Added annTopK support

Author: rjrudin

lib/plan-builder-generated.js

@@ -15,7 +15,13 @@
  */
 'use strict';
 
-/* IMPORTANT: Do not edit. This file is generated. */
+/* 
+The contents of this file are mostly an output of the Optic code generator. 
+But some parts have been manually modified. So care needs to be taken when 
+running the code generator to not lose any manual modifications, each of
+which should either be commented or should be obvious, such as changes in 
+version numbers for "since" annotations. 
+*/
 
 const types    = require('./server-types-generated.js');
 const bldrbase = require('./plan-builder-base.js');
@@ -7470,6 +7476,28 @@ class PlanModifyPlan extends PlanPreparePlan {
   constructor(prior, ns, fn, args) {
     super(prior, ns, fn, args);
   }
+
+  /**
+    * This method facilitates Approximate Nearest Neighbor vector search. It searches for K nearest neighbor vector embeddings that are stored in database given a query vector. Provides a client interface to a server function. See {@link http://docs.marklogic.com/ModifyPlan.prototype.annTopK|ModifyPlan.prototype.annTopK}
+    * @method planBuilder.ModifyPlan#annTopK
+    * @since 3.6.0
+    * @param { XsInt } [k] - This positive integer k is the top-K rows to return as a result of the index lookup.
+    * @param { PlanColumn } [vectorColumn] - The column representing the vector ann-indexed column to perform the index lookup against. The columns can be named with a string or a column parameter function such as op:col.
+    * @param { VecVector } [queryVector] - This specifies the query vector to perform the index lookup with.
+    * @param { PlanColumn } [distance] - The column is the output column that returns the values of the distance metric of the vectors retrieved from the index associated with vectorColumn and the qv. The columns can be named with a string or a column parameter function such as op.col.
+    * @param { XsFloat } [queryTolerance] - This 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.
+    * @returns { planBuilder.ModifyPlan }
+    */
+annTopK(...args) {
+    const namer = bldrbase.getNamer(args, 'k');
+    const paramdefs = [['k', [types.XsInt], true, false], ['vector-column', [PlanColumn], true, false], ['query-vector', [types.VecVector], true, false], ['distance', [PlanColumn], false, false], ['query-tolerance', [types.XsFloat], false, false]];
+    const checkedArgs = (namer !== null) ?
+        bldrbase.makeNamedArgs(namer, 'PlanModifyPlan.annTopK', 3, new Set(['k', 'vector-column', 'query-vector', 'distance', 'query-tolerance']), paramdefs, args) :
+        bldrbase.makePositionalArgs('PlanModifyPlan.annTopK', 3, false, paramdefs, args);
+    return new PlanModifyPlan(this, 'op', 'ann-top-k', checkedArgs);
+
+    }
+
 /**
     * This function adds new columns or modifies existing columns based on expressions while preserving existing unmodified columns in the row set. Provides a client interface to a server function. See {@link http://docs.marklogic.com/ModifyPlan.prototype.bind|ModifyPlan.prototype.bind}
     * @method planBuilder.ModifyPlan#bind

test-app/build.gradle

@@ -1,5 +1,17 @@
 // Only used for setting up the test AppServer in MarkLogic
 plugins {
   id 'net.saliman.properties' version '1.5.2'
-  id "com.marklogic.ml-gradle" version "4.6.0"
+  id "com.marklogic.ml-gradle" version "5.0.0"
 }
+
+tasks.register("addMarkLogic12SchemasIfNecessary", com.marklogic.gradle.task.MarkLogicTask) {
+  description = "If testing against MarkLogic 12, include schemas that will not work on MarkLogic 11 or earlier."
+  doLast {
+    def version = new com.marklogic.mgmt.resource.clusters.ClusterManager(getManageClient()).getVersion()
+    if (version.startsWith("12.")) {
+      mlAppConfig.getSchemaPaths().add(new File(getProjectDir(), "src/main/ml-schemas-12").getAbsolutePath())
+    }
+  }
+}
+mlDeploy.dependsOn addMarkLogic12SchemasIfNecessary
+mlLoadSchemas.dependsOn addMarkLogic12SchemasIfNecessary

test-app/docker-compose.yaml

@@ -1,10 +1,10 @@
-version: '1.0'
 name: node-client
 
 services:
 
   marklogic:
-    image: "marklogicdb/marklogic-db:11.1.0-centos-1.1.0"
+    image: "progressofficial/marklogic-db:latest"
+    # image: "ml-docker-db-dev-tierpoint.bed-artifactory.bedford.progress.com/marklogic/marklogic-server-ubi:latest-12"
     platform: linux/amd64
     environment:
       - INSTALL_CONVERTERS=true

test-app/src/main/ml-data/optic/vectors/alice.json

@@ -0,0 +1,11 @@
+{
+	"person": {
+		"name": "Alice",
+		"summary": "About Alice",
+		"embedding": [
+			1.1,
+			2.2,
+			3.3
+		]
+	}
+}

test-app/src/main/ml-data/optic/vectors/bob.json

@@ -0,0 +1,11 @@
+{
+	"person": {
+		"name": "Bob",
+		"summary": "Below Bob",
+		"embedding": [
+			9.9,
+			8.8,
+			7.7
+		]
+	}
+}

test-app/src/main/ml-data/optic/vectors/permissions.properties

@@ -0,0 +1 @@
+*=rest-reader,read,rest-writer,update

test-app/src/main/ml-schemas-12/tde/vector.xml

@@ -0,0 +1,31 @@
+<template xmlns="http://marklogic.com/xdmp/tde">
+	<context>/person</context>
+	<directories><directory>/optic/vectors/</directory></directories>
+	<rows>
+		<row>
+			<schema-name>vectors</schema-name>
+			<view-name>persons</view-name>
+			<columns>
+				<column>
+					<name>name</name>
+					<scalar-type>string</scalar-type>
+					<val>name</val>
+				</column>
+
+				<column>
+					<name>summary</name>
+					<scalar-type>string</scalar-type>
+					<val>summary</val>
+				</column>
+
+				<column>
+					<name>embedding</name>
+					<scalar-type>vector</scalar-type>
+					<val>vec:vector(embedding)</val>
+					<dimension>3</dimension>
+					<invalid-values>reject</invalid-values>
+				</column>
+			</columns>
+		</row>
+	</rows>
+</template>

test-basic/annTopK.js

@@ -0,0 +1,48 @@
+/*
+ * Copyright © 2024 MarkLogic Corporation. All Rights Reserved.
+ */
+'use strict';
+
+const marklogic = require('../');
+const p = marklogic.planBuilder;
+
+const pbb = require('./plan-builder-base');
+const assert = require('assert');
+const testlib = require('../etc/test-lib');
+let serverConfiguration = {};
+const execPlan = pbb.execPlan;
+
+describe('tests for annTopK', function () {
+    before(function (done) {
+        try {
+            testlib.findServerConfiguration(serverConfiguration);
+            setTimeout(() => {
+                if (serverConfiguration.serverVersion < 12) {
+                    this.skip();
+                }
+                done();
+            }, 3000);
+        } catch (error) {
+            done(error);
+        }
+    });
+
+    it('happy path', function (done) {
+        execPlan(p
+            .fromView('vectors', 'persons', '')
+            .annTopK(10, p.col('embedding'), p.vec.vector([1.1, 2.2, 3.3]), p.col('distance'), 0.5)
+            .orderBy(p.col('name'))
+        )
+            .then(function (response) {
+                const rows = response.rows;
+                assert(rows.length === 2, 'Expecting both rows in the view to be returned.');
+                assert(rows[0].name.value === 'Alice');
+                assert(rows[0].distance.type === 'xs:float', 'Verifying that the distance column was populated.');
+                assert(rows[1].name.value === 'Bob');
+                assert(rows[1].distance.type === 'xs:float',  'Verifying that the distance column was populated.');
+                done();
+            })
+            .catch(done);
+    });
+
+});