apply review suggestion and made some small readability improvements

Author: MBkkt

axiom/optimizer/ToGraph.cpp

@@ -1622,107 +1622,100 @@ void ToGraph::makeUnionDistributionAndStats(
   }
 }
 
-void ToGraph::translateUnion(const lp::SetNode& set) {
-  auto* const setDt = currentDt_;
-  auto initialRenames = std::move(renames_);
-  QGVector<DerivedTableP> children;
-  bool isLeftLeaf = true;
-  const auto topSetOp = set.operation();
+void ToGraph::translateUnionInput(
+    const folly::F14FastMap<std::string, ExprCP>& renames,
+    const lp::LogicalPlanNode& input,
+    bool& isLeftLeaf) {
+  renames_ = renames;
 
-  auto isUnionLike =
+  auto* setDt = currentDt_;
+
+  auto maybeFlatten =
       [&](const lp::LogicalPlanNode& node) -> const lp::SetNode* {
-    if (node.kind() == lp::NodeKind::kSet) {
-      const auto* set = node.asUnchecked<lp::SetNode>();
-      if (topSetOp == set->operation()) {
-        // Same set operation can be flattened.
-        return set;
-      }
-      if (topSetOp == lp::SetOperation::kUnion &&
-          set->operation() == lp::SetOperation::kUnionAll) {
-        // UNION ALL can be flattened into UNION.
-        return set;
-      }
+    if (node.kind() != lp::NodeKind::kSet) {
+      return nullptr;
+    }
+    const auto* set = node.asUnchecked<lp::SetNode>();
+    if (setDt->setOp == set->operation()) {
+      // Same set operation can be flattened.
+      return set;
+    }
+    if (setDt->setOp == lp::SetOperation::kUnion &&
+        set->operation() == lp::SetOperation::kUnionAll) {
+      // UNION ALL can be flattened into UNION.
+      return set;
     }
-
     return nullptr;
   };
+  if (const auto* setNode = maybeFlatten(input)) {
+    for (const auto& child : setNode->inputs()) {
+      translateUnionInput(renames, *child, isLeftLeaf);
+    }
+  } else {
+    currentDt_ = newDt();
+    auto* queryDt = makeUnordered(input, kAllAllowedInDt);
+    VELOX_DCHECK_NULL(queryDt);
+    auto* newDt = currentDt_;
+    currentDt_ = setDt;
+
+    const auto& type = input.outputType();
 
-  // TODO: use deducing this lambda when C++23 is available.
-  std::function<void(const lp::LogicalPlanNode&)> addChild;
+    if (isLeftLeaf) {
+      // This is the left leaf of a union tree.
+      for (auto i : usedChannels(input)) {
+        const auto& name = type->nameOf(i);
 
-  addChild = [&](const lp::LogicalPlanNode& input) {
-    renames_ = initialRenames;
+        ExprCP inner = translateColumn(name);
+        newDt->exprs.push_back(inner);
 
-    if (auto* setNode = isUnionLike(input)) {
-      for (auto& child : setNode->inputs()) {
-        addChild(*child);
+        // The top dt has the same columns as all the unioned dts.
+        const auto* columnName = toName(name);
+        auto* outer =
+            make<Column>(columnName, setDt, inner->value(), columnName);
+        setDt->columns.push_back(outer);
+        newDt->columns.push_back(outer);
       }
+      isLeftLeaf = false;
     } else {
-      currentDt_ = newDt();
-      auto* queryDt = makeUnordered(input, kAllAllowedInDt);
-      VELOX_DCHECK_NULL(queryDt);
-      auto* newDt = currentDt_;
-
-      const auto& type = input.outputType();
-
-      if (isLeftLeaf) {
-        // This is the left leaf of a union tree.
-        for (auto i : usedChannels(input)) {
-          const auto& name = type->nameOf(i);
-
-          ExprCP inner = translateColumn(name);
-          newDt->exprs.push_back(inner);
-
-          // The top dt has the same columns as all the unioned dts.
-          const auto* columnName = toName(name);
-          auto* outer =
-              make<Column>(columnName, setDt, inner->value(), columnName);
-          setDt->columns.push_back(outer);
-          newDt->columns.push_back(outer);
-        }
-        isLeftLeaf = false;
-      } else {
-        for (auto i : usedChannels(input)) {
-          ExprCP inner = translateColumn(type->nameOf(i));
-          newDt->exprs.push_back(inner);
-        }
-
-        // Same outward facing columns as the top dt of union.
-        newDt->columns = setDt->columns;
+      for (auto i : usedChannels(input)) {
+        ExprCP inner = translateColumn(type->nameOf(i));
+        newDt->exprs.push_back(inner);
       }
 
-      newDt->makeInitialPlan();
-      children.push_back(newDt);
+      // Same outward facing columns as the top dt of union.
+      newDt->columns = setDt->columns;
     }
-  };
 
-  addChild(set);
-  currentDt_ = setDt;
+    newDt->makeInitialPlan();
+    setDt->children.push_back(newDt);
+  }
+}
 
-  setDt->children = std::move(children);
-  setDt->setOp = set.operation();
+void ToGraph::translateUnion(const lp::SetNode& set) {
+  auto renames = std::move(renames_);
 
+  auto* setDt = currentDt_;
+  setDt->setOp = set.operation();
+  bool isLeftLeaf = true;
+  translateUnionInput(renames, set, isLeftLeaf);
   makeUnionDistributionAndStats(setDt);
 
-  renames_ = std::move(initialRenames);
+  renames_ = std::move(renames);
   for (const auto* column : setDt->columns) {
     renames_[column->name()] = column;
   }
 }
 
 DerivedTableP ToGraph::makeQueryGraph(const lp::LogicalPlanNode& logicalPlan) {
   markAllSubfields(logicalPlan);
-
-  currentDt_ = newDt();
-  auto* queryDt = makeQueryGraph(logicalPlan, kAllAllowedInDt);
-  VELOX_DCHECK_NULL(queryDt);
+  wrapInDt(logicalPlan);
   return currentDt_;
 }
 
 DerivedTableP ToGraph::makeUnordered(
-    const lp::LogicalPlanNode& input,
+    const lp::LogicalPlanNode& node,
     uint64_t allowedInDt) {
-  auto* outerDt = makeQueryGraph(input, allowedInDt);
+  auto* outerDt = makeQueryGraph(node, allowedInDt);
   if (currentDt_->hasOrderBy() && !currentDt_->hasLimit()) {
     currentDt_->orderKeys.clear();
     currentDt_->orderTypes.clear();
@@ -1731,24 +1724,28 @@ DerivedTableP ToGraph::makeUnordered(
 }
 
 DerivedTableP ToGraph::makeStream(
-    const lp::LogicalPlanNode& input,
+    const lp::LogicalPlanNode& node,
     uint64_t allowedInDt) {
-  auto* outerDt = makeQueryGraph(input, allowedInDt);
+  auto* outerDt = makeQueryGraph(node, allowedInDt);
   if (currentDt_->hasLimit()) {
-    finalizeDt(input, outerDt);
+    finalizeDt(node, outerDt);
     return nullptr;
   }
   return outerDt;
 }
 
+void ToGraph::wrapInDt(const lp::LogicalPlanNode& node) {
+  currentDt_ = newDt();
+  auto* queryDt = makeQueryGraph(node, kAllAllowedInDt);
+  VELOX_DCHECK_NULL(queryDt);
+}
+
 DerivedTableP ToGraph::makeQueryGraph(
     const lp::LogicalPlanNode& node,
     uint64_t allowedInDt) {
   if (!contains(allowedInDt, node.kind())) {
     auto* outerDt = currentDt_;
-    currentDt_ = newDt();
-    auto* queryDt = makeQueryGraph(node, kAllAllowedInDt);
-    VELOX_DCHECK_NULL(queryDt);
+    wrapInDt(node);
     return outerDt;
   }
 
@@ -1793,20 +1790,22 @@ DerivedTableP ToGraph::makeQueryGraph(
     }
     case lp::NodeKind::kJoin: {
       const auto& join = *node.asUnchecked<lp::JoinNode>();
+      const auto& left = *join.left();
+      const auto& right = *join.right();
       // TODO Allow mixing Unnest with Join in a single DT.
       // https://github.com/facebookincubator/axiom/issues/286
       allowedInDt = makeDtIf(allowedInDt, lp::NodeKind::kUnnest);
       allowedInDt = makeDtIf(allowedInDt, lp::NodeKind::kAggregate);
       allowedInDt = makeDtIf(allowedInDt, lp::NodeKind::kLimit);
-      if (auto* outerDt = makeUnordered(*join.left(), allowedInDt)) {
-        finalizeDt(*join.left(), outerDt);
+      if (auto* outerDt = makeUnordered(left, allowedInDt)) {
+        finalizeDt(left, outerDt);
       }
       if (join.joinType() != lp::JoinType::kInner ||
           queryCtx()->optimization()->options().syntacticJoinOrder) {
         allowedInDt = makeDtIf(allowedInDt, lp::NodeKind::kJoin);
       }
-      if (auto* outerDt = makeUnordered(*join.right(), allowedInDt)) {
-        finalizeDt(*join.right(), outerDt);
+      if (auto* outerDt = makeUnordered(right, allowedInDt)) {
+        finalizeDt(right, outerDt);
       }
       translateJoin(join);
       return nullptr;

axiom/optimizer/ToGraph.h

@@ -200,14 +200,33 @@ class ToGraph {
   // to the left.
   void canonicalizeCall(Name& name, ExprVector& args);
 
-  DerivedTableP makeStream(
-      const logical_plan::LogicalPlanNode& input,
+  // Helper function to call makeQueryGraph
+  // but removes order by from currentDt_.
+  // This is useful for cases where we want to ignore the order
+  // Example: set operations, join, write, aggregation inputs.
+  DerivedTableP makeUnordered(
+      const logical_plan::LogicalPlanNode& node,
       uint64_t allowedInDt);
 
-  DerivedTableP makeUnordered(
-      const logical_plan::LogicalPlanNode& input,
+  // Helper function to call makeQueryGraph but finishes currentDt_
+  // if it will have limit.
+  // This is useful for cases where we want to apply limit before our node.
+  // Example: filter and sort input.
+  DerivedTableP makeStream(
+      const logical_plan::LogicalPlanNode& node,
       uint64_t allowedInDt);
 
+  // Helper function to create new DerivedTable
+  // and call makeQueryGraph to fill it.
+  void wrapInDt(const logical_plan::LogicalPlanNode& node);
+
+  // Fill currentDt_ from 'node'.
+  // If 'allowedInDt' does not allow 'node', wraps in a new DerivedTable.
+  // Returns nullptr if 'node' is fully processed in currentDt_.
+  // Otherwise returns the "outer" DerivedTable, which needs to be
+  // finalized by the caller.
+  // This is useful for cases where we want to control from caller
+  // do we need to drop order by/limit/aggregation or not.
   DerivedTableP makeQueryGraph(
       const logical_plan::LogicalPlanNode& node,
       uint64_t allowedInDt);
@@ -291,6 +310,11 @@ class ToGraph {
       DerivedTableP setDt,
       DerivedTableP innerDt = nullptr);
 
+  void translateUnionInput(
+      const folly::F14FastMap<std::string, ExprCP>& renames,
+      const logical_plan::LogicalPlanNode& input,
+      bool& isLeftLeaf);
+
   void translateUnion(const logical_plan::SetNode& set);
 
   DerivedTableP translateUnnest(