Add some extra PlanNode functionality.

velox/core/PlanNode.h

@@ -135,6 +135,25 @@ class PlanNode {
   /// The name of the plan node, used in toString.
   virtual std::string_view name() const = 0;
 
+  /// Recursively checks the node tree for a first node that satisfy a given
+  /// condition. Returns pointer to the node if found, nullptr if not.
+  static const PlanNode* visit(
+      const PlanNode* node,
+      const std::function<bool(const PlanNode* node)>& predicate) {
+    if (predicate(node)) {
+      return node;
+    }
+
+    // Recursively go further through the sources.
+    for (const auto& source : node->sources()) {
+      const auto* ret = PlanNode::visit(source.get(), predicate);
+      if (ret != nullptr) {
+        return ret;
+      }
+    }
+    return nullptr;
+  }
+
  private:
   /// The details of the plan node in textual format.
   virtual void addDetails(std::stringstream& stream) const = 0;
@@ -516,6 +535,14 @@ class AggregationNode : public PlanNode {
     return "Aggregation";
   }
 
+  bool isFinal() const {
+    return step_ == Step::kFinal;
+  }
+
+  bool isSingle() const {
+    return step_ == Step::kSingle;
+  }
+
  private:
   void addDetails(std::stringstream& stream) const override;
 

velox/core/tests/CMakeLists.txt

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-add_executable(velox_core_test TestQueryConfig.cpp TestString.cpp
-                               TestTypeAnalysis.cpp)
+add_executable(velox_core_test TestPlanNode.cpp TestQueryConfig.cpp
+                               TestString.cpp TestTypeAnalysis.cpp)
 
 add_test(velox_core_test velox_core_test)
 

velox/core/tests/TestPlanNode.cpp

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <gtest/gtest.h>
+
+#include "velox/core/PlanNode.h"
+
+using namespace ::facebook::velox;
+using namespace ::facebook::velox::core;
+
+TEST(TestPlanNode, visit) {
+  std::shared_ptr<const Type> bigIntType =
+      std::make_shared<ScalarType<TypeKind::BIGINT>>();
+  auto rowType = std::make_shared<RowType>(
+      std::vector<std::string>{"name1"},
+      std::vector<std::shared_ptr<const Type>>{bigIntType});
+
+  std::shared_ptr<connector::ConnectorTableHandle> tableHandle;
+  std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
+      assignments;
+
+  std::shared_ptr<PlanNode> tableScan3 =
+      std::make_shared<TableScanNode>("3", rowType, tableHandle, assignments);
+  std::shared_ptr<PlanNode> tableScan2 =
+      std::make_shared<TableScanNode>("2", rowType, tableHandle, assignments);
+
+  std::vector<FieldAccessTypedExprPtr> sortingKeys;
+  std::vector<SortOrder> sortingOrders;
+  std::shared_ptr<PlanNode> localMerge1 = std::make_shared<LocalMergeNode>(
+      "1",
+      sortingKeys,
+      sortingOrders,
+      std::vector<PlanNodePtr>{tableScan2, tableScan3});
+
+  std::vector<std::string> names;
+  std::vector<TypedExprPtr> projections;
+  std::shared_ptr<PlanNode> project0 =
+      std::make_shared<ProjectNode>("0", names, projections, localMerge1);
+
+  EXPECT_EQ(
+      tableScan3.get(),
+      PlanNode::visit(project0.get(), [](const PlanNode* node) {
+        return node->id() == "3";
+      }));
+
+  EXPECT_EQ(
+      project0.get(), PlanNode::visit(project0.get(), [](const PlanNode* node) {
+        return node->name() == "Project";
+      }));
+
+  EXPECT_EQ(nullptr, PlanNode::visit(project0.get(), [](const PlanNode* node) {
+              return node->name() == "Unknown";
+            }));
+}