From 8e6a14ee438845c838c9048c18c10ecdc46753d1 Mon Sep 17 00:00:00 2001
From: Jonathan Shook <jshook@gmail.com>
Date: Thu, 8 Feb 2024 16:18:22 -0600
Subject: [PATCH] enable knn

---
 .../dispensers/CreateIndexOpDispenser.java    |  2 +
 .../main/resources/activities/osvectors.yaml  | 60 +++++++++++++++++--
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java
index a3dd5b626..29ddb09f4 100644
--- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java
+++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java
@@ -23,6 +23,7 @@ import org.opensearch.client.json.JsonData;
 import org.opensearch.client.opensearch.OpenSearchClient;
 import org.opensearch.client.opensearch._types.mapping.*;
 import org.opensearch.client.opensearch.indices.CreateIndexRequest;
+import org.opensearch.client.opensearch.indices.IndexSettings;
 
 import java.util.Map;
 import java.util.function.LongFunction;
@@ -41,6 +42,7 @@ public class CreateIndexOpDispenser extends BaseOpenSearchOpDispenser {
         CreateIndexRequest.Builder eb = new CreateIndexRequest.Builder();
         LongFunction<CreateIndexRequest.Builder> bfunc =
             l -> new CreateIndexRequest.Builder()
+                .settings(b -> b.knn(true))
                 .index(targetF.apply(1));
         bfunc = op.enhanceFunc(bfunc, "mappings", Map.class, this::resolveTypeMapping);
 
diff --git a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml
index 9add87b00..011f8c035 100644
--- a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml
+++ b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml
@@ -1,16 +1,68 @@
 description: |
   basic test for vectors in open search
   https://www.elastic.co/guide/en/elasticsearch/reference/current/rest-apis.html
+  template vars:
+     TEMPLATE(indexname,vectors_index)
+     TEMPLATE(dimensions,25)
+
+params:
+  driver: opensearch
+  instrument: true
 
 scenarios:
-  schema: run driver=opensearch tags=block:create threads===1 cycles===undef
+  # schema: run tags=block:schema threads===1 cycles===undef
+  # rampup: run tags=block:rampup threads=auto cycles=TEMPLATE(rampup_cycles,10)
+
+  vectors:
+    drop: run tags='block:drop' labels='target:opensearch' threads===1 cycles===UNDEF
+    schema: run tags='block:schema' labels='target:astra' threads===1 cycles===UNDEF
+    rampup: run tags='block:rampup' labels='target:astra' threads=1 cycles=TEMPLATE(trainsize)
+      # threads=auto
+      #  errors=counter,warn,log
+      #    search_and_index: >-
+      #      run alias=search_and_index tags='block:search_and_index,optype=select' labels='target:astra'
+      #      cycles=TEMPLATE(testsize) errors=count,retry rate=100
+    #      errors=counter,warn threads=auto
 
 bindings:
+  id: ToString()
+  test_floatlist: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/test");
+  relevant_indices: HdfFileToIntArray("testdata/TEMPLATE(dataset).hdf5", "/neighbors")
+  distance_floatlist: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/distance")
+  train_floatlist: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/train");
+  synthetic_vectors: HashedFloatVectors(TEMPLATE(dimensions));
 
 blocks:
-  create:
+  drop:
     ops:
-      index:
-        create_index: TEMPLATE(indexname,vectors_index)
+      drop_index:
+        delete_index: TEMPLATE(indexname,vectors_index)
+
+  schema:
+    ops:
+      create_index:
+        create_index: TEMPLATE(indexname)
         mappings:
           m1: v1
+  bulkrampup:
+    ops:
+      bulk_index:
+        bulk_index:
+          - '{"key":"{id}","value":"{train_floatlist}"'
+          -
+
+  rampup:
+    ops:
+      index:
+        index: TEMPLATE(indexname)
+        document:
+          key: "{id}"
+          value: "{train_floatlist}"
+
+
+
+#        document: |
+#          {
+#            "key": "{id}",
+#            "value": {train_floatlist}
+#          }