align adapter-azureaisearch to API changes

This commit is contained in:
Jonathan Shook 2024-10-24 14:20:30 -05:00
parent a7e5580203
commit 109ca07b2d
17 changed files with 145 additions and 126 deletions

View File

@ -47,13 +47,18 @@
<version>${revision}</version> <version>${revision}</version>
<scope>compile</scope> <scope>compile</scope>
</dependency> </dependency>
<dependency> <dependency>
<groupId>com.azure</groupId> <groupId>com.azure</groupId>
<artifactId>azure-search-documents</artifactId> <artifactId>azure-search-documents</artifactId>
<version>11.7.0</version>
</dependency> </dependency>
<dependency> <dependency>
<groupId>com.azure</groupId> <groupId>com.azure</groupId>
<artifactId>azure-identity</artifactId> <artifactId>azure-identity</artifactId>
<version>1.13.2</version>
</dependency> </dependency>
</dependencies> </dependencies>
</project> </project>

View File

@ -17,6 +17,7 @@ package io.nosqlbench.adapter.azureaisearch;
import java.util.function.Function; import java.util.function.Function;
import java.util.function.LongFunction;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp; import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp;
import io.nosqlbench.adapters.api.activityimpl.OpMapper; import io.nosqlbench.adapters.api.activityimpl.OpMapper;
@ -31,19 +32,19 @@ import io.nosqlbench.nb.api.labels.NBLabels;
import static io.nosqlbench.adapter.azureaisearch.AzureAISearchAdapterUtils.AZURE_AI_SEARCH; import static io.nosqlbench.adapter.azureaisearch.AzureAISearchAdapterUtils.AZURE_AI_SEARCH;
@Service(value = DriverAdapter.class, selector = AZURE_AI_SEARCH) @Service(value = DriverAdapter.class, selector = AZURE_AI_SEARCH)
public class AzureAISearchDriverAdapter extends BaseDriverAdapter<AzureAISearchBaseOp<?>, AzureAISearchSpace> { public class AzureAISearchDriverAdapter extends BaseDriverAdapter<AzureAISearchBaseOp, AzureAISearchSpace> {
public AzureAISearchDriverAdapter(NBComponent parentComponent, NBLabels labels) { public AzureAISearchDriverAdapter(NBComponent parentComponent, NBLabels labels) {
super(parentComponent, labels); super(parentComponent, labels);
} }
@Override @Override
public OpMapper<AzureAISearchBaseOp<?>> getOpMapper() { public OpMapper<AzureAISearchBaseOp,AzureAISearchSpace> getOpMapper() {
return new AzureAISearchOpMapper(this); return new AzureAISearchOpMapper(this);
} }
@Override @Override
public Function<String, ? extends AzureAISearchSpace> getSpaceInitializer(NBConfiguration cfg) { public LongFunction<AzureAISearchSpace> getSpaceInitializer(NBConfiguration cfg) {
return (s) -> new AzureAISearchSpace(s, cfg); return (s) -> new AzureAISearchSpace(s, cfg);
} }

View File

@ -32,7 +32,9 @@ import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.templating.ParsedOp; import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.nosqlbench.engine.api.templating.TypeAndTarget; import io.nosqlbench.engine.api.templating.TypeAndTarget;
public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp<?>> { import java.util.function.LongFunction;
public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp, AzureAISearchSpace> {
private static final Logger logger = LogManager.getLogger(AzureAISearchOpMapper.class); private static final Logger logger = LogManager.getLogger(AzureAISearchOpMapper.class);
private final AzureAISearchDriverAdapter adapter; private final AzureAISearchDriverAdapter adapter;
@ -40,7 +42,8 @@ public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp<?>> {
* Create a new {@code AzureAISearchOpMapper} implementing the {@link OpMapper}. * Create a new {@code AzureAISearchOpMapper} implementing the {@link OpMapper}.
* interface. * interface.
* *
* @param adapter The associated {@link AzureAISearchDriverAdapter} * @param adapter
* The associated {@link AzureAISearchDriverAdapter}
*/ */
public AzureAISearchOpMapper(AzureAISearchDriverAdapter adapter) { public AzureAISearchOpMapper(AzureAISearchDriverAdapter adapter) {
this.adapter = adapter; this.adapter = adapter;
@ -50,12 +53,15 @@ public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp<?>> {
* Given an instance of a {@link ParsedOp} returns the appropriate * Given an instance of a {@link ParsedOp} returns the appropriate
* {@link AzureAISearchBaseOpDispenser} subclass. * {@link AzureAISearchBaseOpDispenser} subclass.
* *
* @param op The {@link ParsedOp} to be evaluated * @param op
* The {@link ParsedOp} to be evaluated
* @param spaceInitF
* @return The correct {@link AzureAISearchBaseOpDispenser} subclass based on * @return The correct {@link AzureAISearchBaseOpDispenser} subclass based on
* the op type * the op type
*/ */
@Override @Override
public OpDispenser<? extends AzureAISearchBaseOp<?>> apply(ParsedOp op) { public OpDispenser<AzureAISearchBaseOp> apply(ParsedOp op, LongFunction<AzureAISearchSpace> spaceInitF) {
TypeAndTarget<AzureAISearchOpType, String> typeAndTarget = op.getTypeAndTarget(AzureAISearchOpType.class, TypeAndTarget<AzureAISearchOpType, String> typeAndTarget = op.getTypeAndTarget(AzureAISearchOpType.class,
String.class, "type", "target"); String.class, "type", "target");
logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'"); logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'");
@ -65,11 +71,14 @@ public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp<?>> {
case create_or_update_index -> case create_or_update_index ->
new AzureAISearchCreateOrUpdateIndexOpDispenser(adapter, op, typeAndTarget.targetFunction); new AzureAISearchCreateOrUpdateIndexOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_indexes -> new AzureAISearchListIndexesOpDispenser(adapter, op, typeAndTarget.targetFunction); case list_indexes -> new AzureAISearchListIndexesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case upload_documents -> new AzureAISearchUploadDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction); case upload_documents ->
case search_documents -> new AzureAISearchSearchDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction); new AzureAISearchUploadDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
case search_documents ->
new AzureAISearchSearchDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
// default -> throw new RuntimeException( // default -> throw new RuntimeException(
// "Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " + "mapping parsed op " + op); // "Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " + "mapping parsed op " + op);
}; };
} }
} }

View File

@ -20,6 +20,7 @@ import java.io.IOException;
import java.nio.file.Files; import java.nio.file.Files;
import java.nio.file.Paths; import java.nio.file.Paths;
import io.nosqlbench.adapters.api.activityimpl.uniform.BaseSpace;
import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger; import org.apache.logging.log4j.Logger;
@ -50,9 +51,8 @@ import io.nosqlbench.nb.api.config.standard.Param;
* "https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/search/azure-search-documents/">Azure * "https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/search/azure-search-documents/">Azure
* AI Search Java searchIndexClient</a> * AI Search Java searchIndexClient</a>
*/ */
public class AzureAISearchSpace implements AutoCloseable { public class AzureAISearchSpace extends BaseSpace {
private final static Logger logger = LogManager.getLogger(AzureAISearchSpace.class); private final static Logger logger = LogManager.getLogger(AzureAISearchSpace.class);
private final String name;
private final NBConfiguration cfg; private final NBConfiguration cfg;
protected SearchIndexClient searchIndexClient; protected SearchIndexClient searchIndexClient;
@ -62,11 +62,11 @@ public class AzureAISearchSpace implements AutoCloseable {
* contextual information needed to interact with the <b>Azure AI Search</b> * contextual information needed to interact with the <b>Azure AI Search</b>
* database instance. * database instance.
* *
* @param name The name of this space * @param idx The name of this space
* @param cfg The configuration ({@link NBConfiguration}) for this nb run * @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/ */
public AzureAISearchSpace(String name, NBConfiguration cfg) { public AzureAISearchSpace(long idx, NBConfiguration cfg) {
this.name = name; super(idx);
this.cfg = cfg; this.cfg = cfg;
} }

View File

@ -26,33 +26,32 @@ import com.azure.search.documents.indexes.SearchIndexClient;
import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.CycleOp; import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.CycleOp;
public abstract class AzureAISearchBaseOp<T> implements CycleOp<Object> { public abstract class AzureAISearchBaseOp<REQUEST, RESULT> implements CycleOp<RESULT> {
protected final static Logger logger = LogManager.getLogger(AzureAISearchBaseOp.class); protected final static Logger logger = LogManager.getLogger(AzureAISearchBaseOp.class);
protected final SearchIndexClient searchIndexClient; protected final SearchIndexClient searchIndexClient;
protected final T request; protected final REQUEST request;
protected final LongFunction<Object> apiCall; protected final LongFunction<RESULT> apiCall;
public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, T requestParam) { public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, REQUEST requestParam) {
this.searchIndexClient = searchIndexClient; this.searchIndexClient = searchIndexClient;
this.request = requestParam; this.request = requestParam;
this.apiCall = this::applyOp; this.apiCall = this::applyOp;
} }
public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, T requestParam, LongFunction<Object> call) { public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, REQUEST requestParam, LongFunction<RESULT> call) {
this.searchIndexClient = searchIndexClient; this.searchIndexClient = searchIndexClient;
this.request = requestParam; this.request = requestParam;
this.apiCall = call; this.apiCall = call;
} }
@Override @Override
public final Object apply(long value) { public final RESULT apply(long value) {
logger.trace(() -> "applying op: " + this); logger.trace(() -> "applying op: " + this);
try { try {
Object result = applyOp(value); RESULT result = applyOp(value);
return result; return result;
} catch (Exception e) { } catch (Exception e) {
if (e instanceof RuntimeException rte) { if (e instanceof RuntimeException rte) {
@ -63,7 +62,7 @@ public abstract class AzureAISearchBaseOp<T> implements CycleOp<Object> {
} }
}; };
public abstract Object applyOp(long value); public abstract RESULT applyOp(long value);
@Override @Override
public String toString() { public String toString() {

View File

@ -17,11 +17,13 @@ package io.nosqlbench.adapter.azureaisearch.ops;
import com.azure.search.documents.SearchClient; import com.azure.search.documents.SearchClient;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.util.AutocompletePagedIterable;
import com.azure.search.documents.util.SearchPagedIterable;
public abstract class AzureAISearchClientBaseOp<T> extends AzureAISearchBaseOp<T> { public abstract class AzureAISearchClientBaseOp<REQUEST, RESULT> extends AzureAISearchBaseOp<REQUEST, RESULT> {
protected final SearchClient searchClient; protected final SearchClient searchClient;
public AzureAISearchClientBaseOp(SearchIndexClient searchIdxClient, SearchClient searchClnt, T requestParam) { public AzureAISearchClientBaseOp(SearchIndexClient searchIdxClient, SearchClient searchClnt, REQUEST requestParam) {
super(searchIdxClient, requestParam); super(searchIdxClient, requestParam);
this.searchClient = searchClnt; this.searchClient = searchClnt;
} }

View File

@ -18,14 +18,14 @@ package io.nosqlbench.adapter.azureaisearch.ops;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.models.SearchIndex; import com.azure.search.documents.indexes.models.SearchIndex;
public class AzureAISearchCreateOrUpdateIndexOp extends AzureAISearchBaseOp<SearchIndex> { public class AzureAISearchCreateOrUpdateIndexOp extends AzureAISearchBaseOp<SearchIndex,SearchIndex> {
public AzureAISearchCreateOrUpdateIndexOp(SearchIndexClient client, SearchIndex request) { public AzureAISearchCreateOrUpdateIndexOp(SearchIndexClient client, SearchIndex request) {
super(client, request); super(client, request);
} }
@Override @Override
public Object applyOp(long value) { public SearchIndex applyOp(long value) {
SearchIndex createResponse = null; SearchIndex createResponse = null;
try { try {
createResponse = searchIndexClient.createOrUpdateIndex(request); createResponse = searchIndexClient.createOrUpdateIndex(request);

View File

@ -17,14 +17,14 @@ package io.nosqlbench.adapter.azureaisearch.ops;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
public class AzureAISearchDeleteIndexOp extends AzureAISearchBaseOp<String> { public class AzureAISearchDeleteIndexOp extends AzureAISearchBaseOp<String,String> {
public AzureAISearchDeleteIndexOp(SearchIndexClient client, String request) { public AzureAISearchDeleteIndexOp(SearchIndexClient client, String request) {
super(client, request); super(client, request);
} }
@Override @Override
public Object applyOp(long value) { public String applyOp(long value) {
try { try {
searchIndexClient.deleteIndex(request); searchIndexClient.deleteIndex(request);
logger.debug("Successfully deleted the index: {}", request); logger.debug("Successfully deleted the index: {}", request);

View File

@ -19,14 +19,14 @@ import com.azure.core.http.rest.PagedIterable;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.models.SearchIndex; import com.azure.search.documents.indexes.models.SearchIndex;
public class AzureAISearchListIndexesOp extends AzureAISearchBaseOp<String> { public class AzureAISearchListIndexesOp extends AzureAISearchBaseOp<String,String> {
public AzureAISearchListIndexesOp(SearchIndexClient client, String request) { public AzureAISearchListIndexesOp(SearchIndexClient client, String request) {
super(client, request); super(client, request);
} }
@Override @Override
public Object applyOp(long value) { public String applyOp(long value) {
try { try {
PagedIterable<SearchIndex> response = searchIndexClient.listIndexes(); PagedIterable<SearchIndex> response = searchIndexClient.listIndexes();
response.forEach((index) -> { response.forEach((index) -> {

View File

@ -19,9 +19,10 @@ import com.azure.core.util.Context;
import com.azure.search.documents.SearchClient; import com.azure.search.documents.SearchClient;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.models.SearchOptions; import com.azure.search.documents.models.SearchOptions;
import com.azure.search.documents.util.AutocompletePagedIterable;
import com.azure.search.documents.util.SearchPagedIterable; import com.azure.search.documents.util.SearchPagedIterable;
public class AzureAISearchSearchDocumentsOp extends AzureAISearchClientBaseOp<SearchOptions> { public class AzureAISearchSearchDocumentsOp extends AzureAISearchClientBaseOp<SearchOptions,SearchPagedIterable> {
public AzureAISearchSearchDocumentsOp(SearchIndexClient searchIndexClient, SearchClient searchClient, public AzureAISearchSearchDocumentsOp(SearchIndexClient searchIndexClient, SearchClient searchClient,
SearchOptions request) { SearchOptions request) {
@ -29,7 +30,7 @@ public class AzureAISearchSearchDocumentsOp extends AzureAISearchClientBaseOp<Se
} }
@Override @Override
public Object applyOp(long value) { public SearchPagedIterable applyOp(long value) {
SearchPagedIterable searchDocsResponse = null; SearchPagedIterable searchDocsResponse = null;
try { try {
searchDocsResponse = searchClient.search(null, // we've not implemented other complex searches yet here. searchDocsResponse = searchClient.search(null, // we've not implemented other complex searches yet here.

View File

@ -22,7 +22,7 @@ import com.azure.search.documents.SearchDocument;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.models.IndexDocumentsResult; import com.azure.search.documents.models.IndexDocumentsResult;
public class AzureAISearchUploadDocumentsOp extends AzureAISearchClientBaseOp<SearchDocument> { public class AzureAISearchUploadDocumentsOp extends AzureAISearchClientBaseOp<SearchDocument,IndexDocumentsResult> {
public AzureAISearchUploadDocumentsOp(SearchIndexClient searchIndexClient, SearchClient searchClient, public AzureAISearchUploadDocumentsOp(SearchIndexClient searchIndexClient, SearchClient searchClient,
SearchDocument request) { SearchDocument request) {
@ -30,7 +30,7 @@ public class AzureAISearchUploadDocumentsOp extends AzureAISearchClientBaseOp<Se
} }
@Override @Override
public Object applyOp(long value) { public IndexDocumentsResult applyOp(long value) {
IndexDocumentsResult uploadDocsResponse = null; IndexDocumentsResult uploadDocsResponse = null;
try { try {
uploadDocsResponse = searchClient.uploadDocuments(List.of(request)); uploadDocsResponse = searchClient.uploadDocuments(List.of(request));

View File

@ -28,13 +28,13 @@ import io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter; import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapters.api.templating.ParsedOp; import io.nosqlbench.adapters.api.templating.ParsedOp;
public abstract class AzureAISearchBaseOpDispenser<T> public abstract class AzureAISearchBaseOpDispenser<REQUEST,RESULT>
extends BaseOpDispenser<AzureAISearchBaseOp<T>, AzureAISearchSpace> { extends BaseOpDispenser<AzureAISearchBaseOp, AzureAISearchSpace> {
protected final LongFunction<AzureAISearchSpace> azureAISearchSpaceFunction; protected final LongFunction<AzureAISearchSpace> azureAISearchSpaceFunction;
protected final LongFunction<SearchIndexClient> clientFunction; protected final LongFunction<SearchIndexClient> clientFunction;
private final LongFunction<? extends AzureAISearchBaseOp<T>> opF; private final LongFunction<? extends AzureAISearchBaseOp<REQUEST,RESULT>> opF;
private final LongFunction<T> paramF; private final LongFunction<REQUEST> paramF;
@SuppressWarnings("rawtypes") @SuppressWarnings("rawtypes")
protected AzureAISearchBaseOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op, protected AzureAISearchBaseOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
@ -57,14 +57,14 @@ public abstract class AzureAISearchBaseOpDispenser<T>
return (AzureAISearchDriverAdapter) adapter; return (AzureAISearchDriverAdapter) adapter;
} }
public abstract LongFunction<T> getParamFunc(LongFunction<SearchIndexClient> clientF, ParsedOp op, public abstract LongFunction<REQUEST> getParamFunc(LongFunction<SearchIndexClient> clientF, ParsedOp op,
LongFunction<String> targetF); LongFunction<String> targetF);
public abstract LongFunction<AzureAISearchBaseOp<T>> createOpFunc(LongFunction<T> paramF, public abstract LongFunction<AzureAISearchBaseOp<REQUEST,RESULT>> createOpFunc(LongFunction<REQUEST> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF); LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF);
@Override @Override
public AzureAISearchBaseOp<T> getOp(long value) { public AzureAISearchBaseOp<REQUEST,RESULT> getOp(long value) {
return opF.apply(value); return opF.apply(value);
} }

View File

@ -53,7 +53,7 @@ import io.nosqlbench.nb.api.errors.OpConfigError;
* "https://learn.microsoft.com/en-us/java/api/com.azure.search.documents?view=azure-java-stable">Index * "https://learn.microsoft.com/en-us/java/api/com.azure.search.documents?view=azure-java-stable">Index
* docs</a> * docs</a>
*/ */
public class AzureAISearchCreateOrUpdateIndexOpDispenser extends AzureAISearchBaseOpDispenser<SearchIndex> { public class AzureAISearchCreateOrUpdateIndexOpDispenser extends AzureAISearchBaseOpDispenser<SearchIndex,SearchIndex> {
private SearchField searchField; private SearchField searchField;
private VectorSearchProfile vsProfile; private VectorSearchProfile vsProfile;
@ -290,7 +290,7 @@ public class AzureAISearchCreateOrUpdateIndexOpDispenser extends AzureAISearchBa
} }
@Override @Override
public LongFunction<AzureAISearchBaseOp<SearchIndex>> createOpFunc(LongFunction<SearchIndex> paramF, public LongFunction<AzureAISearchBaseOp<SearchIndex,SearchIndex>> createOpFunc(LongFunction<SearchIndex> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) { LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchCreateOrUpdateIndexOp(clientF.apply(l), paramF.apply(l)); return l -> new AzureAISearchCreateOrUpdateIndexOp(clientF.apply(l), paramF.apply(l));
} }

View File

@ -34,7 +34,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
* "https://learn.microsoft.com/en-us/rest/api/searchservice/">REST * "https://learn.microsoft.com/en-us/rest/api/searchservice/">REST
* API</a>. * API</a>.
*/ */
public class AzureAISearchDeleteIndexOpDispenser extends AzureAISearchBaseOpDispenser<String> { public class AzureAISearchDeleteIndexOpDispenser extends AzureAISearchBaseOpDispenser<String,String> {
public AzureAISearchDeleteIndexOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op, public AzureAISearchDeleteIndexOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) { LongFunction<String> targetF) {
@ -48,7 +48,7 @@ public class AzureAISearchDeleteIndexOpDispenser extends AzureAISearchBaseOpDisp
} }
@Override @Override
public LongFunction<AzureAISearchBaseOp<String>> createOpFunc(LongFunction<String> paramF, public LongFunction<AzureAISearchBaseOp<String,String>> createOpFunc(LongFunction<String> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) { LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchDeleteIndexOp(clientF.apply(l), paramF.apply(l)); return l -> new AzureAISearchDeleteIndexOp(clientF.apply(l), paramF.apply(l));
} }

View File

@ -27,7 +27,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
/** /**
* Lists the indexes available. * Lists the indexes available.
*/ */
public class AzureAISearchListIndexesOpDispenser extends AzureAISearchBaseOpDispenser<String> { public class AzureAISearchListIndexesOpDispenser extends AzureAISearchBaseOpDispenser<String,String> {
public AzureAISearchListIndexesOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op, public AzureAISearchListIndexesOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) { LongFunction<String> targetF) {
super(adapter, op, targetF); super(adapter, op, targetF);
@ -40,7 +40,7 @@ public class AzureAISearchListIndexesOpDispenser extends AzureAISearchBaseOpDisp
} }
@Override @Override
public LongFunction<AzureAISearchBaseOp<String>> createOpFunc(LongFunction<String> paramF, public LongFunction<AzureAISearchBaseOp<String,String>> createOpFunc(LongFunction<String> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) { LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchListIndexesOp(clientF.apply(l), paramF.apply(l)); return l -> new AzureAISearchListIndexesOp(clientF.apply(l), paramF.apply(l));
} }

View File

@ -26,6 +26,7 @@ import com.azure.search.documents.models.VectorQuery;
import com.azure.search.documents.models.VectorSearchOptions; import com.azure.search.documents.models.VectorSearchOptions;
import com.azure.search.documents.models.VectorizedQuery; import com.azure.search.documents.models.VectorizedQuery;
import com.azure.search.documents.util.SearchPagedIterable;
import io.nosqlbench.adapter.azureaisearch.AzureAISearchDriverAdapter; import io.nosqlbench.adapter.azureaisearch.AzureAISearchDriverAdapter;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp; import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchSearchDocumentsOp; import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchSearchDocumentsOp;
@ -40,7 +41,7 @@ import io.nosqlbench.nb.api.errors.OpConfigError;
* "https://learn.microsoft.com/en-us/azure/search/vector-search-how-to-query?tabs=query-2024-07-01%2Cfilter-2024-07-01%2Cbuiltin-portal#vector-query-request">How * "https://learn.microsoft.com/en-us/azure/search/vector-search-how-to-query?tabs=query-2024-07-01%2Cfilter-2024-07-01%2Cbuiltin-portal#vector-query-request">How
* to query/vector search</a> * to query/vector search</a>
*/ */
public class AzureAISearchSearchDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchOptions> { public class AzureAISearchSearchDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchOptions, SearchPagedIterable> {
public AzureAISearchSearchDocumentsOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op, public AzureAISearchSearchDocumentsOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) { LongFunction<String> targetF) {
super(adapter, op, targetF); super(adapter, op, targetF);
@ -70,7 +71,7 @@ public class AzureAISearchSearchDocumentsOpDispenser extends AzureAISearchBaseOp
} }
@Override @Override
public LongFunction<AzureAISearchBaseOp<SearchOptions>> createOpFunc(LongFunction<SearchOptions> paramF, public LongFunction<AzureAISearchBaseOp<SearchOptions,SearchPagedIterable>> createOpFunc(LongFunction<SearchOptions> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) { LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchSearchDocumentsOp(clientF.apply(l), return l -> new AzureAISearchSearchDocumentsOp(clientF.apply(l),
clientF.apply(l).getSearchClient(targetF.apply(l)), paramF.apply(l)); clientF.apply(l).getSearchClient(targetF.apply(l)), paramF.apply(l));

View File

@ -21,6 +21,7 @@ import java.util.function.LongFunction;
import com.azure.search.documents.SearchDocument; import com.azure.search.documents.SearchDocument;
import com.azure.search.documents.indexes.SearchIndexClient; import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.models.IndexDocumentsResult;
import io.nosqlbench.adapter.azureaisearch.AzureAISearchDriverAdapter; import io.nosqlbench.adapter.azureaisearch.AzureAISearchDriverAdapter;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp; import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchUploadDocumentsOp; import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchUploadDocumentsOp;
@ -34,7 +35,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
* "https://learn.microsoft.com/en-us/java/api/com.azure.search.documents?view=azure-java-stable">Index * "https://learn.microsoft.com/en-us/java/api/com.azure.search.documents?view=azure-java-stable">Index
* docs</a> * docs</a>
*/ */
public class AzureAISearchUploadDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchDocument> { public class AzureAISearchUploadDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchDocument, IndexDocumentsResult> {
public AzureAISearchUploadDocumentsOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op, public AzureAISearchUploadDocumentsOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) { LongFunction<String> targetF) {
super(adapter, op, targetF); super(adapter, op, targetF);
@ -55,7 +56,7 @@ public class AzureAISearchUploadDocumentsOpDispenser extends AzureAISearchBaseOp
} }
@Override @Override
public LongFunction<AzureAISearchBaseOp<SearchDocument>> createOpFunc(LongFunction<SearchDocument> paramF, public LongFunction<AzureAISearchBaseOp<SearchDocument,IndexDocumentsResult>> createOpFunc(LongFunction<SearchDocument> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) { LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchUploadDocumentsOp(clientF.apply(l), return l -> new AzureAISearchUploadDocumentsOp(clientF.apply(l),
clientF.apply(l).getSearchClient(targetF.apply(l)), paramF.apply(l)); clientF.apply(l).getSearchClient(targetF.apply(l)), paramF.apply(l));