impl of delete all and create collection with class

This commit is contained in:
Mark Wolters 2024-06-12 16:57:27 -04:00
parent 70cd4a0dd4
commit 4618f4d405
8 changed files with 208 additions and 16 deletions

View File

@ -69,6 +69,8 @@ public class DataApiOpMapper implements OpMapper<DataApiBaseOp> {
case count_documents -> new DataApiCountDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
case replace_one -> new DataApiReplaceOneOpDispenser(adapter, op, typeAndTarget.targetFunction);
case find_one_and_replace -> new DataApiFindOneAndReplaceOpDispenser(adapter, op, typeAndTarget.targetFunction);
case delete_all -> new DataApiDeleteAllOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_collection_with_class -> new DataApiCreateCollectionWithClassOpDispenser(adapter, op, typeAndTarget.targetFunction);
};
}
}

View File

@ -39,26 +39,12 @@ public class DataApiCreateCollectionOpDispenser extends DataApiOpDispenser {
private LongFunction<DataApiCreateCollectionOp> createOpFunction(ParsedOp op) {
return (l) -> {
CollectionOptions.CollectionOptionsBuilder optionsBldr = CollectionOptions.builder();
Optional<LongFunction<Integer>> dimFunc = op.getAsOptionalFunction("dimensions", Integer.class);
if (dimFunc.isPresent()) {
LongFunction<Integer> af = dimFunc.get();
optionsBldr.vectorDimension(af.apply(l));
}
// COSINE("cosine"),
// EUCLIDEAN("euclidean"),
// DOT_PRODUCT("dot_product");
Optional<LongFunction<String>> simFunc = op.getAsOptionalFunction("similarity", String.class);
if (simFunc.isPresent()) {
LongFunction<String> sf = simFunc.get();
optionsBldr.vectorSimilarity(SimilarityMetric.fromValue(sf.apply(l)));
}
DataApiCreateCollectionOp dataApiCreateCollectionOp =
new DataApiCreateCollectionOp(
spaceFunction.apply(l).getDatabase(),
targetFunction.apply(l),
optionsBldr.build());
this.getCollectionOptionsFromOp(op, l)
);
return dataApiCreateCollectionOp;
};

View File

@ -0,0 +1,61 @@
/*
* 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.dataapi.opdispensers;
import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiCreateCollectionWithClassOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.function.LongFunction;
public class DataApiCreateCollectionWithClassOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiCreateCollectionWithClassOpDispenser.class);
private final LongFunction<DataApiCreateCollectionWithClassOp> opFunction;
public DataApiCreateCollectionWithClassOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiCreateCollectionWithClassOp> createOpFunction(ParsedOp op) {
return (l) -> new DataApiCreateCollectionWithClassOp(
spaceFunction.apply(l).getDatabase(),
targetFunction.apply(l),
this.getCollectionOptionsFromOp(op, l),
getCreateClass(op, l)
);
}
private Class<?> getCreateClass(ParsedOp op, long l) {
String className = op.getAsFunctionOr("createClass", "com.datastax.astra.client.model.Document").apply(l);
try {
return Class.forName(className);
} catch (ClassNotFoundException e) {
throw new RuntimeException(e);
}
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -0,0 +1,53 @@
/*
* 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.dataapi.opdispensers;
import com.datastax.astra.client.Database;
import io.nosqlbench.adapter.dataapi.DataApiDriverAdapter;
import io.nosqlbench.adapter.dataapi.ops.DataApiBaseOp;
import io.nosqlbench.adapter.dataapi.ops.DataApiDeleteAllOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.function.LongFunction;
public class DataApiDeleteAllOpDispenser extends DataApiOpDispenser {
private static final Logger logger = LogManager.getLogger(DataApiDeleteAllOpDispenser.class);
private final LongFunction<DataApiDeleteAllOp> opFunction;
public DataApiDeleteAllOpDispenser(DataApiDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
this.opFunction = createOpFunction(op);
}
private LongFunction<DataApiDeleteAllOp> createOpFunction(ParsedOp op) {
return (l) -> {
Database db = spaceFunction.apply(l).getDatabase();
return new DataApiDeleteAllOp(
db,
db.getCollection(targetFunction.apply(l))
);
};
}
@Override
public DataApiBaseOp getOp(long value) {
return opFunction.apply(value);
}
}

View File

@ -188,4 +188,19 @@ public abstract class DataApiOpDispenser extends BaseOpDispenser<DataApiBaseOp,
return projection;
}
protected CollectionOptions getCollectionOptionsFromOp(ParsedOp op, long l) {
CollectionOptions.CollectionOptionsBuilder optionsBldr = CollectionOptions.builder();
Optional<LongFunction<Integer>> dimFunc = op.getAsOptionalFunction("dimensions", Integer.class);
if (dimFunc.isPresent()) {
LongFunction<Integer> af = dimFunc.get();
optionsBldr.vectorDimension(af.apply(l));
}
Optional<LongFunction<String>> simFunc = op.getAsOptionalFunction("similarity", String.class);
if (simFunc.isPresent()) {
LongFunction<String> sf = simFunc.get();
optionsBldr.vectorSimilarity(SimilarityMetric.fromValue(sf.apply(l)));
}
return optionsBldr.build();
}
}

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.dataapi.ops;
import com.datastax.astra.client.Database;
import com.datastax.astra.client.model.CollectionOptions;
public class DataApiCreateCollectionWithClassOp extends DataApiBaseOp {
private final String collectionName;
private final CollectionOptions options;
private final Class<?> clazz;
public DataApiCreateCollectionWithClassOp(Database db, String collectionName, CollectionOptions options, Class<?> clazz) {
super(db);
this.collectionName = collectionName;
this.options = options;
this.clazz = clazz;
}
@Override
public Object apply(long value) {
return db.createCollection(collectionName, options, clazz);
}
}

View File

@ -0,0 +1,35 @@
/*
* 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.dataapi.ops;
import com.datastax.astra.client.Collection;
import com.datastax.astra.client.Database;
import com.datastax.astra.client.model.Filter;
public class DataApiDeleteAllOp extends DataApiBaseOp {
private final Collection collection;
public DataApiDeleteAllOp(Database db, Collection collection) {
super(db);
this.collection = collection;
}
@Override
public Object apply(long value) {
return collection.deleteAll();
}
}

View File

@ -40,4 +40,6 @@ public enum DataApiOpType {
count_documents,
replace_one,
find_one_and_replace,
delete_all,
create_collection_with_class,
}