Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Validate attribute vector data types for grouping requests. #33021

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,36 @@ 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 isSingleRawOrBoolAttribute(AttributesConfig.Attribute attribute) {
var datatype = attribute.datatype();
return (datatype == AttributesConfig.Attribute.Datatype.RAW ||
datatype == AttributesConfig.Attribute.Datatype.BOOL) &&
attribute.collectiontype() == AttributesConfig.Attribute.Collectiontype.SINGLE;
}

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

private void verifyCompatibleAttributeTypes(String keyAttributeName,
Expand All @@ -92,14 +118,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(), true);
verifyHasAttribute(mapLookup.getValueAttribute(), true);
if (mapLookup.hasKeySourceAttribute()) {
verifyHasAttribute(mapLookup.getKeySourceAttribute());
verifyHasAttribute(mapLookup.getKeySourceAttribute(), true);
verifyCompatibleAttributeTypes(mapLookup.getKeyAttribute(), mapLookup.getKeySourceAttribute());
}
} else if (exp instanceof AttributeValue) {
verifyHasAttribute(((AttributeValue) exp).getAttributeName());
verifyHasAttribute(((AttributeValue) exp).getAttributeName(), false);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -145,11 +145,79 @@ 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());
}
}

private static void supported_attribute_type_works(String name, AttributesConfig.Attribute.Datatype.Enum datatype) {
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())))");

}

@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() {
supported_attribute_type_works("raw", AttributesConfig.Attribute.Datatype.RAW);
}

@Test
void bool_attribute_is_ok() {
supported_attribute_type_works("mybool", AttributesConfig.Attribute.Datatype.BOOL);
}

@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 +240,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
Loading