Refactor replacement of grouping expression attribute nodes.

2 files changed, 36 insertions, 6 deletions

diff --git a/searchlib/src/vespa/searchlib/aggregation/modifiers.cpp b/searchlib/src/vespa/searchlib/aggregation/modifiers.cpp
index 6b6f2a83436..6ef3b6921d5 100644
--- a/searchlib/src/vespa/searchlib/aggregation/modifiers.cpp
+++ b/searchlib/src/vespa/searchlib/aggregation/modifiers.cpp
@@ -10,17 +10,20 @@ using namespace search::expression;
 
 namespace search::aggregation {
 
-bool Attribute2DocumentAccessor::check(const vespalib::Identifiable &obj) const
+bool AttributeNodeReplacer::check(const vespalib::Identifiable &obj) const
 {
     return obj.getClass().inherits(GroupingLevel::classId) || obj.getClass().inherits(AggregationResult::classId) || obj.getClass().inherits(MultiArgFunctionNode::classId);
 }
 
-void Attribute2DocumentAccessor::execute(vespalib::Identifiable &obj)
+void AttributeNodeReplacer::execute(vespalib::Identifiable &obj)
 {
     if (obj.getClass().inherits(GroupingLevel::classId)) {
         GroupingLevel & g(static_cast<GroupingLevel &>(obj));
         if (g.getExpression().getRoot()->inherits(AttributeNode::classId)) {
-            g.setExpression(std::make_unique<DocumentFieldNode>(static_cast<const AttributeNode &>(*g.getExpression().getRoot()).getAttributeName()));
+            auto replacementNode = getReplacementNode(static_cast<const AttributeNode &>(*g.getExpression().getRoot()));
+            if (replacementNode) {
+                g.setExpression(std::move(replacementNode));
+            }
         } else {
             g.getExpression().getRoot()->select(*this, *this);
         }
@@ -30,7 +33,10 @@ void Attribute2DocumentAccessor::execute(vespalib::Identifiable &obj)
         ExpressionNode * e(a.getExpression());
         if (e) {
             if (e->inherits(AttributeNode::classId)) {
-                a.setExpression(std::make_unique<DocumentFieldNode>(static_cast<const AttributeNode &>(*e).getAttributeName()));
+                auto replacementNode = getReplacementNode(static_cast<const AttributeNode &>(*e));
+                if (replacementNode) {
+                    a.setExpression(std::move(replacementNode));
+                }
             } else {
                 e->select(*this, *this);
             }
@@ -40,7 +46,10 @@ void Attribute2DocumentAccessor::execute(vespalib::Identifiable &obj)
         for(size_t i(0), m(v.size()); i < m; i++) {
             ExpressionNode::CP & e(v[i]);
             if (e->inherits(AttributeNode::classId)) {
-                e.reset(new DocumentFieldNode(static_cast<const AttributeNode &>(*e).getAttributeName()));
+                auto replacementNode = getReplacementNode(static_cast<const AttributeNode &>(*e));
+                if (replacementNode) {
+                    e.reset(replacementNode.release());
+                }
             } else {
                 e->select(*this, *this);
             }
@@ -48,6 +57,12 @@ void Attribute2DocumentAccessor::execute(vespalib::Identifiable &obj)
     }
 }
 
+std::unique_ptr<ExpressionNode>
+Attribute2DocumentAccessor::getReplacementNode(const AttributeNode &attributeNode)
+{
+    return std::make_unique<DocumentFieldNode>(attributeNode.getAttributeName());
+}
+
 }
 
 // this function was added by ../../forcelink.sh
diff --git a/searchlib/src/vespa/searchlib/aggregation/modifiers.h b/searchlib/src/vespa/searchlib/aggregation/modifiers.h
index 27f4b3ce6a4..0120cb4eac9 100644
--- a/searchlib/src/vespa/searchlib/aggregation/modifiers.h
+++ b/searchlib/src/vespa/searchlib/aggregation/modifiers.h
@@ -3,14 +3,29 @@
 
 #include <vespa/vespalib/objects/objectoperation.h>
 #include <vespa/vespalib/objects/objectpredicate.h>
+#include <memory>
+
+namespace search::expression {
+
+class ExpressionNode;
+class AttributeNode;
+
+}
 
 namespace search::aggregation {
 
-class Attribute2DocumentAccessor : public vespalib::ObjectOperation, public vespalib::ObjectPredicate
+class AttributeNodeReplacer : public vespalib::ObjectOperation, public vespalib::ObjectPredicate
 {
 private:
     void execute(vespalib::Identifiable &obj) override;
     bool check(const vespalib::Identifiable &obj) const override;
+    virtual std::unique_ptr<search::expression::ExpressionNode> getReplacementNode(const search::expression::AttributeNode &attributeNode) = 0;
+};
+
+class Attribute2DocumentAccessor : public AttributeNodeReplacer
+{
+private:
+    std::unique_ptr<search::expression::ExpressionNode> getReplacementNode(const search::expression::AttributeNode &attributeNode) override;
 };
 
 }