[Address Comment] Add logical planner test for call distributed procedure

Author: hantangwangd

presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergLogicalPlanner.java

@@ -98,6 +98,7 @@
 import static com.facebook.presto.sql.planner.assertions.MatchResult.match;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyNot;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.callDistributedProcedure;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.exchange;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.expression;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.filter;
@@ -107,8 +108,13 @@
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.strictProject;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.strictTableScan;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.tableFinish;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values;
 import static com.facebook.presto.sql.planner.optimizations.PlanNodeSearcher.searchFrom;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Scope.LOCAL;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Scope.REMOTE_STREAMING;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Type.GATHER;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Type.REPARTITION;
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
@@ -730,6 +736,50 @@ public void testThoroughlyPushdownForTableWithUnsupportedSpecsWhoseDataAllDelete
         }
     }
 
+    @Test
+    public void testCallDistributedProcedureOnPartitionedTable()
+    {
+        String tableName = "partition_table_for_call_distributed_procedure";
+        try {
+            assertUpdate("CREATE TABLE " + tableName + " (c1 integer, c2 varchar) with (partitioning = ARRAY['c1'])");
+            assertUpdate("INSERT INTO " + tableName + " values(1, 'foo'), (2, 'bar')", 2);
+            assertUpdate("INSERT INTO " + tableName + " values(3, 'foo'), (4, 'bar')", 2);
+            assertUpdate("INSERT INTO " + tableName + " values(5, 'foo'), (6, 'bar')", 2);
+
+            assertPlan(format("CALL system.rewrite_data_files(table_name => '%s', schema => '%s')", tableName, getSession().getSchema().get()),
+                    output(tableFinish(exchange(REMOTE_STREAMING, GATHER,
+                            callDistributedProcedure(
+                                    exchange(LOCAL, GATHER,
+                                            exchange(REMOTE_STREAMING, REPARTITION,
+                                                    strictTableScan(tableName, identityMap("c1", "c2")))))))));
+
+            // Do not support the filter that couldn't be enforced totally by tableScan
+            assertQueryFails(format("CALL system.rewrite_data_files(table_name => '%s', schema => '%s', filter => 'c2 > ''bar''')", tableName, getSession().getSchema().get()),
+                    "Unexpected FilterNode found in plan; probably connector was not able to handle provided WHERE expression");
+
+            // Support the filter that could be enforced totally by tableScan
+            assertPlan(getSession(), format("CALL system.rewrite_data_files(table_name => '%s', schema => '%s', filter => 'c1 > 3')", tableName, getSession().getSchema().get()),
+                    output(tableFinish(exchange(REMOTE_STREAMING, GATHER,
+                            callDistributedProcedure(
+                                    exchange(LOCAL, GATHER,
+                                            exchange(REMOTE_STREAMING, REPARTITION,
+                                                    strictTableScan(tableName, identityMap("c1", "c2")))))))),
+                    plan -> assertTableLayout(
+                            plan,
+                            tableName,
+                            withColumnDomains(ImmutableMap.of(
+                                    new Subfield(
+                                            "c1",
+                                            ImmutableList.of()),
+                                    Domain.create(ValueSet.ofRanges(greaterThan(INTEGER, 3L)), false))),
+                            TRUE_CONSTANT,
+                            ImmutableSet.of("c1")));
+        }
+        finally {
+            assertUpdate("DROP TABLE " + tableName);
+        }
+    }
+
     @DataProvider(name = "timezones")
     public Object[][] timezones()
     {

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java

@@ -1198,8 +1198,15 @@ protected Scope visitCall(Call call, Optional<Scope> scope)
             if (analysis.isDescribe()) {
                 return createAndAssignScope(call, scope);
             }
-            QualifiedObjectName procedureName = analysis.getProcedureName()
-                    .orElse(createQualifiedObjectName(session, call, call.getName(), metadata));
+            Optional<QualifiedObjectName> procedureNameOptional = analysis.getProcedureName();
+            QualifiedObjectName procedureName;
+            if (!procedureNameOptional.isPresent()) {
+                procedureName = createQualifiedObjectName(session, call, call.getName(), metadata);
+                analysis.setProcedureName(Optional.of(procedureName));
+            }
+            else {
+                procedureName = procedureNameOptional.get();
+            }
             ConnectorId connectorId = metadata.getCatalogHandle(session, procedureName.getCatalogName())
                     .orElseThrow(() -> new SemanticException(MISSING_CATALOG, call, "Catalog %s does not exist", procedureName.getCatalogName()));
 

presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java

@@ -43,6 +43,7 @@
 import com.facebook.presto.spi.plan.SemiJoinNode;
 import com.facebook.presto.spi.plan.SortNode;
 import com.facebook.presto.spi.plan.SpatialJoinNode;
+import com.facebook.presto.spi.plan.TableFinishNode;
 import com.facebook.presto.spi.plan.TableWriterNode;
 import com.facebook.presto.spi.plan.TopNNode;
 import com.facebook.presto.spi.plan.UnionNode;
@@ -60,6 +61,7 @@
 import com.facebook.presto.sql.planner.iterative.GroupReference;
 import com.facebook.presto.sql.planner.plan.ApplyNode;
 import com.facebook.presto.sql.planner.plan.AssignUniqueId;
+import com.facebook.presto.sql.planner.plan.CallDistributedProcedureNode;
 import com.facebook.presto.sql.planner.plan.EnforceSingleRowNode;
 import com.facebook.presto.sql.planner.plan.ExchangeNode;
 import com.facebook.presto.sql.planner.plan.GroupIdNode;
@@ -691,6 +693,16 @@ public static PlanMatchPattern enforceSingleRow(PlanMatchPattern source)
         return node(EnforceSingleRowNode.class, source);
     }
 
+    public static PlanMatchPattern callDistributedProcedure(PlanMatchPattern source)
+    {
+        return node(CallDistributedProcedureNode.class, source);
+    }
+
+    public static PlanMatchPattern tableFinish(PlanMatchPattern source)
+    {
+        return node(TableFinishNode.class, source);
+    }
+
     public static PlanMatchPattern tableWriter(List<String> columns, List<String> columnNames, PlanMatchPattern source)
     {
         return node(TableWriterNode.class, source).with(new TableWriterMatcher(columns, columnNames));