Allow grouping on bool attribute.

vespa-engine · Dec 10, 2024 · 5d77c94 · 5d77c94
1 parent da44071
commit 5d77c94
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 14 deletions.
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/GroupingValidator.java b/container-search/src/main/java/com/yahoo/search/grouping/GroupingValidator.java
@@ -76,23 +76,25 @@ static private boolean isPrimitiveAttribute(AttributesConfig.Attribute attribute
                 datatype == AttributesConfig.Attribute.Datatype.DOUBLE;
     }
 
-    static private boolean isSingleRawAttribute(AttributesConfig.Attribute attribute) {
-        return  attribute.datatype() == AttributesConfig.Attribute.Datatype.RAW &&
+    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 allowSingleRawAttribute) {
+    private void verifyHasAttribute(String attributeName, boolean isMapLookup) {
         var attribute = attributes.get(attributeName);
         if (attribute == null) {
             throw new UnavailableAttributeException(clusterName, attributeName);
         }
-        if (isPrimitiveAttribute(attribute) || (isSingleRawAttribute(attribute) && allowSingleRawAttribute)) {
+        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() + "'" +
-                (allowSingleRawAttribute ? "" : " for map lookup"));
+                (isMapLookup ? " for map lookup" : ""));
     }
 
     private void verifyCompatibleAttributeTypes(String keyAttributeName,
@@ -116,14 +118,14 @@ private class MyVisitor implements ExpressionVisitor {
         @Override
         public void visitExpression(GroupingExpression exp) {
             if (exp instanceof AttributeMapLookupValue mapLookup) {
-                verifyHasAttribute(mapLookup.getKeyAttribute(), false);
-                verifyHasAttribute(mapLookup.getValueAttribute(), false);
+                verifyHasAttribute(mapLookup.getKeyAttribute(), true);
+                verifyHasAttribute(mapLookup.getValueAttribute(), true);
                 if (mapLookup.hasKeySourceAttribute()) {
-                    verifyHasAttribute(mapLookup.getKeySourceAttribute(), false);
+                    verifyHasAttribute(mapLookup.getKeySourceAttribute(), true);
                     verifyCompatibleAttributeTypes(mapLookup.getKeyAttribute(), mapLookup.getKeySourceAttribute());
                 }
             } else if (exp instanceof AttributeValue) {
-                verifyHasAttribute(((AttributeValue) exp).getAttributeName(), true);
+                verifyHasAttribute(((AttributeValue) exp).getAttributeName(), false);
             }
         }
     }

diff --git a/container-search/src/test/java/com/yahoo/search/grouping/GroupingValidatorTestCase.java b/container-search/src/test/java/com/yahoo/search/grouping/GroupingValidatorTestCase.java
@@ -159,6 +159,14 @@ private static void unsupported_attribute_type_throws(String name, AttributesCon
         }
     }
 
+    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);
@@ -176,11 +184,12 @@ void reference_attribute_throws() {
 
     @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())))");
+        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