From 8c83b17249b265edc546d647ec603fb314acf875 Mon Sep 17 00:00:00 2001 From: Rob Rudin Date: Mon, 25 Nov 2024 16:07:31 -0500 Subject: [PATCH] MLE-17147 Added annTopK support --- lib/plan-builder-generated.js | 30 +++++++++++- test-app/build.gradle | 14 +++++- test-app/docker-compose.yaml | 4 +- .../src/main/ml-data/optic/vectors/alice.json | 11 +++++ .../src/main/ml-data/optic/vectors/bob.json | 11 +++++ .../optic/vectors/permissions.properties | 1 + .../src/main/ml-schemas-12/tde/vector.xml | 31 ++++++++++++ test-basic/annTopK.js | 48 +++++++++++++++++++ 8 files changed, 146 insertions(+), 4 deletions(-) create mode 100644 test-app/src/main/ml-data/optic/vectors/alice.json create mode 100644 test-app/src/main/ml-data/optic/vectors/bob.json create mode 100644 test-app/src/main/ml-data/optic/vectors/permissions.properties create mode 100644 test-app/src/main/ml-schemas-12/tde/vector.xml create mode 100644 test-basic/annTopK.js diff --git a/lib/plan-builder-generated.js b/lib/plan-builder-generated.js index e20c2007..896c461c 100755 --- a/lib/plan-builder-generated.js +++ b/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 diff --git a/test-app/build.gradle b/test-app/build.gradle index b8dea2a2..607bb9a8 100644 --- a/test-app/build.gradle +++ b/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 diff --git a/test-app/docker-compose.yaml b/test-app/docker-compose.yaml index 2ca2b473..4af00806 100644 --- a/test-app/docker-compose.yaml +++ b/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 diff --git a/test-app/src/main/ml-data/optic/vectors/alice.json b/test-app/src/main/ml-data/optic/vectors/alice.json new file mode 100644 index 00000000..8dc6f233 --- /dev/null +++ b/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 + ] + } +} diff --git a/test-app/src/main/ml-data/optic/vectors/bob.json b/test-app/src/main/ml-data/optic/vectors/bob.json new file mode 100644 index 00000000..15aed943 --- /dev/null +++ b/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 + ] + } +} diff --git a/test-app/src/main/ml-data/optic/vectors/permissions.properties b/test-app/src/main/ml-data/optic/vectors/permissions.properties new file mode 100644 index 00000000..c9778549 --- /dev/null +++ b/test-app/src/main/ml-data/optic/vectors/permissions.properties @@ -0,0 +1 @@ +*=rest-reader,read,rest-writer,update diff --git a/test-app/src/main/ml-schemas-12/tde/vector.xml b/test-app/src/main/ml-schemas-12/tde/vector.xml new file mode 100644 index 00000000..3c6de4ec --- /dev/null +++ b/test-app/src/main/ml-schemas-12/tde/vector.xml @@ -0,0 +1,31 @@ + diff --git a/test-basic/annTopK.js b/test-basic/annTopK.js new file mode 100644 index 00000000..738d0609 --- /dev/null +++ b/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); + }); + +}); \ No newline at end of file