-
Notifications
You must be signed in to change notification settings - Fork 13
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add describeIndexStats integration test (#47)
## Problem The java sdk lacks integration test for describeIndexStats functionality ## Solution Added integration test for both future and sync client to describeIndexStats. Note this test does include Upsert operation but I'll be testing it rigorously in UpsertAndFetch integration test. ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [X] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan Ran integration tests Note: PineconeClientLiveIntegTest will go away once I add all of the integration tests for data plane. PineconeIndexOperationsClientIntegrationTest will be moved under control plane integration test sub-folder and renamed to CreateAndDeleteIndex, but I'll be adding more tests in their if needed.
- Loading branch information
1 parent
792cb4f
commit d1546e1
Showing
4 changed files
with
259 additions
and
43 deletions.
There are no files selected for viewing
60 changes: 60 additions & 0 deletions
60
src/integration/java/io/pinecone/helpers/BuildUpsertRequest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,60 @@ | ||
package io.pinecone.helpers; | ||
|
||
import com.google.common.primitives.Floats; | ||
import com.google.protobuf.Struct; | ||
import com.google.protobuf.Value; | ||
import io.pinecone.proto.SparseValues; | ||
import io.pinecone.proto.UpsertRequest; | ||
import io.pinecone.proto.Vector; | ||
|
||
import java.util.ArrayList; | ||
import java.util.Arrays; | ||
import java.util.List; | ||
|
||
public class BuildUpsertRequest { | ||
private static final float[][] upsertData = {{1.0F, 2.0F, 3.0F}, {4.0F, 5.0F, 6.0F}, {7.0F, 8.0F, 9.0F}}; | ||
|
||
public static UpsertRequest buildRequiredUpsertRequest() { | ||
String namespace = RandomStringBuilder.build("ns", 8); | ||
List<String> upsertIds = Arrays.asList("v1", "v2", "v3"); | ||
List<Vector> upsertVectors = new ArrayList<>(); | ||
|
||
for (int i = 0; i < upsertData.length; i++) { | ||
upsertVectors.add(Vector.newBuilder() | ||
.addAllValues(Floats.asList(upsertData[i])) | ||
.setMetadata(Struct.newBuilder() | ||
.putFields("some_field", Value.newBuilder().setNumberValue(i).build()) | ||
.build()) | ||
.setId(upsertIds.get(i)) | ||
.build()); | ||
} | ||
|
||
return UpsertRequest.newBuilder() | ||
.addAllVectors(upsertVectors) | ||
.setNamespace(namespace) | ||
.build(); | ||
} | ||
|
||
public static UpsertRequest buildOptionalUpsertRequest() { | ||
String namespace = RandomStringBuilder.build("ns", 8); | ||
List<String> hybridsIds = Arrays.asList("v4", "v5", "v6"); | ||
List<Vector> hybridVectors = new ArrayList<>(); | ||
List<Integer> sparseIndices = Arrays.asList(0, 1, 2); | ||
List<Float> sparseValues = Arrays.asList(0.11f, 0.22f, 0.33f); | ||
for (int i = 0; i < hybridsIds.size(); i++) { | ||
hybridVectors.add( | ||
Vector.newBuilder() | ||
.addAllValues(Floats.asList(upsertData[i])) | ||
.setSparseValues( | ||
SparseValues.newBuilder().addAllIndices(sparseIndices).addAllValues(sparseValues).build() | ||
) | ||
.setId(hybridsIds.get(i)) | ||
.build()); | ||
} | ||
|
||
return UpsertRequest.newBuilder() | ||
.addAllVectors(hybridVectors) | ||
.setNamespace(namespace) | ||
.build(); | ||
} | ||
} |
59 changes: 59 additions & 0 deletions
59
src/integration/java/io/pinecone/helpers/IndexManager.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
package io.pinecone.helpers; | ||
|
||
import io.pinecone.*; | ||
import io.pinecone.model.CreateIndexRequest; | ||
import io.pinecone.model.IndexMeta; | ||
|
||
import java.io.IOException; | ||
import java.util.List; | ||
|
||
public class IndexManager { | ||
public PineconeConnection createIndexIfNotExists(int dimension) throws IOException, InterruptedException { | ||
boolean createNewIndex = false; | ||
String indexName = ""; | ||
PineconeClientConfig config = new PineconeClientConfig() | ||
.withApiKey(System.getenv("PINECONE_API_KEY")) | ||
.withEnvironment(System.getenv("PINECONE_ENVIRONMENT")); | ||
PineconeIndexOperationClient controlPlaneClient = new PineconeIndexOperationClient(config); | ||
List<String> indexList = controlPlaneClient.listIndexes(); | ||
|
||
if (!indexList.isEmpty()) { | ||
indexName = indexList.get(0); | ||
IndexMeta indexMeta = isIndexReady(indexName, controlPlaneClient); | ||
if (indexMeta.getDatabase().getDimension() != dimension) { | ||
createNewIndex = true; | ||
} | ||
} | ||
|
||
if (createNewIndex) { | ||
indexName = RandomStringBuilder.build("index-name", 8); | ||
CreateIndexRequest createIndexRequest = new CreateIndexRequest() | ||
.withIndexName(indexName) | ||
.withDimension(dimension) | ||
.withMetric("euclidean"); | ||
controlPlaneClient.createIndex(createIndexRequest); | ||
} | ||
|
||
PineconeClient dataPlaneClient = new PineconeClient(config); | ||
IndexMeta indexMeta = controlPlaneClient.describeIndex(indexName); | ||
String host = indexMeta.getStatus().getHost(); | ||
|
||
return dataPlaneClient.connect( | ||
new PineconeConnectionConfig() | ||
.withConnectionUrl("https://" + host)); | ||
} | ||
|
||
public static IndexMeta isIndexReady(String indexName, PineconeIndexOperationClient indexOperationClient) | ||
throws IOException, InterruptedException { | ||
IndexMeta indexMeta; | ||
while (true) { | ||
indexMeta = indexOperationClient.describeIndex(indexName); | ||
if (indexMeta.getStatus().getState().equalsIgnoreCase("ready")) { | ||
break; | ||
} | ||
|
||
Thread.sleep(1000); | ||
} | ||
return indexMeta; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.