Implement a new function node (AttributeMapLookupNode) for doing lookup in a map attribute.

3 files changed, 28 insertions, 1 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/grouping/request/AttributeMapLookupValue.java b/container-search/src/main/java/com/yahoo/search/grouping/request/AttributeMapLookupValue.java
index 82c4b6763d8..281e6e53b36 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/request/AttributeMapLookupValue.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/request/AttributeMapLookupValue.java
@@ -6,6 +6,7 @@ package com.yahoo.search.grouping.request;
  *
  * It evaluates to the value found using the given key for the lookup in that attribute.
  * The key is either specified explicitly or found via a key source attribute.
+ * Two underlying attributes are used to represent the map attribute (the key and value attributes).
  *
  * @author geirst
  */
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/vespa/ExpressionConverter.java b/container-search/src/main/java/com/yahoo/search/grouping/vespa/ExpressionConverter.java
index 95384fb12d3..d017fe2edb3 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/vespa/ExpressionConverter.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/vespa/ExpressionConverter.java
@@ -102,6 +102,7 @@ import com.yahoo.searchlib.expression.AddFunctionNode;
 import com.yahoo.searchlib.expression.AggregationRefNode;
 import com.yahoo.searchlib.expression.AndFunctionNode;
 import com.yahoo.searchlib.expression.ArrayAtLookupNode;
+import com.yahoo.searchlib.expression.AttributeMapLookupNode;
 import com.yahoo.searchlib.expression.AttributeNode;
 import com.yahoo.searchlib.expression.BucketResultNode;
 import com.yahoo.searchlib.expression.CatFunctionNode;
@@ -265,7 +266,14 @@ class ExpressionConverter {
             return addArguments(new AndFunctionNode(), (AndFunction)exp);
         }
         if (exp instanceof AttributeMapLookupValue) {
-            return new AttributeNode(((AttributeMapLookupValue)exp).getAttributeName());
+            AttributeMapLookupValue mapLookup = (AttributeMapLookupValue) exp;
+            if (mapLookup.hasKeySourceAttribute()) {
+                return AttributeMapLookupNode.fromKeySourceAttribute(mapLookup.getAttributeName(),
+                        mapLookup.getKeyAttribute(), mapLookup.getValueAttribute(), mapLookup.getKeySourceAttribute());
+            } else {
+                return AttributeMapLookupNode.fromKey(mapLookup.getAttributeName(),
+                        mapLookup.getKeyAttribute(), mapLookup.getValueAttribute(), mapLookup.getKey());
+            }
         }
         if (exp instanceof AttributeValue) {
             return new AttributeNode(((AttributeValue)exp).getAttributeName());
diff --git a/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java b/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java
index b8571aacca4..c64c4d624f2 100644
--- a/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java
@@ -681,6 +681,24 @@ public class RequestBuilderTestCase {
         assertOutput(test);
     }
 
+    @Test
+    public void requireThatAttributeMapLookupNodeIsCreatedFromKey() {
+        RequestTest test = new RequestTest();
+        test.expectedOutput = AttributeMapLookupNode.fromKey("map{\"my_key\"}", "map.key", "map.value", "my_key").toString();
+        test.request = "all(group(map{\"my_key\"}) each(output(count())))";
+        test.outputWriter = (groupingList, transform) -> groupingList.get(0).getLevels().get(0).getExpression().toString();
+        assertOutput(test);
+    }
+
+    @Test
+    public void requireThatAttributeMapLookupNodeIsCreatedFromKeySourceAttribute() {
+        RequestTest test = new RequestTest();
+        test.expectedOutput = AttributeMapLookupNode.fromKeySourceAttribute("map{attribute(key_source)}", "map.key", "map.value", "key_source").toString();
+        test.request = "all(group(map{attribute(key_source)}) each(output(count())))";
+        test.outputWriter = (groupingList, transform) -> groupingList.get(0).getLevels().get(0).getExpression().toString();
+        assertOutput(test);
+    }
+
     private static CompositeContinuation newComposite(EncodableContinuation... conts) {
         CompositeContinuation ret = new CompositeContinuation();
         for (EncodableContinuation cont : conts) {