Skip to content

Commit

Permalink
Add Usage to QueryResponseWithUnsignedIndices (#74)
Browse files Browse the repository at this point in the history
## Problem
`Usage` is now returned for data plane operations against serverless
indexes. This was included in the codebase in this PR where we replaced
the `vector_service.proto` file with generated classes that included the
most recent changes to the data plane including read units:
#66

However, the newer `QueryResponseWithUnsignedIndices` class does not map
`usage` so it gets dropped for query requests using this class.

## Solution
Update `QueryResponseWithUnsignedIndices` to handle `usage`.
I also added `toString()` overrides to all of the classes in
`unsigned_indices_model` since debugging/logging was difficult without
it.

## Type of Change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [X] 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)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan
Integration tests for this are pending, I believe @rohanshah18 has
integration test changes in flight so I didn't make any changes there.
We should add assertions on `queryResponse.getUsage().getReadUnits()`.

I've run data plane operations locally using the examples folder to
validate that we're seeing usage and read units on data plane responses.
  • Loading branch information
austin-denoble authored Mar 7, 2024
1 parent 5d087cc commit 86a5ae7
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import io.pinecone.proto.QueryResponse;
import io.pinecone.proto.ScoredVector;
import io.pinecone.proto.Usage;

import java.util.ArrayList;
import java.util.Collections;
Expand All @@ -11,14 +12,18 @@ public class QueryResponseWithUnsignedIndices {

private List<ScoredVectorWithUnsignedIndices> matches;
private String namespace;
private Usage usage;


public QueryResponseWithUnsignedIndices(QueryResponse queryResponse) {
if (queryResponse == null) {
this.matches = Collections.emptyList();
this.namespace = "";
this.usage = null;
} else {
this.matches = convertToScoredVectorWithUnsignedIndices(queryResponse.getMatchesList());
this.namespace = queryResponse.getNamespace();
this.usage = queryResponse.getUsage();
}
}

Expand Down Expand Up @@ -52,4 +57,30 @@ public String getNamespace() {
public void setNamespace(String namespace) {
this.namespace = namespace;
}

public Usage getUsage() { return usage; }

public void setUsage(Usage usage) { this.usage = usage; }

/**
* Convert the given object to string with each line indented by 4 spaces
* (except the first line).
*/
private String toIndentedString(Object o) {
if (o == null) {
return "null";
}
return o.toString().replace("\n", "\n ");
}

@Override
public String toString() {
final StringBuilder sb = new StringBuilder();
sb.append("class QueryResponseWithUnsignedIndices {").append("\n");
sb.append(" matches: ").append(toIndentedString(matches)).append("\n");
sb.append(" namespace: ").append(toIndentedString(namespace)).append("\n");
sb.append(" usage: ").append(usage).append("\n");
sb.append("}");
return sb.toString();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -69,4 +69,28 @@ public SparseValuesWithUnsignedIndices getSparseValuesWithUnsignedIndices() {
public void setSparseValuesWithUnsignedIndices(SparseValuesWithUnsignedIndices sparseValuesWithUnsignedIndices) {
this.sparseValuesWithUnsignedIndices = sparseValuesWithUnsignedIndices;
}

/**
* Convert the given object to string with each line indented by 4 spaces
* (except the first line).
*/
private String toIndentedString(Object o) {
if (o == null) {
return "null";
}
return o.toString().replace("\n", "\n ");
}

@Override
public String toString() {
final StringBuilder sb = new StringBuilder();
sb.append("ScoredVectorWithUnsignedIndices {").append("\n");
sb.append(" score: ").append(toIndentedString(score)).append("\n");
sb.append(" id: ").append(toIndentedString(id)).append("\n");
sb.append(" values: ").append(toIndentedString(values)).append("\n");
sb.append(" metadata: ").append(toIndentedString(metadata)).append("\n");
sb.append(" sparseValuesWithUnsignedIndices: ").append(toIndentedString(sparseValuesWithUnsignedIndices)).append("\n");
sb.append("}");
return sb.toString();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -42,4 +42,25 @@ public List<Float> getValuesList() {
public void setValues(List<Float> values) {
this.values = values;
}

/**
* Convert the given object to string with each line indented by 4 spaces
* (except the first line).
*/
private String toIndentedString(Object o) {
if (o == null) {
return "null";
}
return o.toString().replace("\n", "\n ");
}

@Override
public String toString() {
final StringBuilder sb = new StringBuilder();
sb.append("SparseValuesWithUnsignedIndices {").append("\n");
sb.append(" indicesWithUnsigned32Int: ").append(toIndentedString(indicesWithUnsigned32Int)).append("\n");
sb.append(" values: ").append(toIndentedString(values)).append("\n");
sb.append("}");
return sb.toString();
}
}

0 comments on commit 86a5ae7

Please sign in to comment.