Apply Pavel review suggestions

Author: MBkkt

axiom/optimizer/Optimization.cpp

@@ -482,6 +482,9 @@ bool isIndexColocated(
           input->distribution().maybeNeedsShuffle(distribution)) {
     return !*needsShuffle;
   }
+  // TODO: Future partition type should be accounted.
+  // TODO: This code actually doesn't feel right.
+  //       We should check it when we will add indexes.
   // True if 'input' is partitioned so that each partitioning key is joined to
   // the corresponding partition key in 'info'.
   for (size_t i = 0; i < input->distribution().partition.size(); ++i) {

axiom/optimizer/Schema.cpp

@@ -367,22 +367,23 @@ std::optional<bool> Distribution::maybeNeedsShuffle(
     // One is gather, the other is not, repartitioning needed.
     return true;
   }
-  if (partition.size() != desired.partition.size()) {
-    // Different number of partition keys, repartitioning needed.
-    // TODO: Future partition type should decide this.
-    return true;
-  }
   return std::nullopt;
 }
 
 bool Distribution::needsShuffle(const Distribution& desired) const {
   if (auto maybe = maybeNeedsShuffle(desired)) {
     return *maybe;
   }
+  // TODO: Future partition type should decide this.
+  // For an example range partitioning may not need shuffle if
+  // the ranges are compatible, e.g. "a, b, c" and "a, b".
+  if (partition.size() != desired.partition.size()) {
+    // Different number of partition keys, repartitioning needed.
+    return true;
+  }
   for (size_t i = 0; i < partition.size(); ++i) {
     if (!partition[i]->sameOrEqual(*desired.partition[i])) {
       // Different partition key, repartitioning needed.
-      // TODO: Future partition type should decide this.
       return true;
     }
   }