Small refactoring for better readability based on disscussions with Andrey and Pavel

Author: MBkkt

axiom/optimizer/Optimization.cpp

@@ -477,9 +477,9 @@ bool isIndexColocated(
     const IndexInfo& info,
     const ExprVector& lookupValues,
     const RelationOpPtr& input) {
-  const auto& distribution = info.index->distribution;
+  const auto& desired = info.index->distribution;
   if (const auto needsShuffle =
-          input->distribution().maybeNeedsShuffle(distribution)) {
+          input->distribution().maybeNeedsShuffle(desired)) {
     return !*needsShuffle;
   }
   // TODO: Future partition type should be accounted.
@@ -490,8 +490,7 @@ bool isIndexColocated(
   for (size_t i = 0; i < input->distribution().partition.size(); ++i) {
     auto nthKey = position(lookupValues, *input->distribution().partition[i]);
     if (nthKey == kNotFound ||
-        info.schemaColumn(info.lookupKeys[nthKey]) !=
-            distribution.partition[i]) {
+        info.schemaColumn(info.lookupKeys[nthKey]) != desired.partition[i]) {
       return false;
     }
   }
@@ -600,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
@@ -1303,7 +1302,7 @@ void Optimization::addJoin(
   if (toTry.size() == 2 && candidate.tables.size() == 1) {
     VELOX_DCHECK(!options_.syntacticJoinOrder);
     if (toTry[0].isWorse(toTry[1])) {
-      std::swap(toTry[0], toTry[1]);
+      toTry[0] = std::move(toTry[1]);
       toTry.pop_back();
     } else if (toTry[1].isWorse(toTry[0])) {
       toTry.pop_back();
@@ -1548,9 +1547,11 @@ 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.totalBytes += otherCost.totalBytes;
+    firstState.cost.transferBytes += otherCost.transferBytes;
   }
   if (distinct) {
     firstState.addCost(*distinct);

axiom/optimizer/Plan.cpp

@@ -114,8 +114,9 @@ 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.unitCost += op.cost().unitCost * cost.fanout;
   cost.fanout *= op.cost().fanout;
   cost.totalBytes += op.cost().totalBytes;
   cost.transferBytes += op.cost().transferBytes;
@@ -339,6 +340,7 @@ PlanP PlanSet::addPlan(RelationOpPtr plan, PlanState& state) {
 }
 
 PlanP PlanSet::best(const Distribution& desired, bool& needsShuffle) {
+  // TODO: Consider desired order here too.
   PlanP best = nullptr;
   PlanP match = nullptr;
   float bestCost = -1;
@@ -347,8 +349,7 @@ PlanP PlanSet::best(const Distribution& desired, 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) {
@@ -370,8 +371,9 @@ PlanP PlanSet::best(const Distribution& desired, 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;
     }
   }

axiom/optimizer/RelationOp.cpp

@@ -24,13 +24,6 @@
 #include "velox/expression/ScopedVarSetter.h"
 
 namespace facebook::axiom::optimizer {
-
-void Cost::add(const Cost& other) {
-  inputCardinality += other.inputCardinality;
-  fanout += other.fanout;
-  setupCost += other.setupCost;
-}
-
 namespace {
 
 const auto& relTypeNames() {
@@ -216,8 +209,7 @@ void RelationOp::printCost(bool detail, std::stringstream& out) const {
   auto ctx = queryCtx();
   if (ctx && ctx->contextPlan()) {
     auto plan = ctx->contextPlan();
-    auto totalCost = plan->cost.unitCost + plan->cost.setupCost;
-    auto pct = 100 * cost_.inputCardinality * cost_.unitCost / totalCost;
+    auto pct = 100 * cost_.cost() / plan->cost.cost();
     out << " " << std::fixed << std::setprecision(2) << pct << "% ";
   }
   if (detail) {
@@ -353,7 +345,7 @@ Join::Join(
   cost_.inputCardinality = inputCardinality();
   cost_.fanout = fanout;
 
-  const float buildSize = right->cost().inputCardinality;
+  const float buildSize = right->resultCardinality();
   const auto numRightColumns =
       static_cast<float>(right->input()->columns().size());
   auto rowCost = numRightColumns * Costs::kHashExtractColumnCost;
@@ -535,7 +527,7 @@ Aggregation::Aggregation(
   // being unique after n values is 1 - (1/d)^n.
   auto nOut = cardinality -
       cardinality *
-          std::pow(1.0F - (1.0F / cardinality), input_->resultCardinality());
+          std::pow(1.0F - (1.0F / cardinality), cost_.inputCardinality);
 
   cost_.fanout = nOut / cost_.inputCardinality;
   const auto numGrouppingKeys = static_cast<float>(groupingKeys.size());
@@ -802,8 +794,7 @@ UnionAll::UnionAll(RelationOpPtrVector inputsVector)
     : RelationOp{RelType::kUnionAll, nullptr, Distribution{}, inputsVector[0]->columns()},
       inputs{std::move(inputsVector)} {
   for (auto& input : inputs) {
-    cost_.inputCardinality +=
-        input->cost().inputCardinality * input->cost().fanout;
+    cost_.inputCardinality += input->resultCardinality();
   }
 
   cost_.fanout = 1;

axiom/optimizer/RelationOp.h

@@ -54,17 +54,21 @@ struct Cost {
   // or index access.
   float setupCost{0};
 
-  /// Total cost of the operation.
+  /// Total cost of the a RelationOp or Plan.
   float cost() const {
-    return setupCost + inputCardinality * unitCost;
+    return setupCost + unitCost * inputCardinality;
   }
 
-  /// Total cost including cost of shuffling the output rows.
+  /// Returns the estimated number of result rows.
+  float resultCardinality() const {
+    return fanout * inputCardinality;
+  }
+
+  /// Total cost a RelationOp or Plan plus cost of shuffling the result rows.
   /// TODO: I'm not sure our shuffle cost model is right.
   /// Because broadcast vs gather vs partitioned cost are same now.
   float cost(float shuffleCostPerRow) const {
-    const auto outputCardinality = inputCardinality * fanout;
-    return cost() + outputCardinality * shuffleCostPerRow;
+    return cost() + shuffleCostPerRow * resultCardinality();
   }
 
   // Estimate of total data volume  for a hash join build or group/order
@@ -79,8 +83,6 @@ struct Cost {
   // amount of spill is 'totalBytes' - 'peakResidentBytes'.
   float peakResidentBytes{0};
 
-  void add(const Cost& other);
-
   /// If 'isUnit' shows the cost/cardinality for one row, else for
   /// 'inputCardinality' rows.
   std::string toString(bool detail, bool isUnit = false) const;
@@ -204,7 +206,7 @@ class RelationOp {
 
   /// Returns the number of output rows.
   float resultCardinality() const {
-    return cost_.inputCardinality * cost_.fanout;
+    return cost_.resultCardinality();
   }
 
   float inputCardinality() const {

axiom/optimizer/Schema.cpp

@@ -371,7 +371,7 @@ std::optional<bool> Distribution::maybeNeedsShuffle(
 }
 
 bool Distribution::needsShuffle(const Distribution& desired) const {
-  if (auto needsShuffle = maybeNeedsShuffle(desired)) {
+  if (const auto needsShuffle = maybeNeedsShuffle(desired)) {
     return *needsShuffle;
   }
   // TODO: Future partition type should decide this.

axiom/optimizer/ToVelox.cpp

@@ -1488,8 +1488,7 @@ void ToVelox::makePredictionAndHistory(
     const velox::core::PlanNodeId& id,
     const RelationOp* op) {
   nodeHistory_[id] = op->historyKey();
-  prediction_[id] = NodePrediction{
-      .cardinality = op->cost().inputCardinality * op->cost().fanout};
+  prediction_[id] = NodePrediction{.cardinality = op->resultCardinality()};
 }
 
 velox::core::PlanNodePtr ToVelox::makeFragment(