fix: Fix distribution, order, layouts and plan enumeration

Author: MBkkt

axiom/optimizer/JoinSample.cpp

@@ -219,8 +219,8 @@ std::pair<float, float> sampleJoin(
     const ExprVector& leftKeys,
     SchemaTableCP right,
     const ExprVector& rightKeys) {
-  const uint64_t leftRows = left->numRows();
-  const uint64_t rightRows = right->numRows();
+  const auto leftRows = left->cardinality;
+  const auto rightRows = right->cardinality;
 
   const auto leftCard = keyCardinality(leftKeys);
   const auto rightCard = keyCardinality(rightKeys);

axiom/optimizer/Optimization.cpp

@@ -427,10 +427,9 @@ RelationOpPtr repartitionForAgg(const RelationOpPtr& plan, PlanState& state) {
 
   const auto* agg = state.dt->aggregation;
 
-  // If no grouping and not yet gathered on a single node, add a gather before
-  // final agg.
-  if (agg->groupingKeys().empty() &&
-      !plan->distribution().distributionType.isGather) {
+  // If no grouping and not yet gathered on a single node,
+  // add a gather before final agg.
+  if (agg->groupingKeys().empty()) {
     auto* gather =
         make<Repartition>(plan, Distribution::gather(), plan->columns());
     state.addCost(*gather);
@@ -478,33 +477,20 @@ bool isIndexColocated(
     const IndexInfo& info,
     const ExprVector& lookupValues,
     const RelationOpPtr& input) {
-  const auto& distribution = info.index->distribution;
-  if (distribution.isBroadcast) {
-    return true;
-  }
-
+  const auto& desired = info.index->distribution;
+  if (const auto needsShuffle =
+          input->distribution().maybeNeedsShuffle(desired)) {
+    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'.
-  if (input->distribution().distributionType != distribution.distributionType) {
-    return false;
-  }
-
-  if (input->distribution().partition.empty()) {
-    return false;
-  }
-
-  if (input->distribution().partition.size() != distribution.partition.size()) {
-    return false;
-  }
-
-  for (auto i = 0; i < input->distribution().partition.size(); ++i) {
+  for (size_t i = 0; i < input->distribution().partition.size(); ++i) {
     auto nthKey = position(lookupValues, *input->distribution().partition[i]);
-    if (nthKey != kNotFound) {
-      if (info.schemaColumn(info.lookupKeys.at(nthKey)) !=
-          distribution.partition.at(i)) {
-        return false;
-      }
-    } else {
+    if (nthKey == kNotFound ||
+        info.schemaColumn(info.lookupKeys[nthKey]) != desired.partition[i]) {
       return false;
     }
   }
@@ -613,7 +599,7 @@ PlanObjectSet availableColumns(BaseTableCP baseTable, ColumnGroupCP index) {
 }
 
 bool isBroadcastableSize(PlanP build, PlanState& /*state*/) {
-  return build->cost.fanout < 100'000;
+  return build->cost.resultCardinality() < 100'000;
 }
 
 // The 'other' side gets shuffled to align with 'input'. If 'input' is not
@@ -1024,9 +1010,7 @@ void Optimization::joinByHash(
         candidate.join->isBroadcastableType() &&
         isBroadcastableSize(buildPlan, state)) {
       auto* broadcast = make<Repartition>(
-          buildInput,
-          Distribution::broadcast(plan->distribution().distributionType),
-          buildInput->columns());
+          buildInput, Distribution::broadcast(), buildInput->columns());
       buildState.addCost(*broadcast);
       buildInput = broadcast;
     } else {
@@ -1316,10 +1300,12 @@ void Optimization::addJoin(
 
   // If one is much better do not try the other.
   if (toTry.size() == 2 && candidate.tables.size() == 1) {
+    VELOX_DCHECK(!options_.syntacticJoinOrder);
     if (toTry[0].isWorse(toTry[1])) {
-      toTry.erase(toTry.begin());
+      toTry[0] = std::move(toTry[1]);
+      toTry.pop_back();
     } else if (toTry[1].isWorse(toTry[0])) {
-      toTry.erase(toTry.begin() + 1);
+      toTry.pop_back();
     }
   }
   result.insert(result.end(), toTry.begin(), toTry.end());
@@ -1348,7 +1334,7 @@ RelationOpPtr Optimization::placeSingleRowDt(
   memoKey.tables.add(subquery);
   memoKey.columns.unionObjects(subquery->columns);
 
-  const auto broadcast = Distribution::broadcast({});
+  const auto broadcast = Distribution::broadcast();
 
   PlanObjectSet empty;
   bool needsShuffle = false;
@@ -1545,11 +1531,7 @@ Distribution somePartition(const RelationOpPtrVector& inputs) {
     }
   }
 
-  DistributionType distributionType;
-  distributionType.numPartitions =
-      queryCtx()->optimization()->runnerOptions().numWorkers;
-
-  return {distributionType, std::move(columns)};
+  return {DistributionType{}, std::move(columns)};
 }
 
 // Adds the costs in the input states to the first state and if 'distinct' is
@@ -1565,9 +1547,12 @@ PlanP unionPlan(
   for (auto i = 1; i < states.size(); ++i) {
     const auto& otherCost = states[i].cost;
     fullyImported.intersect(inputPlans[i]->fullyImported);
-    firstState.cost.add(otherCost);
-    // The input cardinality is not additive, the fanout and other metrics are.
-    firstState.cost.inputCardinality -= otherCost.inputCardinality;
+    // The input cardinality is not additive.
+    firstState.cost.setupCost += otherCost.setupCost;
+    firstState.cost.unitCost += otherCost.unitCost;
+    firstState.cost.fanout += otherCost.fanout;
+    firstState.cost.totalBytes += otherCost.totalBytes;
+    firstState.cost.transferBytes += otherCost.transferBytes;
   }
   if (distinct) {
     firstState.addCost(*distinct);

axiom/optimizer/Plan.cpp

@@ -102,12 +102,6 @@ Plan::Plan(RelationOpPtr op, const PlanState& state)
       columns(exprColumns(state.targetExprs)),
       fullyImported(state.dt->fullyImported) {}
 
-bool Plan::isStateBetter(const PlanState& state, float perRowMargin) const {
-  return cost.unitCost * cost.inputCardinality + cost.setupCost >
-      state.cost.unitCost * state.cost.inputCardinality + state.cost.setupCost +
-      perRowMargin * state.cost.fanout;
-}
-
 std::string Plan::printCost() const {
   return cost.toString(true, false);
 }
@@ -120,9 +114,10 @@ std::string Plan::toString(bool detail) const {
 }
 
 void PlanState::addCost(RelationOp& op) {
-  cost.unitCost += cost.inputCardinality * cost.fanout * op.cost().unitCost;
+  VELOX_DCHECK_EQ(cost.inputCardinality, 1);
   cost.setupCost += op.cost().setupCost;
-  cost.fanout *= op.cost().fanout;
+  cost.unitCost += op.cost().unitCost * op.cost().inputCardinality;
+  cost.fanout = op.cost().resultCardinality();
   cost.totalBytes += op.cost().totalBytes;
   cost.transferBytes += op.cost().transferBytes;
 }
@@ -276,74 +271,76 @@ std::string PlanState::printPlan(RelationOpPtr op, bool detail) const {
 }
 
 PlanP PlanSet::addPlan(RelationOpPtr plan, PlanState& state) {
-  int32_t replaceIndex = -1;
-  const float shuffleCostPerRow =
-      shuffleCost(plan->columns()) * state.cost.fanout;
-
-  if (!plans.empty()) {
-    // Compare with existing. If there is one with same distribution and new is
-    // better, replace. If there is one with a different distribution and the
-    // new one can produce the same distribution by repartition, for cheaper,
-    // add the new one and delete the old one.
-    for (auto i = 0; i < plans.size(); ++i) {
-      auto old = plans[i].get();
-      if (state.input != old->input) {
-        continue;
-      }
+  const float shuffleCostPerRow = shuffleCost(plan->columns());
 
-      const bool newIsBetter = old->isStateBetter(state);
-      const bool newIsBetterWithShuffle =
-          old->isStateBetter(state, shuffleCostPerRow);
-      const bool sameDist =
-          old->op->distribution().isSamePartition(plan->distribution());
-      const bool sameOrder =
-          old->op->distribution().isSameOrder(plan->distribution());
-      if (sameDist && sameOrder) {
-        if (newIsBetter) {
-          replaceIndex = i;
-          continue;
-        }
-        // There's a better one with same dist and partition.
-        return nullptr;
-      }
+  // Determine is old plan worse the new one in all aspects.
+  auto isWorse = [&](const Plan& old) {
+    if (plan->distribution().needsSort(old.op->distribution())) {
+      // New plan needs a sort to match the old one, so cannot compare.
+      return false;
+    }
+    const bool needsShuffle =
+        plan->distribution().needsShuffle(old.op->distribution());
+    return old.cost.cost() >
+        state.cost.cost(needsShuffle ? shuffleCostPerRow : 0);
+  };
 
-      if (newIsBetterWithShuffle && old->op->distribution().orderKeys.empty()) {
-        // Old plan has no order and is worse than new plus shuffle. Can't win.
-        // Erase.
-        queryCtx()->optimization()->trace(
-            OptimizerOptions::kExceededBest,
-            state.dt->id(),
-            old->cost,
-            *old->op);
-        plans.erase(plans.begin() + i);
-        --i;
-        continue;
-      }
+  // Determine is old plan better than the new one in all aspects.
+  auto isBetter = [&](const Plan& old) {
+    if (old.op->distribution().needsSort(plan->distribution())) {
+      // Old plan needs a sort to match the new one, so cannot compare.
+      return false;
+    }
+    const bool needsShuffle =
+        old.op->distribution().needsShuffle(plan->distribution());
+    return state.cost.cost() >
+        old.cost.cost(needsShuffle ? shuffleCost(old.op->columns()) : 0);
+  };
 
-      if (plan->distribution().orderKeys.empty() &&
-          !old->isStateBetter(state, -shuffleCostPerRow)) {
-        // New has no order and old would beat it even after adding shuffle.
-        return nullptr;
-      }
+  // Compare with existing plans.
+  const auto plansSize = plans.size();
+  enum {
+    kFoundWorse = -1,
+    kNone = 0,
+    kFoundBetter = 1,
+  };
+  auto found = kNone;
+  for (size_t i = 0; i < plans.size(); ++i) {
+    const auto& old = *plans[i];
+    if (old.input != state.input) {
+      // Different plans, cannot compare.
+      continue;
     }
+    if (isWorse(old)) {
+      // Remove old plan, it is worse than the new one in all aspects.
+      queryCtx()->optimization()->trace(
+          OptimizerOptions::kExceededBest, state.dt->id(), old.cost, *old.op);
+      std::swap(plans[i], plans.back());
+      plans.pop_back();
+      --i;
+      found = kFoundWorse;
+    } else if (found == kNone && isBetter(old)) {
+      // Old plan is better than the new one in all aspects.
+      found = kFoundBetter;
+    }
+  }
+  if (found == kFoundBetter) {
+    // No existing plan was worse than the new one in all aspects,
+    // and at least one existing plan is better than the new one in all aspects.
+    // So don't add the new plan.
+    return nullptr;
   }
 
   auto newPlan = std::make_unique<Plan>(std::move(plan), state);
   auto* result = newPlan.get();
-  auto newPlanCost =
-      result->cost.unitCost + result->cost.setupCost + shuffleCostPerRow;
-  if (bestCostWithShuffle == 0 || newPlanCost < bestCostWithShuffle) {
-    bestCostWithShuffle = newPlanCost;
-  }
-  if (replaceIndex >= 0) {
-    plans[replaceIndex] = std::move(newPlan);
-  } else {
-    plans.push_back(std::move(newPlan));
-  }
+  bestCostWithShuffle =
+      std::min(bestCostWithShuffle, result->cost.cost(shuffleCostPerRow));
+  plans.push_back(std::move(newPlan));
   return result;
 }
 
-PlanP PlanSet::best(const Distribution& distribution, bool& needsShuffle) {
+PlanP PlanSet::best(const Distribution& desired, bool& needsShuffle) {
+  // TODO: Consider desired order here too.
   PlanP best = nullptr;
   PlanP match = nullptr;
   float bestCost = -1;
@@ -352,8 +349,7 @@ PlanP PlanSet::best(const Distribution& distribution, bool& needsShuffle) {
   const bool single = isSingleWorker();
 
   for (const auto& plan : plans) {
-    const float cost =
-        plan->cost.fanout * plan->cost.unitCost + plan->cost.setupCost;
+    const float cost = plan->cost.cost();
 
     auto update = [&](PlanP& current, float& currentCost) {
       if (!current || cost < currentCost) {
@@ -363,7 +359,7 @@ PlanP PlanSet::best(const Distribution& distribution, bool& needsShuffle) {
     };
 
     update(best, bestCost);
-    if (!single && plan->op->distribution().isSamePartition(distribution)) {
+    if (!single && !plan->op->distribution().needsShuffle(desired)) {
       update(match, matchCost);
     }
   }
@@ -375,8 +371,9 @@ PlanP PlanSet::best(const Distribution& distribution, bool& needsShuffle) {
   }
 
   if (match) {
-    const float shuffle = shuffleCost(best->op->columns()) * best->cost.fanout;
-    if (matchCost <= bestCost + shuffle) {
+    const float bestCostWithShuffle =
+        best->cost.cost(shuffleCost(best->op->columns()));
+    if (matchCost <= bestCostWithShuffle) {
       return match;
     }
   }
@@ -489,12 +486,14 @@ std::string JoinCandidate::toString() const {
 }
 
 bool NextJoin::isWorse(const NextJoin& other) const {
-  float shuffle =
-      plan->distribution().isSamePartition(other.plan->distribution())
-      ? 0
-      : plan->cost().fanout * shuffleCost(plan->columns());
-  return cost.unitCost + cost.setupCost + shuffle >
-      other.cost.unitCost + other.cost.setupCost;
+  if (other.plan->distribution().needsSort(plan->distribution())) {
+    // 'other' needs a sort to match 'plan', so cannot compare.
+    return false;
+  }
+  const auto needsShuffle =
+      other.plan->distribution().needsShuffle(plan->distribution());
+  return cost.cost() >
+      other.cost.cost(needsShuffle ? shuffleCost(other.plan->columns()) : 0);
 }
 
 size_t MemoKey::hash() const {

axiom/optimizer/Plan.h

@@ -36,10 +36,6 @@ using HashBuildVector = std::vector<HashBuildCP>;
 struct Plan {
   Plan(RelationOpPtr op, const PlanState& state);
 
-  /// True if 'state' has a lower cost than 'this'. If 'perRowMargin' is given,
-  /// then 'other' must win by margin per row.
-  bool isStateBetter(const PlanState& state, float perRowMargin = 0) const;
-
   /// Root of the plan tree.
   const RelationOpPtr op;
 
@@ -81,13 +77,13 @@ struct PlanSet {
   /// Interesting equivalent plans.
   std::vector<std::unique_ptr<Plan>> plans;
 
-  /// Cost of lowest-cost plan plus shuffle. If a cutoff is applicable, nothing
-  /// more expensive than this should be tried.
-  float bestCostWithShuffle{0};
+  /// Cost of lowest-cost plan plus shuffle. If a cutoff is applicable,
+  /// nothing more expensive than this should be tried.
+  float bestCostWithShuffle{std::numeric_limits<float>::infinity()};
 
-  /// Returns the best plan that produces 'distribution'. If the best plan has
-  /// some other distribution, sets 'needsShuffle ' to true.
-  PlanP best(const Distribution& distribution, bool& needsShuffle);
+  /// Returns the best plan that produces 'desired' distribution.
+  /// If the best plan has some other distribution, sets 'needsShuffle' to true.
+  PlanP best(const Distribution& desired, bool& needsShuffle);
 
   /// Retruns the best plan when we're ok with any distribution.
   PlanP best() {
@@ -275,8 +271,11 @@ struct PlanState {
   /// True if the costs accumulated so far are so high that this should not be
   /// explored further.
   bool isOverBest() const {
-    return hasCutoff && plans.bestCostWithShuffle != 0 &&
-        cost.unitCost + cost.setupCost > plans.bestCostWithShuffle;
+    /// This isn't conservative. Because it's possible that we explore some
+    /// completely new plan with non-compatible input/distribution to any old
+    /// plan. This plan if not this condition will be added to plans and later
+    /// can become part of the best plan.
+    return hasCutoff && cost.cost() > plans.bestCostWithShuffle;
   }
 
   void debugSetFirstTable(int32_t id);