Skip to content

Commit

Permalink
Validate attribute vector data types for grouping requests.
Browse files Browse the repository at this point in the history
  • Loading branch information
toregge committed Dec 10, 2024
1 parent 63cb8f5 commit da44071
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,34 @@ public Result search(Query query, Execution execution) {
return execution.search(query);
}

private void verifyHasAttribute(String attributeName) {
if (!attributes.containsKey(attributeName)) {
static private boolean isPrimitiveAttribute(AttributesConfig.Attribute attribute) {
var datatype = attribute.datatype();
return datatype == AttributesConfig.Attribute.Datatype.INT8 ||
datatype == AttributesConfig.Attribute.Datatype.INT16 ||
datatype == AttributesConfig.Attribute.Datatype.INT32 ||
datatype == AttributesConfig.Attribute.Datatype.INT64 ||
datatype == AttributesConfig.Attribute.Datatype.STRING ||
datatype == AttributesConfig.Attribute.Datatype.FLOAT ||
datatype == AttributesConfig.Attribute.Datatype.DOUBLE;
}

static private boolean isSingleRawAttribute(AttributesConfig.Attribute attribute) {
return attribute.datatype() == AttributesConfig.Attribute.Datatype.RAW &&
attribute.collectiontype() == AttributesConfig.Attribute.Collectiontype.SINGLE;
}

private void verifyHasAttribute(String attributeName, boolean allowSingleRawAttribute) {
var attribute = attributes.get(attributeName);
if (attribute == null) {
throw new UnavailableAttributeException(clusterName, attributeName);
}
if (isPrimitiveAttribute(attribute) || (isSingleRawAttribute(attribute) && allowSingleRawAttribute)) {
return;
}
throw new IllegalInputException("Grouping request references attribute '" +
attributeName + "' with unsupported data type '" + attribute.datatype() +
"' collection type '" + attribute.collectiontype() + "'" +
(allowSingleRawAttribute ? "" : " for map lookup"));
}

private void verifyCompatibleAttributeTypes(String keyAttributeName,
Expand All @@ -92,14 +116,14 @@ private class MyVisitor implements ExpressionVisitor {
@Override
public void visitExpression(GroupingExpression exp) {
if (exp instanceof AttributeMapLookupValue mapLookup) {
verifyHasAttribute(mapLookup.getKeyAttribute());
verifyHasAttribute(mapLookup.getValueAttribute());
verifyHasAttribute(mapLookup.getKeyAttribute(), false);
verifyHasAttribute(mapLookup.getValueAttribute(), false);
if (mapLookup.hasKeySourceAttribute()) {
verifyHasAttribute(mapLookup.getKeySourceAttribute());
verifyHasAttribute(mapLookup.getKeySourceAttribute(), false);
verifyCompatibleAttributeTypes(mapLookup.getKeyAttribute(), mapLookup.getKeySourceAttribute());
}
} else if (exp instanceof AttributeValue) {
verifyHasAttribute(((AttributeValue) exp).getAttributeName());
verifyHasAttribute(((AttributeValue) exp).getAttributeName(), true);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -145,11 +145,70 @@ void key_source_attribute_with_multi_value_collection_type_throws() {
}
}

private static void unsupported_attribute_type_throws(String name, AttributesConfig.Attribute.Datatype.Enum datatype) {
try {
var builder = new AttributesConfig.Builder();
builder.attribute(new AttributesConfig.Attribute.Builder().name(name).datatype(datatype));
validateGrouping(new AttributesConfig(builder),
"all(group(" + name + ") each(output(count())))");
fail("Expected exception");
} catch (IllegalArgumentException e) {
assertEquals("Grouping request references attribute '" + name + "' " +
"with unsupported data type '" + datatype.toString() + "' collection type 'SINGLE'",
e.getMessage());
}
}

@Test
void tensor_attribute_throws() {
unsupported_attribute_type_throws("tensor", AttributesConfig.Attribute.Datatype.TENSOR);
}

@Test
void predicate_attribute_throws() {
unsupported_attribute_type_throws("predicate", AttributesConfig.Attribute.Datatype.PREDICATE);
}

@Test
void reference_attribute_throws() {
unsupported_attribute_type_throws("reference", AttributesConfig.Attribute.Datatype.REFERENCE);
}

@Test
void raw_attribute_is_ok() {
var builder = new AttributesConfig.Builder();
builder.attribute(new AttributesConfig.Attribute.Builder().name("raw")
.datatype(AttributesConfig.Attribute.Datatype.RAW));
validateGrouping(new AttributesConfig(builder),
"all(group(raw) each(output(count())))");
}

@Test
void raw_attribute_cannot_be_used_for_lookup() {
try {
var builder = new AttributesConfig.Builder();
builder.attribute(new AttributesConfig.Attribute.Builder().name("map.key")
.datatype(AttributesConfig.Attribute.Datatype.RAW));
builder.attribute(new AttributesConfig.Attribute.Builder().name("map.value")
.datatype(AttributesConfig.Attribute.Datatype.STRING));
builder.attribute(new AttributesConfig.Attribute.Builder().name("key_source")
.datatype(AttributesConfig.Attribute.Datatype.RAW));
validateGrouping(new AttributesConfig(builder),
"all(group(map{attribute(key_source)}) each(output(count())))");
fail("Expected exception");
} catch (IllegalArgumentException e) {
assertEquals("Grouping request references attribute 'map.key' " +
"with unsupported data type 'RAW' collection type 'SINGLE' for map lookup",
e.getMessage());
}
}

private static AttributesConfig setupMismatchingKeySourceAttribute(boolean matchingDataType) {
AttributesConfig.Builder builder = new AttributesConfig.Builder();
builder.attribute(new AttributesConfig.Attribute.Builder().name("map.key")
.datatype(AttributesConfig.Attribute.Datatype.Enum.STRING));
builder.attribute(new AttributesConfig.Attribute.Builder().name("map.value"));
builder.attribute(new AttributesConfig.Attribute.Builder().name("map.value")
.datatype(AttributesConfig.Attribute.Datatype.STRING));
builder.attribute(new AttributesConfig.Attribute.Builder().name("key_source")
.datatype(matchingDataType ? AttributesConfig.Attribute.Datatype.Enum.STRING :
AttributesConfig.Attribute.Datatype.Enum.INT32)
Expand All @@ -172,7 +231,7 @@ private static AttributesConfig createAttributesConfig(Collection<String> attrib
AttributesConfig.Builder builder = new AttributesConfig.Builder();
for (String attributeName : attributeNames) {
builder.attribute(new AttributesConfig.Attribute.Builder()
.name(attributeName));
.name(attributeName).datatype(AttributesConfig.Attribute.Datatype.INT32));
}
return new AttributesConfig(builder);
}
Expand Down

0 comments on commit da44071

Please sign in to comment.