Replace tableScan API with tableScanBuilder

facebookincubator · Dec 5, 2023 · 7a3f414 · 7a3f414
1 parent 975ca3a
commit 7a3f414
Show file tree

Hide file tree

Showing 7 changed files with 337 additions and 162 deletions.
diff --git a/velox/connectors/fuzzer/tests/FuzzerConnectorTest.cpp b/velox/connectors/fuzzer/tests/FuzzerConnectorTest.cpp
@@ -31,8 +31,12 @@ TEST_F(FuzzerConnectorTest, singleSplit) {
   const size_t numRows = 100;
   auto type = ROW({BIGINT(), DOUBLE(), VARCHAR()});
 
-  auto plan =
-      PlanBuilder().tableScan(type, makeFuzzerTableHandle(), {}).planNode();
+  auto plan = PlanBuilder()
+                  .startTableScan()
+                  .outputType(type)
+                  .tableHandle(makeFuzzerTableHandle())
+                  .endTableScan()
+                  .planNode();
 
   exec::test::AssertQueryBuilder(plan)
       .split(makeFuzzerSplit(numRows))
@@ -43,8 +47,12 @@ TEST_F(FuzzerConnectorTest, floatingPoints) {
   const size_t numRows = 1000;
   auto type = ROW({REAL(), DOUBLE()});
 
-  auto plan =
-      PlanBuilder().tableScan(type, makeFuzzerTableHandle(), {}).planNode();
+  auto plan = PlanBuilder()
+                  .startTableScan()
+                  .outputType(type)
+                  .tableHandle(makeFuzzerTableHandle())
+                  .endTableScan()
+                  .planNode();
 
   exec::test::AssertQueryBuilder(plan)
       .split(makeFuzzerSplit(numRows))
@@ -59,8 +67,12 @@ TEST_F(FuzzerConnectorTest, complexTypes) {
       REAL(),
   });
 
-  auto plan =
-      PlanBuilder().tableScan(type, makeFuzzerTableHandle(), {}).planNode();
+  auto plan = PlanBuilder()
+                  .startTableScan()
+                  .outputType(type)
+                  .tableHandle(makeFuzzerTableHandle())
+                  .endTableScan()
+                  .planNode();
 
   exec::test::AssertQueryBuilder(plan)
       .split(makeFuzzerSplit(numRows))
@@ -72,8 +84,12 @@ TEST_F(FuzzerConnectorTest, multipleSplits) {
   const size_t numSplits = 10;
   auto type = ROW({BIGINT(), DOUBLE(), VARCHAR()});
 
-  auto plan =
-      PlanBuilder().tableScan(type, makeFuzzerTableHandle(), {}).planNode();
+  auto plan = PlanBuilder()
+                  .startTableScan()
+                  .outputType(type)
+                  .tableHandle(makeFuzzerTableHandle())
+                  .endTableScan()
+                  .planNode();
 
   exec::test::AssertQueryBuilder(plan)
       .splits(makeFuzzerSplits(rowsPerSplit, numSplits))
@@ -89,8 +105,12 @@ TEST_F(FuzzerConnectorTest, randomTypes) {
   for (size_t i = 0; i < iterations; ++i) {
     auto type = VectorFuzzer({}, pool()).randRowType();
 
-    auto plan =
-        PlanBuilder().tableScan(type, makeFuzzerTableHandle(), {}).planNode();
+    auto plan = PlanBuilder()
+                    .startTableScan()
+                    .outputType(type)
+                    .tableHandle(makeFuzzerTableHandle())
+                    .endTableScan()
+                    .planNode();
     exec::test::AssertQueryBuilder(plan)
         .splits(makeFuzzerSplits(rowsPerSplit, numSplits))
         .assertTypeAndNumRows(type, rowsPerSplit * numSplits);
@@ -101,14 +121,18 @@ TEST_F(FuzzerConnectorTest, reproducible) {
   const size_t numRows = 100;
   auto type = ROW({BIGINT(), ARRAY(INTEGER()), VARCHAR()});
 
-  auto plan1 =
-      PlanBuilder()
-          .tableScan(type, makeFuzzerTableHandle(/*fuzerSeed=*/1234), {})
-          .planNode();
-  auto plan2 =
-      PlanBuilder()
-          .tableScan(type, makeFuzzerTableHandle(/*fuzerSeed=*/1234), {})
-          .planNode();
+  auto plan1 = PlanBuilder()
+                   .startTableScan()
+                   .outputType(type)
+                   .tableHandle(makeFuzzerTableHandle(/*fuzerSeed=*/1234))
+                   .endTableScan()
+                   .planNode();
+  auto plan2 = PlanBuilder()
+                   .startTableScan()
+                   .outputType(type)
+                   .tableHandle(makeFuzzerTableHandle(/*fuzerSeed=*/1234))
+                   .endTableScan()
+                   .planNode();
 
   auto results1 = exec::test::AssertQueryBuilder(plan1)
                       .split(makeFuzzerSplit(numRows))

diff --git a/velox/exec/tests/AssertQueryBuilderTest.cpp b/velox/exec/tests/AssertQueryBuilderTest.cpp
@@ -86,10 +86,11 @@ TEST_F(AssertQueryBuilderTest, hiveSplits) {
 
   AssertQueryBuilder(
       PlanBuilder()
-          .tableScan(
-              ROW({"c0", "ds"}, {INTEGER(), VARCHAR()}),
-              makeTableHandle(),
-              assignments)
+          .startTableScan()
+          .outputType(ROW({"c0", "ds"}, {INTEGER(), VARCHAR()}))
+          .tableHandle(makeTableHandle())
+          .assignments(assignments)
+          .endTableScan()
           .planNode(),
       duckDbQueryRunner_)
       .split(HiveConnectorSplitBuilder(file->path)

diff --git a/velox/exec/tests/AsyncConnectorTest.cpp b/velox/exec/tests/AsyncConnectorTest.cpp
@@ -193,7 +193,10 @@ TEST_F(AsyncConnectorTest, basic) {
   auto tableHandle = std::make_shared<TestTableHandle>();
   core::PlanNodeId scanId;
   auto plan = PlanBuilder()
-                  .tableScan(ROW({"a"}, {BIGINT()}), tableHandle, {})
+                  .startTableScan()
+                  .outputType(ROW({"a"}, {BIGINT()}))
+                  .tableHandle(tableHandle)
+                  .endTableScan()
                   .capturePlanNodeId(scanId)
                   .singleAggregation({}, {"min(a)"})
                   .planNode();

diff --git a/velox/exec/tests/HashJoinTest.cpp b/velox/exec/tests/HashJoinTest.cpp
@@ -4123,16 +4123,17 @@ TEST_F(HashJoinTest, dynamicFilters) {
     assignments["b"] = regularColumn("c1", BIGINT());
 
     core::PlanNodeId probeScanId;
-    auto op =
-        PlanBuilder(planNodeIdGenerator, pool_.get())
-            .tableScan(
-                scanOutputType,
-                makeTableHandle(common::test::SubfieldFiltersBuilder().build()),
-                assignments)
-            .capturePlanNodeId(probeScanId)
-            .hashJoin({"a"}, {"u_c0"}, buildSide, "", {"a", "b", "u_c1"})
-            .project({"a", "b + 1", "b + u_c1"})
-            .planNode();
+    auto op = PlanBuilder(planNodeIdGenerator, pool_.get())
+                  .startTableScan()
+                  .outputType(scanOutputType)
+                  .tableHandle(makeTableHandle(
+                      common::test::SubfieldFiltersBuilder().build()))
+                  .assignments(assignments)
+                  .endTableScan()
+                  .capturePlanNodeId(probeScanId)
+                  .hashJoin({"a"}, {"u_c0"}, buildSide, "", {"a", "b", "u_c1"})
+                  .project({"a", "b + 1", "b + u_c1"})
+                  .planNode();
 
     HashJoinBuilder(*pool_, duckDbQueryRunner_, driverExecutor_.get())
         .planNode(std::move(op))
@@ -4889,7 +4890,11 @@ TEST_F(HashJoinTest, dynamicFilterOnPartitionKey) {
   auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
   auto op =
       PlanBuilder(planNodeIdGenerator)
-          .tableScan(outputType, tableHandle, assignments)
+          .startTableScan()
+          .outputType(outputType)
+          .tableHandle(tableHandle)
+          .assignments(assignments)
+          .endTableScan()
           .capturePlanNodeId(probeScanId)
           .hashJoin(
               {"n1_1"},