incremental progress on opensearch

This commit is contained in:
Jonathan Shook 2024-02-07 14:21:27 -06:00
parent 582d22af69
commit a60c43c765
11 changed files with 428 additions and 16 deletions

View File

@ -16,7 +16,8 @@
package io.nosqlbench.adapter.opensearch;
import io.nosqlbench.adapter.opensearch.dispensers.CreateIndexOpDispenser;
import io.nosqlbench.adapter.opensearch.dispensers.*;
import io.nosqlbench.adapter.opensearch.ops.UpdateOp;
import io.nosqlbench.adapters.api.activityimpl.OpDispenser;
import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverSpaceCache;
@ -39,6 +40,10 @@ public class OpenSearchOpMapper implements OpMapper<Op> {
op.getTypeAndTarget(OpenSearchOpTypes.class, String.class, "verb", "index");
return switch (typeAndTarget.enumId) {
case create_index -> new CreateIndexOpDispenser(adapter, op);
case delete_index -> new DeleteIndexOpDispenser(adapter, op);
case index -> new IndexOpDispenser(adapter,op);
case update -> new UpdateOpDispenser(adapter,op);
case delete -> new DeleteOpDispenser(adapter,op);
default -> throw new RuntimeException("Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " +
"mapping parsed op " + op);
};

View File

@ -17,5 +17,9 @@
package io.nosqlbench.adapter.opensearch;
public enum OpenSearchOpTypes {
create_index
create_index,
delete_index,
index,
update,
delete
}

View File

@ -0,0 +1,94 @@
/*
* Copyright (c) 2024 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.opensearch.dispensers;
import io.nosqlbench.adapter.opensearch.OpenSearchAdapter;
import io.nosqlbench.adapter.opensearch.ops.CreateIndexOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
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 java.util.Map;
import java.util.function.LongFunction;
public class DeleteIndexOpDispenser extends BaseOpenSearchOpDispenser {
public DeleteIndexOpDispenser(OpenSearchAdapter adapter, ParsedOp op) {
super(adapter, op);
}
/**
* {@see
* <a href="https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-vector-search.html">doc
* </a>}
* <pre>{@code
* {
* "mappings": {
* "properties": {
* "value": {
* "type": "dense_vector",
* "dims": TEMPLATE(dimensions, 25),
* "index": true,
* "similarity": "TEMPLATE(similarity_function, cosine)"
* },
* "key": {
* "type": "text"
* }
* }
* }
* }}</pre>
*
* @return
*/
@Override
public LongFunction<CreateIndexOp> createOpFunc(LongFunction<OpenSearchClient> clientF, ParsedOp op) {
CreateIndexRequest.Builder eb = new CreateIndexRequest.Builder();
LongFunction<CreateIndexRequest.Builder> bfunc = l -> new CreateIndexRequest.Builder().index("testindex1");
bfunc = op.enhanceFunc(bfunc, "mappings", Map.class, this::resolveTypeMapping);
LongFunction<CreateIndexRequest.Builder> finalBfunc = bfunc;
return (long l) -> new CreateIndexOp(clientF.apply(l), finalBfunc.apply(l).build());
}
// https://opensearch.org/docs/latest/search-plugins/knn/knn-index/
private CreateIndexRequest.Builder resolveTypeMapping(CreateIndexRequest.Builder eb, Map<?, ?> mappings) {
TypeMapping.Builder builder = new TypeMapping.Builder().properties(
Map.of(
"p1",
new Property.Builder().knnVector(new KnnVectorProperty.Builder()
.dimension(23)
.method(
new KnnVectorMethod.Builder()
.name("hnsw")
.engine("faiss")
.spaceType("l2")
.parameters(Map.of("ef_construction", JsonData.of(256),"m",JsonData.of(8)))
.build()
).build()
).build()
))
.indexField(new IndexField.Builder()
.enabled(true).build())
.fieldNames(new FieldNamesField.Builder()
.enabled(true).build()
);
return eb.mappings(b -> builder);
}
}

View File

@ -0,0 +1,94 @@
/*
* Copyright (c) 2024 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.opensearch.dispensers;
import io.nosqlbench.adapter.opensearch.OpenSearchAdapter;
import io.nosqlbench.adapter.opensearch.ops.CreateIndexOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
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 java.util.Map;
import java.util.function.LongFunction;
public class DeleteOpDispenser extends BaseOpenSearchOpDispenser {
public DeleteOpDispenser(OpenSearchAdapter adapter, ParsedOp op) {
super(adapter, op);
}
/**
* {@see
* <a href="https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-vector-search.html">doc
* </a>}
* <pre>{@code
* {
* "mappings": {
* "properties": {
* "value": {
* "type": "dense_vector",
* "dims": TEMPLATE(dimensions, 25),
* "index": true,
* "similarity": "TEMPLATE(similarity_function, cosine)"
* },
* "key": {
* "type": "text"
* }
* }
* }
* }}</pre>
*
* @return
*/
@Override
public LongFunction<CreateIndexOp> createOpFunc(LongFunction<OpenSearchClient> clientF, ParsedOp op) {
CreateIndexRequest.Builder eb = new CreateIndexRequest.Builder();
LongFunction<CreateIndexRequest.Builder> bfunc = l -> new CreateIndexRequest.Builder().index("testindex1");
bfunc = op.enhanceFunc(bfunc, "mappings", Map.class, this::resolveTypeMapping);
LongFunction<CreateIndexRequest.Builder> finalBfunc = bfunc;
return (long l) -> new CreateIndexOp(clientF.apply(l), finalBfunc.apply(l).build());
}
// https://opensearch.org/docs/latest/search-plugins/knn/knn-index/
private CreateIndexRequest.Builder resolveTypeMapping(CreateIndexRequest.Builder eb, Map<?, ?> mappings) {
TypeMapping.Builder builder = new TypeMapping.Builder().properties(
Map.of(
"p1",
new Property.Builder().knnVector(new KnnVectorProperty.Builder()
.dimension(23)
.method(
new KnnVectorMethod.Builder()
.name("hnsw")
.engine("faiss")
.spaceType("l2")
.parameters(Map.of("ef_construction", JsonData.of(256),"m",JsonData.of(8)))
.build()
).build()
).build()
))
.indexField(new IndexField.Builder()
.enabled(true).build())
.fieldNames(new FieldNamesField.Builder()
.enabled(true).build()
);
return eb.mappings(b -> builder);
}
}

View File

@ -0,0 +1,60 @@
/*
* Copyright (c) 2024 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.opensearch.dispensers;
import io.nosqlbench.adapter.opensearch.OpenSearchAdapter;
import io.nosqlbench.adapter.opensearch.ops.IndexOp;
import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.Op;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch._types.OpType;
import org.opensearch.client.opensearch._types.VersionType;
import org.opensearch.client.opensearch.core.IndexRequest;
import java.util.Map;
import java.util.function.LongFunction;
public class IndexOpDispenser extends BaseOpenSearchOpDispenser {
public IndexOpDispenser(OpenSearchAdapter adapter, ParsedOp op) {
super(adapter, op);
}
@Override
public LongFunction<? extends Op> createOpFunc(LongFunction<OpenSearchClient> clientF, ParsedOp op) {
LongFunction<IndexRequest.Builder<?>> func = l -> new IndexRequest.Builder<>();
func = op.enhanceFuncOptionally(func, "index",String.class, IndexRequest.Builder::index);
func = op.enhanceFuncOptionally(func,"id",String.class, IndexRequest.Builder::id);
func = op.enhanceFuncOptionally(func,"ifPrimaryTerm",long.class, IndexRequest.Builder::ifPrimaryTerm);
func = op.enhanceFuncOptionally(func,"ifSeqNo",long.class,IndexRequest.Builder::ifSeqNo);
func = op.enhanceFuncOptionally(func,"pipeline", String.class, IndexRequest.Builder::pipeline);
func = op.enhanceFuncOptionally(func,"routing", String.class, IndexRequest.Builder::routing);
func = op.enhanceFuncOptionally(func,"requireAlias", boolean.class, IndexRequest.Builder::requireAlias);
func = op.enhanceFuncOptionally(func,"version", long.class, IndexRequest.Builder::version);
func = op.enhanceEnumOptionally(func,"opType", OpType.class,IndexRequest.Builder::opType);
func = op.enhanceEnumOptionally(func,"versionType", VersionType.class,IndexRequest.Builder::versionType);
func = op.enhanceFunc(func,"document",Object.class,(b1,d) -> this.bindDocument(b1,d));
// TODO: func = op.enhanceFuncOptionally(func,"timeout",) ...
LongFunction<IndexRequest.Builder<?>> finalFunc = func;
return l -> new IndexOp(clientF.apply(l), finalFunc.apply(l).build());
}
private IndexRequest.Builder<?> bindDocument(IndexRequest.Builder builder, Object docdata) {
return builder.document(docdata);
}
}

View File

@ -0,0 +1,73 @@
/*
* Copyright (c) 2024 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.opensearch.dispensers;
import io.nosqlbench.adapter.opensearch.OpenSearchAdapter;
import io.nosqlbench.adapter.opensearch.ops.CreateIndexOp;
import io.nosqlbench.adapter.opensearch.ops.UpdateOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.opensearch.client.json.JsonData;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch._types.mapping.*;
import org.opensearch.client.opensearch.core.UpdateRequest;
import org.opensearch.client.opensearch.indices.CreateIndexRequest;
import java.util.Map;
import java.util.function.LongFunction;
public class UpdateOpDispenser extends BaseOpenSearchOpDispenser {
public UpdateOpDispenser(OpenSearchAdapter adapter, ParsedOp op) {
super(adapter, op);
}
/**
* {@see
* <a href="https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-vector-search.html">doc
* </a>}
* <pre>{@code
* {
* "mappings": {
* "properties": {
* "value": {
* "type": "dense_vector",
* "dims": TEMPLATE(dimensions, 25),
* "index": true,
* "similarity": "TEMPLATE(similarity_function, cosine)"
* },
* "key": {
* "type": "text"
* }
* }
* }
* }}</pre>
*
* @return
*/
@Override
public LongFunction<UpdateOp> createOpFunc(LongFunction<OpenSearchClient> clientF, ParsedOp op) {
return null;
// LongFunction<UpdateRequest.Builder> bfunc = l -> new UpdateRequest.Builder();
// op.getAsRequiredFunction("type")
// return l -> new UpdateOp(clientF.apply(l),bfunc.apply(l).build());
// bfunc = op.enhanceFunc(bfunc, "mappings", Map.class, this::resolveTypeMapping);
//
// LongFunction<CreateIndexRequest.Builder> finalBfunc = bfunc;
// return (long l) -> new CreateIndexOp(clientF.apply(l), finalBfunc.apply(l).build());
}
}

View File

@ -27,5 +27,18 @@ public abstract class BaseOpenSearchOp implements CycleOp<Object> {
}
@Override
public abstract Object apply(long value);
public final Object apply(long value) {
try {
Object result = applyOp(value);
return result;
} catch (Exception e) {
if (e instanceof RuntimeException rte) {
throw rte;
} else {
throw new RuntimeException(e);
}
}
};
public abstract Object applyOp(long value) throws Exception;
}

View File

@ -31,12 +31,9 @@ public class CreateIndexOp extends BaseOpenSearchOp {
}
@Override
public Object apply(long value) {
try {
public Object applyOp(long value) throws Exception {
CreateIndexResponse response = client.indices().create(rq);
return response;
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}

View File

@ -31,11 +31,7 @@ public class DeleteIndexOp extends BaseOpenSearchOp {
}
@Override
public Object apply(long value) {
try {
public Object applyOp(long value) throws IOException {
return client.indices().delete(rq);
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}

View File

@ -0,0 +1,38 @@
/*
* Copyright (c) 2024 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.opensearch.ops;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch.core.IndexRequest;
import org.opensearch.client.opensearch.core.IndexResponse;
import org.opensearch.client.opensearch.core.UpdateRequest;
import java.io.IOException;
public class IndexOp extends BaseOpenSearchOp {
private final IndexRequest<?> rq;
public IndexOp(OpenSearchClient client, IndexRequest<?> rq) {
super(client);
this.rq = rq;
}
public Object applyOp(long value) throws IOException {
return client.index(rq);
}
}

View File

@ -0,0 +1,38 @@
/*
* Copyright (c) 2024 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.opensearch.ops;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch.core.UpdateRequest;
import java.io.IOException;
public class UpdateOp extends BaseOpenSearchOp {
private final UpdateRequest rq;
private final Class<?> doctype;
public UpdateOp(OpenSearchClient client, UpdateRequest rq, Class<?> doctype) {
super(client);
this.rq = rq;
this.doctype = doctype;
}
@Override
public Object applyOp(long value) throws IOException {
return client.update(rq, doctype);
}
}