prevent an optimization from destroying variables (#5604)

arangod/Aql/OptimizerRules.cpp

@@ -3422,7 +3422,10 @@ void arangodb::aql::collectInClusterRule(Optimizer* opt,
   SmallVector<ExecutionNode*> nodes{a};
   plan->findNodesOfType(nodes, EN::COLLECT, true);
   
+  std::unordered_set<Variable const*> allUsed;
+
   for (auto& node : nodes) {
+    allUsed.clear();
     auto used = node->getVariablesUsedHere();
 
     // found a node we need to replace in the plan
@@ -3438,6 +3441,16 @@ void arangodb::aql::collectInClusterRule(Optimizer* opt,
     while (current != nullptr) {
       bool eligible = true;
       
+      // check if any of the nodes we pass use a variable that will not be
+      // available after we insert a new COLLECT on top of it (note: COLLECT
+      // will eliminate all variables from the scope but its own)
+      for (auto const& it : current->getVariablesUsedHere()) {
+        if (current->getType() != EN::GATHER) {
+          // Gather nodes are taken care of separately below
+          allUsed.emplace(it);
+        }
+      }
+
       for (auto const& it : current->getVariablesSetHere()) {
         if (std::find(used.begin(), used.end(), it) != used.end()) {
           eligible = false;
@@ -3464,6 +3477,18 @@ void arangodb::aql::collectInClusterRule(Optimizer* opt,
         }
 
         if (previous != nullptr) {
+          for (auto const& otherVariable : allUsed) {
+            auto const setHere = collectNode->getVariablesSetHere();
+            if (std::find(setHere.begin(), setHere.end(), otherVariable) == setHere.end()) {
+              eligible = false;
+              break;
+            }
+          }
+
+          if (!eligible) {
+            break;
+          }
+
           bool removeGatherNodeSort = false;
 
           if (collectNode->aggregationMethod() == CollectOptions::CollectMethod::COUNT) {

js/server/tests/aql/aql-optimizer-collect-count.js

@@ -183,6 +183,25 @@ function optimizerCountTestSuite () {
       }
     },
 
+////////////////////////////////////////////////////////////////////////////////
+/// @brief test count
+////////////////////////////////////////////////////////////////////////////////
+
+    testCountTotal : function () {
+      var query = "FOR j IN " + c.name() + " COLLECT WITH COUNT INTO count RETURN count";
+
+      var results = AQL_EXECUTE(query);
+      assertEqual(1, results.json.length);
+      assertEqual(1000, results.json[0]);
+
+      var plan = AQL_EXPLAIN(query).plan;
+      // must not have a SortNode
+      assertEqual(-1, plan.nodes.map(function(node) { return node.type; }).indexOf("SortNode"));
+      if (isCluster) {
+        assertNotEqual(-1, plan.rules.indexOf("collect-in-cluster"));
+      }
+    },
+
 ////////////////////////////////////////////////////////////////////////////////
 /// @brief test count
 ////////////////////////////////////////////////////////////////////////////////
@@ -190,7 +209,26 @@ function optimizerCountTestSuite () {
     testCountTotalNested : function () {
       var query = "FOR i IN 1..2 FOR j IN " + c.name() + " COLLECT WITH COUNT INTO count RETURN count";
 
-      var results = AQL_EXECUTE(query);
+      var results = AQL_EXECUTE(query, null, { optimizer: { rules: ["-interchange-adjacent-enumerations"] } });
+      assertEqual(1, results.json.length);
+      assertEqual(2000, results.json[0]);
+
+      var plan = AQL_EXPLAIN(query).plan;
+      // must not have a SortNode
+      assertEqual(-1, plan.nodes.map(function(node) { return node.type; }).indexOf("SortNode"));
+      if (isCluster) {
+        assertNotEqual(-1, plan.rules.indexOf("collect-in-cluster"));
+      }
+    },
+
+////////////////////////////////////////////////////////////////////////////////
+/// @brief test count
+////////////////////////////////////////////////////////////////////////////////
+
+    testCountTotalNested2 : function () {
+      var query = "FOR j IN " + c.name() + " FOR i IN 1..2 COLLECT WITH COUNT INTO count RETURN count";
+
+      var results = AQL_EXECUTE(query, null, { optimizer: { rules: ["-interchange-adjacent-enumerations"] } });
       assertEqual(1, results.json.length);
       assertEqual(2000, results.json[0]);