From 0764e045fe724e45547dce5c9ae0932a48b43b52 Mon Sep 17 00:00:00 2001
From: Jan <jsteemann@users.noreply.github.com>
Date: Wed, 3 Apr 2019 15:15:46 +0200
Subject: [PATCH] optimize lookups by shard key for smart vertex collections
 (#8641)

---
 arangod/Aql/Collection.cpp     | 21 +++++++++++++++++----
 arangod/Aql/Collection.h       |  4 +++-
 arangod/Aql/OptimizerRules.cpp |  8 ++++----
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/arangod/Aql/Collection.cpp b/arangod/Aql/Collection.cpp
index c38ea17244..0d93301264 100644
--- a/arangod/Aql/Collection.cpp
+++ b/arangod/Aql/Collection.cpp
@@ -150,12 +150,25 @@ std::shared_ptr<std::vector<std::string>> Collection::shardIds(
 }
 
 /// @brief returns the shard keys of a collection
-std::vector<std::string> Collection::shardKeys() const {
+std::vector<std::string> Collection::shardKeys(bool normalize) const {
   auto coll = getCollection();
-  std::vector<std::string> keys;
-  for (auto const& x : coll->shardKeys()) {
-    keys.emplace_back(x);
+  auto const& originalKeys = coll->shardKeys();
+    
+  if (normalize && 
+      coll->isSmart() && coll->type() == TRI_COL_TYPE_DOCUMENT) {
+    // smart vertex collection always has ["_key:"] as shard keys
+    TRI_ASSERT(originalKeys.size() == 1);
+    TRI_ASSERT(originalKeys[0] == "_key:");
+    // now normalize it this to _key
+    return std::vector<std::string>{ StaticStrings::KeyString };
   }
+  
+  std::vector<std::string> keys;
+  keys.reserve(originalKeys.size());
+  for (auto const& key : originalKeys) {
+    keys.emplace_back(key);
+  }
+
   return keys;
 }
 
diff --git a/arangod/Aql/Collection.h b/arangod/Aql/Collection.h
index 2ca6e5a3dc..5f3310f351 100644
--- a/arangod/Aql/Collection.h
+++ b/arangod/Aql/Collection.h
@@ -103,7 +103,9 @@ struct Collection {
   std::shared_ptr<std::vector<std::string>> shardIds(std::unordered_set<std::string> const& includedShards) const;
 
   /// @brief returns the shard keys of a collection
-  std::vector<std::string> shardKeys() const;
+  /// if "normalize" is true, then the shard keys for a smart vertex collection
+  /// will be reported as "_key" instead of "_key:"
+  std::vector<std::string> shardKeys(bool normalize) const;
 
   size_t numberOfShards() const;
 
diff --git a/arangod/Aql/OptimizerRules.cpp b/arangod/Aql/OptimizerRules.cpp
index f84ce3407c..c82276c7a4 100644
--- a/arangod/Aql/OptimizerRules.cpp
+++ b/arangod/Aql/OptimizerRules.cpp
@@ -599,7 +599,7 @@ std::string getSingleShardId(arangodb::aql::ExecutionPlan const* plan,
   }
 
   // note for which shard keys we need to look for
-  auto shardKeys = collection->shardKeys();
+  auto shardKeys = collection->shardKeys(true);
   std::unordered_set<std::string> toFind;
   for (auto const& it : shardKeys) {
     if (it.find('.') != std::string::npos) {
@@ -1462,7 +1462,7 @@ class PropagateConstantAttributesHelper {
                   // don't remove a smart join attribute access!
                   return;
                 } else {
-                  std::vector<std::string> const& shardKeys = logical->shardKeys();
+                  std::vector<std::string> shardKeys = collection->shardKeys(true);
                   if (std::find(shardKeys.begin(), shardKeys.end(), nameAttribute->getString()) != shardKeys.end()) {
                     // don't remove equality lookups on shard keys, as this may prevent
                     // the restrict-to-single-shard rule from being applied later!
@@ -4812,7 +4812,7 @@ class RemoveToEnumCollFinder final : public WalkerWorker<ExecutionNode> {
               break;  // abort . . .
             }
             // check the remove node's collection is sharded over _key
-            std::vector<std::string> shardKeys = rn->collection()->shardKeys();
+            std::vector<std::string> shardKeys = rn->collection()->shardKeys(false);
             if (shardKeys.size() != 1 || shardKeys[0] != StaticStrings::KeyString) {
               break;  // abort . . .
             }
@@ -4833,7 +4833,7 @@ class RemoveToEnumCollFinder final : public WalkerWorker<ExecutionNode> {
             }
 
             // note for which shard keys we need to look for
-            auto shardKeys = rn->collection()->shardKeys();
+            auto shardKeys = rn->collection()->shardKeys(false);
             std::unordered_set<std::string> toFind;
             for (auto const& it : shardKeys) {
               toFind.emplace(it);