CTE reuse WIP

Author: kasiafi

core/trino-main/src/main/java/io/trino/execution/SqlQueryExecution.java

@@ -61,6 +61,8 @@
 import io.trino.sql.analyzer.Analysis;
 import io.trino.sql.analyzer.Analyzer;
 import io.trino.sql.analyzer.AnalyzerFactory;
+import io.trino.sql.newir.FormatOptions;
+import io.trino.sql.newir.Program;
 import io.trino.sql.planner.AdaptivePlanner;
 import io.trino.sql.planner.InputExtractor;
 import io.trino.sql.planner.LogicalPlanner;
@@ -73,6 +75,7 @@
 import io.trino.sql.planner.SplitSourceFactory;
 import io.trino.sql.planner.SubPlan;
 import io.trino.sql.planner.optimizations.AdaptivePlanOptimizer;
+import io.trino.sql.planner.optimizations.ctereuse.CteReuse;
 import io.trino.sql.planner.optimizations.PlanOptimizer;
 import io.trino.sql.planner.plan.OutputNode;
 import io.trino.sql.tree.ExplainAnalyze;
@@ -148,6 +151,7 @@ public class SqlQueryExecution
     private final EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory;
     private final TaskDescriptorStorage taskDescriptorStorage;
     private final PlanOptimizersStatsCollector planOptimizersStatsCollector;
+    private final FormatOptions formatOptions;
 
     private SqlQueryExecution(
             PreparedQuery preparedQuery,
@@ -185,7 +189,8 @@ private SqlQueryExecution(
             SqlTaskManager coordinatorTaskManager,
             ExchangeManagerRegistry exchangeManagerRegistry,
             EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory,
-            TaskDescriptorStorage taskDescriptorStorage)
+            TaskDescriptorStorage taskDescriptorStorage,
+            FormatOptions formatOptions)
     {
         try (SetThreadName _ = new SetThreadName("Query-" + stateMachine.getQueryId())) {
             this.slug = requireNonNull(slug, "slug is null");
@@ -240,6 +245,7 @@ private SqlQueryExecution(
             this.eventDrivenTaskSourceFactory = requireNonNull(eventDrivenTaskSourceFactory, "taskSourceFactory is null");
             this.taskDescriptorStorage = requireNonNull(taskDescriptorStorage, "taskDescriptorStorage is null");
             this.planOptimizersStatsCollector = requireNonNull(planOptimizersStatsCollector, "planOptimizersStatsCollector is null");
+            this.formatOptions = requireNonNull(formatOptions, "formatOptions is null");
         }
     }
 
@@ -503,6 +509,8 @@ private PlanRoot doPlanQuery(CachingTableStatsProvider tableStatsProvider)
         Plan plan = logicalPlanner.plan(analysis);
         queryPlan.set(plan);
 
+        Optional<Program> optimizedProgram = CteReuse.reuseCommonSubqueries(plan, plannerContext, getSession(), formatOptions);
+
         // fragment the plan
         SubPlan fragmentedPlan;
         try (var _ = scopedSpan(tracer, "fragment-plan")) {
@@ -809,6 +817,7 @@ public static class SqlQueryExecutionFactory
         private final ExchangeManagerRegistry exchangeManagerRegistry;
         private final EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory;
         private final TaskDescriptorStorage taskDescriptorStorage;
+        private final FormatOptions formatOptions;
 
         @Inject
         SqlQueryExecutionFactory(
@@ -841,7 +850,8 @@ public static class SqlQueryExecutionFactory
                 SqlTaskManager coordinatorTaskManager,
                 ExchangeManagerRegistry exchangeManagerRegistry,
                 EventDrivenTaskSourceFactory eventDrivenTaskSourceFactory,
-                TaskDescriptorStorage taskDescriptorStorage)
+                TaskDescriptorStorage taskDescriptorStorage,
+                FormatOptions formatOptions)
         {
             this.tracer = requireNonNull(tracer, "tracer is null");
             this.schedulerStats = requireNonNull(schedulerStats, "schedulerStats is null");
@@ -875,6 +885,7 @@ public static class SqlQueryExecutionFactory
             this.exchangeManagerRegistry = requireNonNull(exchangeManagerRegistry, "exchangeManagerRegistry is null");
             this.eventDrivenTaskSourceFactory = requireNonNull(eventDrivenTaskSourceFactory, "eventDrivenTaskSourceFactory is null");
             this.taskDescriptorStorage = requireNonNull(taskDescriptorStorage, "taskDescriptorStorage is null");
+            this.formatOptions = requireNonNull(formatOptions, "formatOptions is null");
         }
 
         @Override
@@ -925,7 +936,8 @@ public QueryExecution createQueryExecution(
                     coordinatorTaskManager,
                     exchangeManagerRegistry,
                     eventDrivenTaskSourceFactory,
-                    taskDescriptorStorage);
+                    taskDescriptorStorage,
+                    formatOptions);
         }
     }
 }

core/trino-main/src/main/java/io/trino/metadata/Metadata.java

@@ -54,6 +54,7 @@
 import io.trino.spi.connector.TableFunctionApplicationResult;
 import io.trino.spi.connector.TableScanRedirectApplicationResult;
 import io.trino.spi.connector.TopNApplicationResult;
+import io.trino.spi.connector.UnificationResult;
 import io.trino.spi.connector.WriterScalingOptions;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.expression.Constant;
@@ -569,6 +570,8 @@ Optional<TopNApplicationResult<TableHandle>> applyTopN(
 
     Optional<TableFunctionApplicationResult<TableHandle>> applyTableFunction(Session session, TableFunctionHandle handle);
 
+    Optional<UnificationResult<TableHandle>> unifyTables(Session session, TableHandle first, TableHandle second);
+
     default void validateScan(Session session, TableHandle table) {}
 
     //

core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java

@@ -88,6 +88,7 @@
 import io.trino.spi.connector.TableFunctionApplicationResult;
 import io.trino.spi.connector.TableScanRedirectApplicationResult;
 import io.trino.spi.connector.TopNApplicationResult;
+import io.trino.spi.connector.UnificationResult;
 import io.trino.spi.connector.WriterScalingOptions;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.expression.Constant;
@@ -2157,6 +2158,24 @@ public Optional<TableFunctionApplicationResult<TableHandle>> applyTableFunction(
                         result.getColumnHandles()));
     }
 
+    @Override
+    public Optional<UnificationResult<TableHandle>> unifyTables(Session session, TableHandle first, TableHandle second)
+    {
+        CatalogHandle catalogHandle = first.catalogHandle();
+        ConnectorTransactionHandle transaction = first.transaction();
+        if (!catalogHandle.equals(second.catalogHandle()) || !transaction.equals(second.transaction())) {
+            return Optional.empty();
+        }
+        ConnectorMetadata metadata = getMetadata(session, catalogHandle);
+
+        return metadata.unifyTables(session.toConnectorSession(catalogHandle), first.connectorHandle(), second.connectorHandle())
+                .map(result -> new UnificationResult<>(
+                        new TableHandle(catalogHandle, result.unifiedHandle(), transaction),
+                        result.firstCompensationFilter(),
+                        result.secondCompensationFilter(),
+                        result.enforcedProperties()));
+    }
+
     private void verifyProjection(TableHandle table, List<ConnectorExpression> projections, List<Assignment> assignments, int expectedProjectionSize)
     {
         projections.forEach(projection -> requireNonNull(projection, "one of the projections is null"));

core/trino-main/src/main/java/io/trino/sql/dialect/trino/Attributes.java

@@ -40,6 +40,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import static io.trino.spi.StandardErrorCode.IR_ERROR;
 import static io.trino.sql.dialect.trino.TrinoDialect.TRINO;

core/trino-main/src/main/java/io/trino/sql/dialect/trino/ProgramBuilder.java

@@ -79,6 +79,16 @@ public static class ValueNameAllocator
     {
         private int label;
 
+        public ValueNameAllocator()
+        {
+            this(0);
+        }
+
+        public ValueNameAllocator(int initialLabel)
+        {
+            this.label = initialLabel;
+        }
+
         public String newName()
         {
             return "%" + label++;

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/AggregateCall.java

@@ -23,7 +23,6 @@
 import io.trino.sql.dialect.trino.Attributes.SortOrderList;
 import io.trino.sql.newir.Block;
 import io.trino.sql.newir.FormatOptions;
-import io.trino.sql.newir.Operation;
 import io.trino.sql.newir.Region;
 import io.trino.sql.newir.Value;
 import io.trino.type.FunctionType;
@@ -51,7 +50,7 @@
 import static java.util.Objects.requireNonNull;
 
 public class AggregateCall
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "aggregate_call";
 

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/Aggregation.java

@@ -48,7 +48,7 @@
 import static java.util.Objects.requireNonNull;
 
 public class Aggregation
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "aggregation";
 
@@ -152,6 +152,18 @@ public Aggregation(
         this.attributes = attributes.buildOrThrow();
     }
 
+    // TODO checks
+    private Aggregation(Result result, Value input, Region aggregateCalls, Region groupingKeysSelector, Region hashSelector, Map<AttributeKey, Object> attributes)
+    {
+        super(TRINO, NAME);
+        this.result = result;
+        this.input = input;
+        this.aggregateCalls = aggregateCalls;
+        this.groupingKeysSelector = groupingKeysSelector;
+        this.hashSelector = hashSelector;
+        this.attributes = ImmutableMap.copyOf(attributes);
+    }
+
     @Override
     public Result result()
     {
@@ -182,6 +194,15 @@ public String prettyPrint(int indentLevel, FormatOptions formatOptions)
         return "pretty aggregation";
     }
 
+    @Override
+    public Operation withArgument(Value newArgument, int index)
+    {
+        // TODO this does not validate the new Regions. Should run the same checks as the constructor.
+        // TODO missing real source attributes - it's incorrect to copy source attributes if the source changed
+        validateArgument(newArgument, index);
+        return new Aggregation(result, newArgument, aggregateCalls, groupingKeysSelector, this.hashSelector, this.attributes);
+    }
+
     @Override
     public boolean equals(Object obj)
     {

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/Comparison.java

@@ -17,7 +17,6 @@
 import io.trino.spi.TrinoException;
 import io.trino.sql.dialect.trino.Attributes.ComparisonOperator;
 import io.trino.sql.newir.FormatOptions;
-import io.trino.sql.newir.Operation;
 import io.trino.sql.newir.Region;
 import io.trino.sql.newir.Value;
 
@@ -35,7 +34,7 @@
 import static java.util.Objects.requireNonNull;
 
 public final class Comparison
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "comparison";
 

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/Constant.java

@@ -17,7 +17,6 @@
 import io.trino.spi.predicate.NullableValue;
 import io.trino.spi.type.Type;
 import io.trino.sql.newir.FormatOptions;
-import io.trino.sql.newir.Operation;
 import io.trino.sql.newir.Region;
 import io.trino.sql.newir.Value;
 
@@ -31,7 +30,7 @@
 import static java.util.Objects.requireNonNull;
 
 public final class Constant
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "constant";
 

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/Exchange.java

@@ -62,7 +62,7 @@
 import static java.util.Objects.requireNonNull;
 
 public class Exchange
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "exchange";
 
@@ -210,6 +210,26 @@ public Exchange(
         this.attributes = attributes.buildOrThrow();
     }
 
+    private Exchange(
+            Result result,
+            List<Value> inputs,
+            List<Region> inputFieldSelectors,
+            Region partitioningBoundArguments,
+            Region partitioningHashSelector,
+            Region orderingSelector,
+            Map<AttributeKey, Object> attributes)
+    {
+        // TODO checks!!!
+        super(TRINO, NAME);
+        this.result = result;
+        this.inputs = inputs;
+        this.inputFieldSelectors = inputFieldSelectors;
+        this.partitioningBoundArguments = partitioningBoundArguments;
+        this.partitioningHashSelector = partitioningHashSelector;
+        this.orderingSelector = orderingSelector;
+        this.attributes = attributes;
+    }
+
     @Override
     public Result result()
     {
@@ -245,6 +265,24 @@ public String prettyPrint(int indentLevel, FormatOptions formatOptions)
         return "pretty exchange";
     }
 
+    @Override
+    public Operation withArgument(Value newArgument, int index)
+    {
+        validateArgument(newArgument, index);
+        ImmutableList.Builder<Value> newInputs = ImmutableList.builder();
+        for (int i = 0; i < inputs.size(); i++) {
+            if (i == index) {
+                newInputs.add(newArgument);
+            }
+            else {
+                newInputs.add(inputs.get(i));
+            }
+        }
+        // TODO add checks in the copy constructor
+        // TODO missing real source attributes - it's incorrect to copy source attributes if the source changed
+        return new Exchange(result, newInputs.build(), inputFieldSelectors, partitioningBoundArguments, partitioningHashSelector, orderingSelector, attributes);
+    }
+
     @Override
     public boolean equals(Object obj)
     {

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/FieldReference.java

@@ -14,6 +14,7 @@
 package io.trino.sql.dialect.trino.operation;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import io.trino.spi.TrinoException;
 import io.trino.spi.type.RowType;
 import io.trino.sql.newir.FormatOptions;
@@ -34,7 +35,7 @@
 import static java.util.Objects.requireNonNull;
 
 public final class FieldReference
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "field_reference";
 
@@ -96,6 +97,19 @@ public String prettyPrint(int indentLevel, FormatOptions formatOptions)
         return "pretty field reference";
     }
 
+    @Override
+    public Operation withArgument(Value newArgument, int index)
+    {
+        validateArgument(newArgument, index);
+        // TODO missing source attributes
+        return new FieldReference(result.name(), newArgument, FIELD_INDEX.getAttribute(attributes), ImmutableMap.of());
+    }
+
+    public Value base()
+    {
+        return base;
+    }
+
     @Override
     public boolean equals(Object obj)
     {

core/trino-main/src/main/java/io/trino/sql/dialect/trino/operation/Filter.java

@@ -36,7 +36,7 @@
 import static java.util.Objects.requireNonNull;
 
 public final class Filter
-        extends Operation
+        extends TrinoOperation
 {
     private static final String NAME = "filter";
 
@@ -61,6 +61,7 @@ public Filter(String resultName, Value input, Block predicate, Map<AttributeKey,
 
         this.input = input;
 
+        // TODO validate block labels
         if (predicate.parameters().size() != 1 ||
                 !trinoType(predicate.parameters().getFirst().type()).equals(relationRowType(trinoType(input.type()))) ||
                 !trinoType(predicate.getReturnedType()).equals(BOOLEAN)) {
@@ -102,6 +103,24 @@ public String prettyPrint(int indentLevel, FormatOptions formatOptions)
         return "pretty filter";
     }
 
+    @Override
+    public Operation withArgument(Value newArgument, int index)
+    {
+        validateArgument(newArgument, index);
+        // TODO missing source attributes
+        return new Filter(result.name(), newArgument, predicate.getOnlyBlock(), ImmutableMap.of());
+    }
+
+    public Value argument()
+    {
+        return input;
+    }
+
+    public Block predicate()
+    {
+        return predicate.getOnlyBlock();
+    }
+
     @Override
     public boolean equals(Object obj)
     {