Compile merge sort comparator

presto-main/src/main/java/com/facebook/presto/operator/MergeOperator.java

@@ -21,6 +21,7 @@
 import com.facebook.presto.spi.block.SortOrder;
 import com.facebook.presto.spi.type.Type;
 import com.facebook.presto.split.RemoteSplit;
+import com.facebook.presto.sql.gen.OrderingCompiler;
 import com.facebook.presto.sql.planner.plan.PlanNodeId;
 import com.google.common.io.Closer;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -56,15 +57,15 @@ public static class MergeOperatorFactory
         private final List<Type> outputTypes;
         private final List<Integer> sortChannels;
         private final List<SortOrder> sortOrder;
-        private final PageWithPositionComparatorFactory pageWithPositionComparatorFactory;
+        private final OrderingCompiler orderingCompiler;
         private boolean closed;
 
         public MergeOperatorFactory(
                 int operatorId,
                 PlanNodeId sourceId,
                 ExchangeClientSupplier exchangeClientSupplier,
                 PagesSerdeFactory serdeFactory,
-                PageWithPositionComparatorFactory pageWithPositionComparatorFactory,
+                OrderingCompiler orderingCompiler,
                 List<Type> types,
                 List<Integer> outputChannels,
                 List<Integer> sortChannels,
@@ -79,7 +80,7 @@ public MergeOperatorFactory(
             this.outputTypes = mappedCopy(outputChannels, types::get);
             this.sortChannels = requireNonNull(sortChannels, "sortChannels is null");
             this.sortOrder = requireNonNull(sortOrder, "sortOrder is null");
-            this.pageWithPositionComparatorFactory = requireNonNull(pageWithPositionComparatorFactory, "pageWithPositionComparatorFactory is null");
+            this.orderingCompiler = requireNonNull(orderingCompiler, "mergeSortComparatorFactory is null");
         }
 
         @Override
@@ -99,7 +100,7 @@ public SourceOperator createOperator(DriverContext driverContext)
                     sourceId,
                     exchangeClientSupplier,
                     serdeFactory.createPagesSerde(),
-                    pageWithPositionComparatorFactory.create(types, sortChannels, sortOrder),
+                    orderingCompiler.compilePageWithPositionComparator(types, sortChannels, sortOrder),
                     outputChannels,
                     outputTypes);
         }

presto-main/src/main/java/com/facebook/presto/operator/SimplePageWithPositionComparator.java

@@ -26,16 +26,6 @@
 public class SimplePageWithPositionComparator
         implements PageWithPositionComparator
 {
-    public static class Factory
-            implements PageWithPositionComparatorFactory
-    {
-        @Override
-        public PageWithPositionComparator create(List<Type> sortTypes, List<Integer> sortChannels, List<SortOrder> sortOrders)
-        {
-            return new SimplePageWithPositionComparator(sortTypes, sortChannels, sortOrders);
-        }
-    }
-
     private final List<Integer> sortChannels;
     private final List<SortOrder> sortOrders;
     private final List<Type> types;

presto-main/src/main/java/com/facebook/presto/operator/exchange/LocalMergeSourceOperator.java

@@ -18,12 +18,12 @@
 import com.facebook.presto.operator.OperatorContext;
 import com.facebook.presto.operator.OperatorFactory;
 import com.facebook.presto.operator.PageWithPositionComparator;
-import com.facebook.presto.operator.PageWithPositionComparatorFactory;
 import com.facebook.presto.operator.WorkProcessor;
 import com.facebook.presto.operator.exchange.LocalExchange.LocalExchangeFactory;
 import com.facebook.presto.spi.Page;
 import com.facebook.presto.spi.block.SortOrder;
 import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.sql.gen.OrderingCompiler;
 import com.facebook.presto.sql.planner.plan.PlanNodeId;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -45,27 +45,27 @@ public static class LocalMergeSourceOperatorFactory
     {
         private final int operatorId;
         private final PlanNodeId planNodeId;
-        private final PageWithPositionComparatorFactory comparatorFactory;
         private final LocalExchangeFactory localExchangeFactory;
         private final List<Type> types;
+        private final OrderingCompiler orderingCompiler;
         private final List<Integer> sortChannels;
         private final List<SortOrder> orderings;
         private boolean closed;
 
         public LocalMergeSourceOperatorFactory(
                 int operatorId,
                 PlanNodeId planNodeId,
-                PageWithPositionComparatorFactory comparatorFactory,
                 LocalExchangeFactory localExchangeFactory,
                 List<Type> types,
+                OrderingCompiler orderingCompiler,
                 List<Integer> sortChannels,
                 List<SortOrder> orderings)
         {
             this.operatorId = operatorId;
             this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
-            this.comparatorFactory = requireNonNull(comparatorFactory, "comparator is null");
             this.localExchangeFactory = requireNonNull(localExchangeFactory, "exchange is null");
             this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
+            this.orderingCompiler = requireNonNull(orderingCompiler, "orderingCompiler is null");
             this.sortChannels = ImmutableList.copyOf(requireNonNull(sortChannels, "sortChannels is null"));
             this.orderings = ImmutableList.copyOf(requireNonNull(orderings, "orderings is null"));
         }
@@ -78,8 +78,7 @@ public Operator createOperator(DriverContext driverContext)
             LocalExchange localExchange = localExchangeFactory.getLocalExchange(driverContext.getLifespan());
 
             OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, LocalMergeSourceOperator.class.getSimpleName());
-            PageWithPositionComparator comparator = comparatorFactory.create(types, sortChannels, orderings);
-
+            PageWithPositionComparator comparator = orderingCompiler.compilePageWithPositionComparator(types, sortChannels, orderings);
             List<LocalExchangeSource> sources = IntStream.range(0, localExchange.getBufferCount())
                     .boxed()
                     .map(index -> localExchange.getNextSource())

presto-main/src/main/java/com/facebook/presto/sql/gen/OrderingCompiler.java

@@ -13,11 +13,14 @@
  */
 package com.facebook.presto.sql.gen;
 
+import com.facebook.presto.operator.PageWithPositionComparator;
 import com.facebook.presto.operator.PagesIndex;
 import com.facebook.presto.operator.PagesIndexComparator;
 import com.facebook.presto.operator.PagesIndexOrdering;
+import com.facebook.presto.operator.SimplePageWithPositionComparator;
 import com.facebook.presto.operator.SimplePagesIndexComparator;
 import com.facebook.presto.operator.SyntheticAddress;
+import com.facebook.presto.spi.Page;
 import com.facebook.presto.spi.block.Block;
 import com.facebook.presto.spi.block.SortOrder;
 import com.facebook.presto.spi.type.Type;
@@ -65,13 +68,25 @@ public class OrderingCompiler
             .maximumSize(1000)
             .build(CacheLoader.from(key -> internalCompilePagesIndexOrdering(key.getSortTypes(), key.getSortChannels(), key.getSortOrders())));
 
+    private final LoadingCache<PagesIndexComparatorCacheKey, PageWithPositionComparator> pageWithPositionComparators = CacheBuilder.newBuilder()
+            .recordStats()
+            .maximumSize(1000)
+            .build(CacheLoader.from(key -> internalCompilePageWithPositionComparator(key.getSortTypes(), key.getSortChannels(), key.getSortOrders())));
+
     @Managed
     @Nested
     public CacheStatsMBean getPagesIndexOrderingsStats()
     {
         return new CacheStatsMBean(pagesIndexOrderings);
     }
 
+    @Managed
+    @Nested
+    public CacheStatsMBean getPageWithPositionsComparatorsStats()
+    {
+        return new CacheStatsMBean(pageWithPositionComparators);
+    }
+
     public PagesIndexOrdering compilePagesIndexOrdering(List<Type> sortTypes, List<Integer> sortChannels, List<SortOrder> sortOrders)
     {
         requireNonNull(sortTypes, "sortTypes is null");
@@ -93,7 +108,7 @@ public PagesIndexOrdering internalCompilePagesIndexOrdering(List<Type> sortTypes
             comparator = pagesHashStrategyClass.getConstructor().newInstance();
         }
         catch (Throwable e) {
-            log.error(e, "Error compiling comparator for channels %s with order %s", sortChannels, sortChannels);
+            log.error(e, "Error compiling comparator for channels %s with order %s", sortChannels, sortOrders);
             comparator = new SimplePagesIndexComparator(sortTypes, sortChannels, sortOrders);
         }
 
@@ -216,6 +231,96 @@ private static void generatePageIndexCompareTo(ClassDefinition classDefinition,
                 .retInt();
     }
 
+    public PageWithPositionComparator compilePageWithPositionComparator(List<Type> sortTypes, List<Integer> sortChannels, List<SortOrder> sortOrders)
+    {
+        requireNonNull(sortTypes, "sortTypes is null");
+        requireNonNull(sortChannels, "sortChannels is null");
+        requireNonNull(sortOrders, "sortOrders is null");
+
+        return pageWithPositionComparators.getUnchecked(new PagesIndexComparatorCacheKey(sortTypes, sortChannels, sortOrders));
+    }
+
+    private PageWithPositionComparator internalCompilePageWithPositionComparator(List<Type> types, List<Integer> sortChannels, List<SortOrder> sortOrders)
+    {
+        PageWithPositionComparator comparator;
+        try {
+            Class<? extends PageWithPositionComparator> pageWithPositionsComparatorClass = generatePageWithPositionComparatorClass(types, sortChannels, sortOrders);
+            comparator = pageWithPositionsComparatorClass.getConstructor().newInstance();
+        }
+        catch (Throwable t) {
+            log.error(t, "Error compiling merge sort comparator for channels %s with order %s", sortChannels, sortChannels);
+            comparator = new SimplePageWithPositionComparator(types, sortChannels, sortOrders);
+        }
+        return comparator;
+    }
+
+    private Class<? extends PageWithPositionComparator> generatePageWithPositionComparatorClass(List<Type> sortTypes, List<Integer> sortChannels, List<SortOrder> sortOrders)
+    {
+        CallSiteBinder callSiteBinder = new CallSiteBinder();
+
+        ClassDefinition classDefinition = new ClassDefinition(
+                a(PUBLIC, FINAL),
+                makeClassName("PageWithPositionComparator"),
+                type(Object.class),
+                type(PageWithPositionComparator.class));
+
+        classDefinition.declareDefaultConstructor(a(PUBLIC));
+
+        generateMergeSortCompareTo(classDefinition, callSiteBinder, sortTypes, sortChannels, sortOrders);
+
+        return defineClass(classDefinition, PageWithPositionComparator.class, callSiteBinder.getBindings(), getClass().getClassLoader());
+    }
+
+    private void generateMergeSortCompareTo(ClassDefinition classDefinition, CallSiteBinder callSiteBinder, List<Type> types, List<Integer> sortChannels, List<SortOrder> sortOrders)
+    {
+        Parameter leftPage = arg("leftPage", Page.class);
+        Parameter leftPosition = arg("leftPosition", int.class);
+        Parameter rightPage = arg("rightPage", Page.class);
+        Parameter rightPosition = arg("rightPosition", int.class);
+        MethodDefinition compareToMethod = classDefinition.declareMethod(a(PUBLIC), "compareTo", type(int.class), leftPage, leftPosition, rightPage, rightPosition);
+
+        for (int i = 0; i < sortChannels.size(); i++) {
+            int sortChannel = sortChannels.get(i);
+            SortOrder sortOrder = sortOrders.get(i);
+
+            BytecodeBlock block = new BytecodeBlock()
+                    .setDescription("compare channel " + sortChannel + " " + sortOrder);
+
+            Type sortType = types.get(sortChannel);
+
+            BytecodeExpression leftBlock = leftPage
+                    .invoke("getBlock", Block.class, constantInt(sortChannel));
+
+            BytecodeExpression rightBlock = rightPage
+                    .invoke("getBlock", Block.class, constantInt(sortChannel));
+
+            block.append(getStatic(SortOrder.class, sortOrder.name())
+                    .invoke("compareBlockValue",
+                            int.class,
+                            ImmutableList.of(Type.class, Block.class, int.class, Block.class, int.class),
+                            constantType(callSiteBinder, sortType),
+                            leftBlock,
+                            leftPosition,
+                            rightBlock,
+                            rightPosition));
+
+            LabelNode equal = new LabelNode("equal");
+            block.comment("if (compare != 0) return compare")
+                    .dup()
+                    .ifZeroGoto(equal)
+                    .retInt()
+                    .visitLabel(equal)
+                    .pop(int.class);
+
+            compareToMethod.getBody().append(block);
+        }
+
+        // values are equal
+        compareToMethod.getBody()
+                .push(0)
+                .retInt();
+    }
+
     private static final class PagesIndexComparatorCacheKey
     {
         private final List<Type> sortTypes;

presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java

@@ -63,7 +63,6 @@
 import com.facebook.presto.operator.ScanFilterAndProjectOperator.ScanFilterAndProjectOperatorFactory;
 import com.facebook.presto.operator.SetBuilderOperator.SetBuilderOperatorFactory;
 import com.facebook.presto.operator.SetBuilderOperator.SetSupplier;
-import com.facebook.presto.operator.SimplePageWithPositionComparator;
 import com.facebook.presto.operator.SourceOperatorFactory;
 import com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialIndexBuilderOperatorFactory;
 import com.facebook.presto.operator.SpatialIndexBuilderOperator.SpatialPredicate;
@@ -113,6 +112,7 @@
 import com.facebook.presto.sql.gen.JoinCompiler;
 import com.facebook.presto.sql.gen.JoinFilterFunctionCompiler;
 import com.facebook.presto.sql.gen.JoinFilterFunctionCompiler.JoinFilterFunctionFactory;
+import com.facebook.presto.sql.gen.OrderingCompiler;
 import com.facebook.presto.sql.gen.PageFunctionCompiler;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.planner.Partitioning.ArgumentBinding;
@@ -290,6 +290,7 @@ public class LocalExecutionPlanner
     private final PagesIndex.Factory pagesIndexFactory;
     private final JoinCompiler joinCompiler;
     private final LookupJoinOperators lookupJoinOperators;
+    private final OrderingCompiler orderingCompiler;
 
     @Inject
     public LocalExecutionPlanner(
@@ -314,7 +315,8 @@ public LocalExecutionPlanner(
             BlockEncodingSerde blockEncodingSerde,
             PagesIndex.Factory pagesIndexFactory,
             JoinCompiler joinCompiler,
-            LookupJoinOperators lookupJoinOperators)
+            LookupJoinOperators lookupJoinOperators,
+            OrderingCompiler orderingCompiler)
     {
         this.queryPerformanceFetcher = requireNonNull(queryPerformanceFetcher, "queryPerformanceFetcher is null");
         this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
@@ -341,6 +343,7 @@ public LocalExecutionPlanner(
         this.pagesIndexFactory = requireNonNull(pagesIndexFactory, "pagesIndexFactory is null");
         this.joinCompiler = requireNonNull(joinCompiler, "joinCompiler is null");
         this.lookupJoinOperators = requireNonNull(lookupJoinOperators, "lookupJoinOperators is null");
+        this.orderingCompiler = requireNonNull(orderingCompiler, "orderingCompiler is null");
     }
 
     public LocalExecutionPlan plan(
@@ -698,7 +701,7 @@ private PhysicalOperation createMergeSource(RemoteSourceNode node, LocalExecutio
                     node.getId(),
                     exchangeClientSupplier,
                     new PagesSerdeFactory(blockEncodingSerde, isExchangeCompressionEnabled(session)),
-                    new SimplePageWithPositionComparator.Factory(),
+                    orderingCompiler,
                     types,
                     outputChannels,
                     sortChannels,
@@ -2266,9 +2269,9 @@ private PhysicalOperation createLocalMerge(ExchangeNode node, LocalExecutionPlan
             OperatorFactory operatorFactory = new LocalMergeSourceOperatorFactory(
                     context.getNextOperatorId(),
                     node.getId(),
-                    new SimplePageWithPositionComparator.Factory(),
                     exchangeFactory,
                     types,
+                    orderingCompiler,
                     sortChannels,
                     orderings);
             return new PhysicalOperation(operatorFactory, layout, context, UNGROUPED_EXECUTION);

presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java

@@ -113,6 +113,7 @@
 import com.facebook.presto.sql.gen.ExpressionCompiler;
 import com.facebook.presto.sql.gen.JoinCompiler;
 import com.facebook.presto.sql.gen.JoinFilterFunctionCompiler;
+import com.facebook.presto.sql.gen.OrderingCompiler;
 import com.facebook.presto.sql.gen.PageFunctionCompiler;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.planner.LocalExecutionPlanner;
@@ -705,7 +706,8 @@ public List<Driver> createDrivers(Session session, @Language("SQL") String sql,
                 blockEncodingManager,
                 new PagesIndex.TestingFactory(false),
                 joinCompiler,
-                new LookupJoinOperators());
+                new LookupJoinOperators(),
+                new OrderingCompiler());
 
         // plan query
         PipelineExecutionStrategy pipelineExecutionStrategy = subplan.getFragment().getPipelineExecutionStrategy();

presto-main/src/test/java/com/facebook/presto/execution/TaskTestUtils.java

@@ -49,6 +49,7 @@
 import com.facebook.presto.sql.gen.ExpressionCompiler;
 import com.facebook.presto.sql.gen.JoinCompiler;
 import com.facebook.presto.sql.gen.JoinFilterFunctionCompiler;
+import com.facebook.presto.sql.gen.OrderingCompiler;
 import com.facebook.presto.sql.gen.PageFunctionCompiler;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.planner.LocalExecutionPlanner;
@@ -165,7 +166,8 @@ public static LocalExecutionPlanner createTestingPlanner()
                 new BlockEncodingManager(new TestingTypeManager()),
                 new PagesIndex.TestingFactory(false),
                 new JoinCompiler(MetadataManager.createTestMetadataManager(), new FeaturesConfig()),
-                new LookupJoinOperators());
+                new LookupJoinOperators(),
+                new OrderingCompiler());
     }
 
     public static TaskInfo updateTask(SqlTask sqlTask, List<TaskSource> taskSources, OutputBuffers outputBuffers)