Skip to content

Commit

Permalink
feat: payload util methods
Browse files Browse the repository at this point in the history
  • Loading branch information
Anush authored and Anush committed Dec 14, 2023
1 parent 6ce7f19 commit 09559b1
Show file tree
Hide file tree
Showing 2 changed files with 325 additions and 0 deletions.
160 changes: 160 additions & 0 deletions src/main/java/io/qdrant/client/PayloadUtil.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
package io.qdrant.client;

import static io.qdrant.client.grpc.JsonWithInt.ListValue;
import static io.qdrant.client.grpc.JsonWithInt.NullValue;
import static io.qdrant.client.grpc.JsonWithInt.Struct;
import static io.qdrant.client.grpc.JsonWithInt.Value;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
* Utility methods for working with Qdrant Payloads.
*/
public class PayloadUtil {

private PayloadUtil () {
}

/**
* Converts a map to a payload map.
*
* @param inputMap The input map to convert.
* @return The converted payload map.
*/
public static Map<String, Value> toPayload(Map<String, Object> inputMap) {
Map<String, Value> map = new HashMap<>();
for (Map.Entry<String, Object> entry : inputMap.entrySet()) {
String fieldName = entry.getKey();
Object value = entry.getValue();

Value.Builder valueBuilder = Value.newBuilder();
setValue(valueBuilder, value);
map.put(fieldName, valueBuilder.build());
}
return map;
}

/**
* Converts a map to a payload struct.
*
* @param inputMap The input map to convert.
* @return The converted payload struct.
*/
public static Struct toPayloadStruct(Map<String, Object> inputMap) {
Struct.Builder structBuilder = Struct.newBuilder();
Map<String, Value> map = toPayload(inputMap);
structBuilder.putAllFields(map);
return structBuilder.build();
}

/**
* Converts a payload struct to a Java Map.
*
* @param struct The payload struct to convert.
* @return The converted hash map.
*/
public static Map<String, Object> toMap(Struct struct) {
Map<String, Object> structMap = toMap(struct.getFieldsMap());
return structMap;
}

/**
* Converts a payload map to a Java Map.
*
* @param payload The payload map to convert.
* @return The converted hash map.
*/
public static Map<String, Object> toMap(Map<String, Value> payload) {
Map<String, Object> hashMap = new HashMap<>();
for (Map.Entry<String, Value> entry : payload.entrySet()) {
String fieldName = entry.getKey();
Value fieldValue = entry.getValue();

Object value = valueToObject(fieldValue);
hashMap.put(fieldName, value);
}
return hashMap;
}

/**
* Sets the value of a Value.Builder based on the given object.
*
* @param valueBuilder The Value.Builder to set the value for.
* @param value The object value to set.
*/
static void setValue(Value.Builder valueBuilder, Object value) {
if (value == null) {
valueBuilder.setNullValue(NullValue.NULL_VALUE);
} else if (value instanceof String) {
valueBuilder.setStringValue((String) value);
} else if (value instanceof Integer) {
valueBuilder.setIntegerValue((Integer) value);
} else if (value instanceof Double) {
valueBuilder.setDoubleValue((Double) value);
} else if (value instanceof Boolean) {
valueBuilder.setBoolValue((Boolean) value);
} else if (value instanceof Map) {
valueBuilder.setStructValue(toPayloadStruct((Map<String, Object>) value));
} else if (value instanceof List) {
valueBuilder.setListValue(listToListValue((List<Object>) value));
}
}

/**
* Converts a list to a ListValue.Builder.
*
* @param list The list to convert.
* @return The converted ListValue.Builder.
*/
static ListValue.Builder listToListValue(List<Object> list) {
ListValue.Builder listValueBuilder = ListValue.newBuilder();

for (Object element : list) {
Value.Builder elementBuilder = Value.newBuilder();
setValue(elementBuilder, element);
listValueBuilder.addValues(elementBuilder.build());
}

return listValueBuilder;
}

/**
* Converts a ListValue to an array of objects.
*
* @param listValue The ListValue to convert.
* @return The converted array of objects.
*/
static Object listValueToList(ListValue listValue) {
return listValue.getValuesList().stream().map(PayloadUtil::valueToObject).toArray();
}

/**
* Converts a Value to an object.
*
* @param value The Value to convert.
* @return The converted object.
*/
static Object valueToObject(Value value) {
if (value.hasStringValue()) {
return value.getStringValue();
} else if (value.hasIntegerValue()) {
// int64 is converted to long
// We need to cast it to int
return (int) value.getIntegerValue();
} else if (value.hasDoubleValue()) {
return value.getDoubleValue();
} else if (value.hasBoolValue()) {
return value.getBoolValue();
} else if (value.hasNullValue()) {
return null;
} else if (value.hasStructValue()) {
return toMap(value.getStructValue());
} else if (value.hasListValue()) {
return listValueToList(value.getListValue());
}

return null;
}
}
165 changes: 165 additions & 0 deletions src/test/java/io/qdrant/client/PayloadUtilTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,165 @@
package io.qdrant.client;

import static org.junit.jupiter.api.Assertions.*;

import static io.qdrant.client.grpc.JsonWithInt.ListValue;
import static io.qdrant.client.grpc.JsonWithInt.NullValue;
import static io.qdrant.client.grpc.JsonWithInt.Struct;
import static io.qdrant.client.grpc.JsonWithInt.Value;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;

class PayloadUtilTest {

@Test
void testToPayloadStruct() {
// Test case 1: Empty input map
Map<String, Object> inputMap = new HashMap<>();
Struct payloadStruct = PayloadUtil.toPayloadStruct(inputMap);
assertTrue(payloadStruct.getFieldsMap().isEmpty());

// Test case 2: Input map with different value types
inputMap.put("name", "Elon");
inputMap.put("age", 52);
inputMap.put("isStudent", true);
payloadStruct = PayloadUtil.toPayloadStruct(inputMap);
assertEquals("Elon", payloadStruct.getFieldsMap().get("name").getStringValue());
assertEquals(52, payloadStruct.getFieldsMap().get("age").getIntegerValue());
assertEquals(true, payloadStruct.getFieldsMap().get("isStudent").getBoolValue());
}

@Test
void testToPayload() {
// Test case 1: Empty input map
Map<String, Object> inputMap = new HashMap<>();
Map<String, Value> payload = PayloadUtil.toPayload(inputMap);
assertTrue(payload.isEmpty());

// Test case 2: Input map with different value types
inputMap.put("name", "Elon");
inputMap.put("age", 52);
inputMap.put("isStudent", true);
payload = PayloadUtil.toPayload(inputMap);
assertEquals("Elon", payload.get("name").getStringValue());
assertEquals(52, payload.get("age").getIntegerValue());
assertEquals(true, payload.get("isStudent").getBoolValue());
}

@Test
void testStructtoMap() {
// Test case 1: Empty struct
Struct.Builder structBuilder = Struct.newBuilder();
Struct struct = structBuilder.build();
Map<String, Object> structMap = PayloadUtil.toMap(struct);
assertTrue(structMap.isEmpty());

// Test case 2: Struct with different value types
structBuilder.putFields("name", Value.newBuilder().setStringValue("Elon").build());
structBuilder.putFields("age", Value.newBuilder().setIntegerValue(52).build());
structBuilder.putFields("isStudent", Value.newBuilder().setBoolValue(true).build());
struct = structBuilder.build();
structMap = PayloadUtil.toMap(struct);
assertEquals("Elon", structMap.get("name"));
assertEquals(52, (int) structMap.get("age"));
assertEquals(true, structMap.get("isStudent"));
}

@Test
void testtoMap() {
// Test case 1: Empty payload
Map<String, Value> payload = new HashMap<>();
Map<String, Object> hashMap = PayloadUtil.toMap(payload);
assertTrue(hashMap.isEmpty());

// Test case 2: Payload with different value types
payload.put("name", Value.newBuilder().setStringValue("Elon").build());
payload.put("age", Value.newBuilder().setIntegerValue(52).build());
payload.put("isStudent", Value.newBuilder().setBoolValue(true).build());
hashMap = PayloadUtil.toMap(payload);
assertEquals("Elon", hashMap.get("name"));
assertEquals(52, hashMap.get("age"));
assertEquals(true, hashMap.get("isStudent"));
}

@Test
void testSetValue() {
// Test case 1: Set string value
Value.Builder valueBuilder = Value.newBuilder();
PayloadUtil.setValue(valueBuilder, "Elon");
assertEquals("Elon", valueBuilder.getStringValue());

// Test case 2: Set integer value
valueBuilder = Value.newBuilder();
PayloadUtil.setValue(valueBuilder, 52);
assertEquals(52, valueBuilder.getIntegerValue());

// Test case 3: Set boolean value
valueBuilder = Value.newBuilder();
PayloadUtil.setValue(valueBuilder, true);
assertEquals(true, valueBuilder.getBoolValue());
}

@Test
void testListToListValue() {
// Test case 1: Empty list
List<Object> list = new ArrayList<>();
ListValue.Builder listValueBuilder = PayloadUtil.listToListValue(list);
assertTrue(listValueBuilder.getValuesList().isEmpty());

// Test case 2: List with different value types
list.add("Elon");
list.add(52);
list.add(true);
listValueBuilder = PayloadUtil.listToListValue(list);
assertEquals("Elon", listValueBuilder.getValuesList().get(0).getStringValue());
assertEquals(52, listValueBuilder.getValuesList().get(1).getIntegerValue());
assertEquals(true, listValueBuilder.getValuesList().get(2).getBoolValue());
}

@Test
void testListValueToList() {
// Test case 1: Empty list value
ListValue.Builder listValueBuilder = ListValue.newBuilder();
ListValue listValue = listValueBuilder.build();
Object[] objectList = (Object[]) PayloadUtil.listValueToList(listValue);
assertTrue(objectList.length == 0);

// Test case 2: List value with different value types
listValueBuilder.addValues(Value.newBuilder().setStringValue("Elon").build());
listValueBuilder.addValues(Value.newBuilder().setIntegerValue(52).build());
listValueBuilder.addValues(Value.newBuilder().setBoolValue(true).build());
listValue = listValueBuilder.build();
objectList = (Object[]) PayloadUtil.listValueToList(listValue);

assertEquals("Elon", objectList[0]);
assertEquals(52, objectList[1]);
assertEquals(true, objectList[2]);
}

@Test
void testValueToObject() {
// Test case 1: String value
Value value = Value.newBuilder().setStringValue("Elon").build();
Object object = PayloadUtil.valueToObject(value);
assertEquals("Elon", object);

// Test case 2: Integer value
value = Value.newBuilder().setIntegerValue(52).build();
object = PayloadUtil.valueToObject(value);
assertEquals(52, object);

// Test case 3: Boolean value
value = Value.newBuilder().setBoolValue(true).build();
object = PayloadUtil.valueToObject(value);
assertEquals(true, object);

// Test case 4: Null value
value = Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build();
object = PayloadUtil.valueToObject(value);
assertNull(object);
}
}

0 comments on commit 09559b1

Please sign in to comment.