fix: Improve sort merge join rule to do merge join when one sided is sorted (#26361)

## Description
The current
[MergeJoinForSortedInputOptimizer](https://github.com/prestodb/presto/blob/master/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/MergeJoinForSortedInputOptimizer.java)
rule does sort merge join only when both inputs are sorted. In this PR,
I made the following changes:
* Make the optimizer to work only for native execution, as java does not
have merge join support in worker side
* Have the optimizer to work for all types of joins rather than just
inner join, as velox supports all types of merge joins
* Do sort merge join if one side is sorted, and add sort on the other
side, when the query is running in presto on spark
* In `GroupedExecutionTagger`, do not fail for merge join node when
grouped execution is not available for presto on spark. This is because
presto on spark spawn as many tasks as the number of partitions of data.
When one side is bucketed, the other side even not bucketed, the join
will get as many tasks as the number of buckets, hence still equivalent
to a bucket by bucket execution

## Motivation and Context
Sort merge join for more cases

## Impact
Improve performance

## Test Plan
Unit tests, and local end to end tests

## Contributor checklist

- [ ] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [ ] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [ ] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [ ] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [ ] Adequate tests were added if applicable.
- [ ] CI passed.
- [ ] If adding new dependencies, verified they have an [OpenSSF
Scorecard](https://securityscorecards.dev/#the-checks) score of 5.0 or
higher (or obtained explicit TSC approval for lower scores).

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== RELEASE NOTES ==

General Changes
* Improve  ``MergeJoinForSortedInputOptimizer`` to do sort merge join when one side of the input is sorted
```

Author: feilong-liu

presto-hive/src/test/java/com/facebook/presto/hive/TestMergeJoinPlan.java

@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.spi.plan.JoinType;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
 import com.facebook.presto.sql.planner.assertions.PlanMatchPattern;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestQueryFramework;
@@ -56,6 +57,12 @@ protected QueryRunner createQueryRunner()
                 Optional.empty());
     }
 
+    @Override
+    protected FeaturesConfig createFeaturesConfig()
+    {
+        return new FeaturesConfig().setNativeExecutionEnabled(true);
+    }
+
     @Test
     public void testJoinType()
     {
@@ -83,19 +90,19 @@ public void testJoinType()
             assertPlan(
                     mergeJoinEnabled(),
                     "select * from test_join_customer_join_type left join test_join_order_join_type on test_join_customer_join_type.custkey = test_join_order_join_type.custkey",
-                    joinPlan("test_join_customer_join_type", "test_join_order_join_type", ImmutableList.of("custkey"), ImmutableList.of("custkey"), LEFT, false));
+                    joinPlan("test_join_customer_join_type", "test_join_order_join_type", ImmutableList.of("custkey"), ImmutableList.of("custkey"), LEFT, true));
 
             // Right join
             assertPlan(
                     mergeJoinEnabled(),
                     "select * from test_join_customer_join_type right join test_join_order_join_type on test_join_customer_join_type.custkey = test_join_order_join_type.custkey",
-                    joinPlan("test_join_customer_join_type", "test_join_order_join_type", ImmutableList.of("custkey"), ImmutableList.of("custkey"), RIGHT, false));
+                    joinPlan("test_join_customer_join_type", "test_join_order_join_type", ImmutableList.of("custkey"), ImmutableList.of("custkey"), RIGHT, true));
 
             // Outer join
             assertPlan(
                     mergeJoinEnabled(),
                     "select * from test_join_customer_join_type full join test_join_order_join_type on test_join_customer_join_type.custkey = test_join_order_join_type.custkey",
-                    joinPlan("test_join_customer_join_type", "test_join_order_join_type", ImmutableList.of("custkey"), ImmutableList.of("custkey"), FULL, false));
+                    joinPlan("test_join_customer_join_type", "test_join_order_join_type", ImmutableList.of("custkey"), ImmutableList.of("custkey"), FULL, true));
         }
         finally {
             queryRunner.execute("DROP TABLE IF EXISTS test_join_customer_join_type");

presto-hive/src/test/java/com/facebook/presto/hive/TestMergeJoinPlanPrestoOnSpark.java

@@ -39,7 +39,6 @@
 import java.util.OptionalInt;
 
 import static com.facebook.presto.SystemSessionProperties.GROUPED_EXECUTION;
-import static com.facebook.presto.SystemSessionProperties.isGroupedExecutionEnabled;
 import static com.facebook.presto.SystemSessionProperties.preferSortMergeJoin;
 import static com.facebook.presto.spi.StandardErrorCode.INVALID_PLAN_ERROR;
 import static com.facebook.presto.spi.connector.ConnectorCapabilities.SUPPORTS_PAGE_SINK_COMMIT;
@@ -58,13 +57,15 @@ class GroupedExecutionTagger
     private final Metadata metadata;
     private final NodePartitioningManager nodePartitioningManager;
     private final boolean groupedExecutionEnabled;
+    private final boolean isPrestoOnSpark;
 
-    public GroupedExecutionTagger(Session session, Metadata metadata, NodePartitioningManager nodePartitioningManager)
+    public GroupedExecutionTagger(Session session, Metadata metadata, NodePartitioningManager nodePartitioningManager, boolean groupedExecutionEnabled, boolean isPrestoOnSpark)
     {
         this.session = requireNonNull(session, "session is null");
         this.metadata = requireNonNull(metadata, "metadata is null");
         this.nodePartitioningManager = requireNonNull(nodePartitioningManager, "nodePartitioningManager is null");
-        this.groupedExecutionEnabled = isGroupedExecutionEnabled(session);
+        this.groupedExecutionEnabled = groupedExecutionEnabled;
+        this.isPrestoOnSpark = isPrestoOnSpark;
     }
 
     @Override
@@ -166,6 +167,15 @@ public GroupedExecutionTagger.GroupedExecutionProperties visitMergeJoin(MergeJoi
             // TODO: This will break the other use case for merge join operating on sorted tables, which requires grouped execution for correctness.
             return GroupedExecutionTagger.GroupedExecutionProperties.notCapable();
         }
+
+        if (isPrestoOnSpark) {
+            GroupedExecutionTagger.GroupedExecutionProperties mergeJoinLeft = node.getLeft().accept(new GroupedExecutionTagger(session, metadata, nodePartitioningManager, true, true), null);
+            GroupedExecutionTagger.GroupedExecutionProperties mergeJoinRight = node.getRight().accept(new GroupedExecutionTagger(session, metadata, nodePartitioningManager, true, true), null);
+            if (mergeJoinLeft.currentNodeCapable || mergeJoinRight.currentNodeCapable) {
+                return GroupedExecutionTagger.GroupedExecutionProperties.notCapable();
+            }
+        }
+
         throw new PrestoException(
                 INVALID_PLAN_ERROR,
                 format("When grouped execution can't be enabled, merge join plan is not valid." +

presto-main-base/src/main/java/com/facebook/presto/sql/planner/GroupedExecutionTagger.java

@@ -52,6 +52,7 @@ public class PlanFragmenter
     private final QueryManagerConfig config;
     private final PlanChecker distributedPlanChecker;
     private final PlanChecker singleNodePlanChecker;
+    private final boolean isPrestoOnSpark;
 
     @Inject
     public PlanFragmenter(Metadata metadata, NodePartitioningManager nodePartitioningManager, QueryManagerConfig queryManagerConfig, FeaturesConfig featuresConfig, PlanCheckerProviderManager planCheckerProviderManager)
@@ -61,6 +62,7 @@ public PlanFragmenter(Metadata metadata, NodePartitioningManager nodePartitionin
         this.config = requireNonNull(queryManagerConfig, "queryManagerConfig is null");
         this.distributedPlanChecker = new PlanChecker(requireNonNull(featuresConfig, "featuresConfig is null"), false, planCheckerProviderManager);
         this.singleNodePlanChecker = new PlanChecker(requireNonNull(featuresConfig, "featuresConfig is null"), true, planCheckerProviderManager);
+        this.isPrestoOnSpark = featuresConfig.isPrestoSparkExecutionEnvironment();
     }
 
     public SubPlan createSubPlans(Session session, Plan plan, boolean noExchange, PlanNodeIdAllocator idAllocator, WarningCollector warningCollector)
@@ -90,7 +92,7 @@ public SubPlan createSubPlans(Session session, Plan plan, boolean noExchange, Pl
         PlanNode root = SimplePlanRewriter.rewriteWith(fragmenter, plan.getRoot(), properties);
 
         SubPlan subPlan = fragmenter.buildRootFragment(root, properties);
-        return finalizeSubPlan(subPlan, config, metadata, nodePartitioningManager, session, noExchange, warningCollector, subPlan.getFragment().getPartitioning());
+        return finalizeSubPlan(subPlan, config, metadata, nodePartitioningManager, session, noExchange, warningCollector, subPlan.getFragment().getPartitioning(), isPrestoOnSpark);
     }
 
     private static class Fragmenter

presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanFragmenter.java

@@ -46,6 +46,7 @@
 import static com.facebook.presto.SystemSessionProperties.getExchangeMaterializationStrategy;
 import static com.facebook.presto.SystemSessionProperties.getQueryMaxStageCount;
 import static com.facebook.presto.SystemSessionProperties.isForceSingleNodeOutput;
+import static com.facebook.presto.SystemSessionProperties.isGroupedExecutionEnabled;
 import static com.facebook.presto.SystemSessionProperties.isRecoverableGroupedExecutionEnabled;
 import static com.facebook.presto.SystemSessionProperties.isSingleNodeExecutionEnabled;
 import static com.facebook.presto.spi.StandardErrorCode.QUERY_HAS_TOO_MANY_STAGES;
@@ -92,12 +93,13 @@ public static SubPlan finalizeSubPlan(
             Session session,
             boolean noExchange,
             WarningCollector warningCollector,
-            PartitioningHandle partitioningHandle)
+            PartitioningHandle partitioningHandle,
+            boolean isPrestoOnSpark)
     {
         subPlan = reassignPartitioningHandleIfNecessary(metadata, session, subPlan, partitioningHandle);
         if (!noExchange && !isSingleNodeExecutionEnabled(session)) {
             // grouped execution is not supported for SINGLE_DISTRIBUTION or SINGLE_NODE_EXECUTION_ENABLED
-            subPlan = analyzeGroupedExecution(session, subPlan, false, metadata, nodePartitioningManager);
+            subPlan = analyzeGroupedExecution(session, subPlan, false, metadata, nodePartitioningManager, isPrestoOnSpark);
         }
 
         checkState(subPlan.getFragment().getId().getId() != ROOT_FRAGMENT_ID || !isForceSingleNodeOutput(session) || subPlan.getFragment().getPartitioning().isSingleNode(), "Root of PlanFragment is not single node");
@@ -148,10 +150,10 @@ private static void sanityCheckFragmentedPlan(
 
      * TODO: We should introduce "query section" and make recoverability analysis done at query section level.
      */
-    private static SubPlan analyzeGroupedExecution(Session session, SubPlan subPlan, boolean parentContainsTableFinish, Metadata metadata, NodePartitioningManager nodePartitioningManager)
+    private static SubPlan analyzeGroupedExecution(Session session, SubPlan subPlan, boolean parentContainsTableFinish, Metadata metadata, NodePartitioningManager nodePartitioningManager, boolean isPrestoOnSpark)
     {
         PlanFragment fragment = subPlan.getFragment();
-        GroupedExecutionTagger.GroupedExecutionProperties properties = fragment.getRoot().accept(new GroupedExecutionTagger(session, metadata, nodePartitioningManager), null);
+        GroupedExecutionTagger.GroupedExecutionProperties properties = fragment.getRoot().accept(new GroupedExecutionTagger(session, metadata, nodePartitioningManager, isGroupedExecutionEnabled(session), isPrestoOnSpark), null);
         if (properties.isSubTreeUseful()) {
             boolean preferDynamic = fragment.getRemoteSourceNodes().stream().allMatch(node -> node.getExchangeType() == REPLICATE)
                     && new HashSet<>(properties.getCapableTableScanNodes()).containsAll(fragment.getTableScanSchedulingOrder());
@@ -185,7 +187,7 @@ private static SubPlan analyzeGroupedExecution(Session session, SubPlan subPlan,
         ImmutableList.Builder<SubPlan> result = ImmutableList.builder();
         boolean containsTableFinishNode = containsTableFinishNode(fragment);
         for (SubPlan child : subPlan.getChildren()) {
-            result.add(analyzeGroupedExecution(session, child, containsTableFinishNode, metadata, nodePartitioningManager));
+            result.add(analyzeGroupedExecution(session, child, containsTableFinishNode, metadata, nodePartitioningManager, isPrestoOnSpark));
         }
         return new SubPlan(fragment, result.build());
     }

presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanFragmenterUtils.java

@@ -955,7 +955,7 @@ public PlanOptimizers(
         // MergeJoinForSortedInputOptimizer can avoid the local exchange for a join operation
         // Should be placed after AddExchanges, but before AddLocalExchange
         // To replace the JoinNode to MergeJoin ahead of AddLocalExchange to avoid adding extra local exchange
-        builder.add(new MergeJoinForSortedInputOptimizer(metadata, featuresConfig.isNativeExecutionEnabled()),
+        builder.add(new MergeJoinForSortedInputOptimizer(metadata, featuresConfig.isNativeExecutionEnabled(), featuresConfig.isPrestoSparkExecutionEnvironment()),
                 new SortMergeJoinOptimizer(metadata, featuresConfig.isNativeExecutionEnabled()));
 
         // Optimizers above this don't understand local exchanges, so be careful moving this.