Include already collected stats in QueryCompletionEvent

core/trino-main/src/main/java/io/trino/cost/CachingTableStatsProvider.java

@@ -13,6 +13,7 @@
  */
 package io.trino.cost;
 
+import com.google.common.collect.ImmutableMap;
 import io.trino.Session;
 import io.trino.metadata.Metadata;
 import io.trino.metadata.TableHandle;
@@ -47,4 +48,9 @@ public TableStatistics getTableStatistics(TableHandle tableHandle)
         }
         return stats;
     }
+
+    public Map<TableHandle, TableStatistics> getCachedTableStatistics()
+    {
+        return ImmutableMap.copyOf(cache);
+    }
 }

core/trino-main/src/main/java/io/trino/cost/ValuesStatsRule.java

@@ -68,16 +68,22 @@ public Optional<PlanNodeStatsEstimate> calculate(ValuesNode node, StatsProvider
         PlanNodeStatsEstimate.Builder statsBuilder = PlanNodeStatsEstimate.builder();
         statsBuilder.setOutputRowCount(node.getRowCount());
 
-        for (int symbolId = 0; symbolId < node.getOutputSymbols().size(); ++symbolId) {
-            Symbol symbol = node.getOutputSymbols().get(symbolId);
-            List<Object> symbolValues = getSymbolValues(
-                    node,
-                    symbolId,
-                    session,
-                    RowType.anonymous(node.getOutputSymbols().stream()
-                            .map(types::get)
-                            .collect(toImmutableList())));
-            statsBuilder.addSymbolStatistics(symbol, buildSymbolStatistics(symbolValues, types.get(symbol)));
+        try {
+            for (int symbolId = 0; symbolId < node.getOutputSymbols().size(); ++symbolId) {
+                Symbol symbol = node.getOutputSymbols().get(symbolId);
+                List<Object> symbolValues = getSymbolValues(
+                        node,
+                        symbolId,
+                        session,
+                        RowType.anonymous(node.getOutputSymbols().stream()
+                                .map(types::get)
+                                .collect(toImmutableList())));
+                statsBuilder.addSymbolStatistics(symbol, buildSymbolStatistics(symbolValues, types.get(symbol)));
+            }
+        }
+        catch (RuntimeException e) {
+            // prevent stats calculations (e.g. division by zero) from causing planning failures
+            return Optional.empty();
         }
 
         return Optional.of(statsBuilder.build());

core/trino-main/src/main/java/io/trino/sql/planner/LogicalPlanner.java

@@ -48,6 +48,7 @@
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.connector.ConnectorTableMetadata;
 import io.trino.spi.security.AccessDeniedException;
+import io.trino.spi.statistics.TableStatistics;
 import io.trino.spi.statistics.TableStatisticsMetadata;
 import io.trino.spi.type.CharType;
 import io.trino.spi.type.Type;
@@ -261,7 +262,7 @@ public Plan plan(Analysis analysis, Stage stage, boolean collectPlanStatistics)
             planSanityChecker.validateIntermediatePlan(root, session, plannerContext, typeAnalyzer, symbolAllocator.getTypes(), warningCollector);
         }
 
-        TableStatsProvider tableStatsProvider = new CachingTableStatsProvider(metadata, session);
+        CachingTableStatsProvider tableStatsProvider = new CachingTableStatsProvider(metadata, session);
 
         if (stage.ordinal() >= OPTIMIZED.ordinal()) {
             try (var ignored = scopedSpan(plannerContext.getTracer(), "optimizer")) {
@@ -280,13 +281,22 @@ public Plan plan(Analysis analysis, Stage stage, boolean collectPlanStatistics)
 
         TypeProvider types = symbolAllocator.getTypes();
 
-        StatsAndCosts statsAndCosts = StatsAndCosts.empty();
+        TableStatsProvider collectTableStatsProvider;
         if (collectPlanStatistics) {
-            StatsProvider statsProvider = new CachingStatsProvider(statsCalculator, session, types, tableStatsProvider);
-            CostProvider costProvider = new CachingCostProvider(costCalculator, statsProvider, Optional.empty(), session, types);
-            try (var ignored = scopedSpan(plannerContext.getTracer(), "plan-stats")) {
-                statsAndCosts = StatsAndCosts.create(root, statsProvider, costProvider);
-            }
+            collectTableStatsProvider = tableStatsProvider;
+        }
+        else {
+            // If stats collection was not requested explicitly, then use statistics
+            // that were fetched and cached during planning.
+            Map<TableHandle, TableStatistics> cachedStatistics = tableStatsProvider.getCachedTableStatistics();
+            collectTableStatsProvider = handle -> cachedStatistics.getOrDefault(handle, TableStatistics.empty());
+        }
+
+        StatsAndCosts statsAndCosts;
+        StatsProvider statsProvider = new CachingStatsProvider(statsCalculator, session, types, collectTableStatsProvider);
+        CostProvider costProvider = new CachingCostProvider(costCalculator, statsProvider, Optional.empty(), session, types);
+        try (var ignored = scopedSpan(plannerContext.getTracer(), "plan-stats")) {
+            statsAndCosts = StatsAndCosts.create(root, statsProvider, costProvider);
         }
         return new Plan(root, types, statsAndCosts);
     }

core/trino-main/src/main/java/io/trino/sql/planner/optimizations/DeterminePartitionCount.java

@@ -132,16 +132,9 @@ public PlanNode optimize(
             return plan;
         }
 
-        try {
-            return determinePartitionCount(plan, session, types, tableStatsProvider, isWriteQuery)
-                    .map(partitionCount -> rewriteWith(new Rewriter(partitionCount), plan))
-                    .orElse(plan);
-        }
-        catch (RuntimeException e) {
-            log.warn(e, "Error occurred when determining hash partition count for query %s", session.getQueryId());
-        }
-
-        return plan;
+        return determinePartitionCount(plan, session, types, tableStatsProvider, isWriteQuery)
+                .map(partitionCount -> rewriteWith(new Rewriter(partitionCount), plan))
+                .orElse(plan);
     }
 
     private Optional<Integer> determinePartitionCount(

core/trino-main/src/main/java/io/trino/sql/planner/planprinter/JsonRenderer.java

@@ -51,7 +51,7 @@ protected JsonRenderedNode renderJson(PlanRepresentation plan, NodeRepresentatio
                 node.getDescriptor(),
                 node.getOutputs(),
                 node.getDetails(),
-                node.getEstimates(plan.getTypes()),
+                node.getEstimates(),
                 children);
     }
 

core/trino-main/src/main/java/io/trino/sql/planner/planprinter/NodeRepresentation.java

@@ -13,14 +13,15 @@
  */
 package io.trino.sql.planner.planprinter;
 
-import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.FormatMethod;
 import io.trino.cost.LocalCostEstimate;
 import io.trino.cost.PlanCostEstimate;
 import io.trino.cost.PlanNodeStatsAndCostSummary;
 import io.trino.cost.PlanNodeStatsEstimate;
+import io.trino.spi.type.Type;
 import io.trino.sql.planner.Symbol;
 import io.trino.sql.planner.TypeProvider;
 import io.trino.sql.planner.plan.PlanFragmentId;
@@ -33,6 +34,7 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
@@ -151,14 +153,16 @@ public Optional<PlanNodeStatsAndCostSummary> getReorderJoinStatsAndCost()
         return reorderJoinStatsAndCost;
     }
 
-    public List<PlanNodeStatsAndCostSummary> getEstimates(TypeProvider typeProvider)
+    public List<PlanNodeStatsAndCostSummary> getEstimates()
     {
         if (getEstimatedStats().stream().allMatch(PlanNodeStatsEstimate::isOutputRowCountUnknown) &&
                 getEstimatedCost().stream().allMatch(c -> c.getRootNodeLocalCostEstimate().equals(LocalCostEstimate.unknown()))) {
             return ImmutableList.of();
         }
 
         ImmutableList.Builder<PlanNodeStatsAndCostSummary> estimates = ImmutableList.builder();
+        TypeProvider typeProvider = TypeProvider.copyOf(outputs.stream()
+                .collect(toImmutableMap(TypedSymbol::getSymbol, TypedSymbol::getTrinoType)));
         for (int i = 0; i < getEstimatedStats().size(); i++) {
             PlanNodeStatsEstimate stats = getEstimatedStats().get(i);
             LocalCostEstimate cost = getEstimatedCost().get(i).getRootNodeLocalCostEstimate();
@@ -181,13 +185,12 @@ public List<PlanNodeStatsAndCostSummary> getEstimates(TypeProvider typeProvider)
     public static class TypedSymbol
     {
         private final Symbol symbol;
-        private final String type;
+        private final Type trinoType;
 
-        @JsonCreator
-        public TypedSymbol(@JsonProperty("symbol") Symbol symbol, @JsonProperty("type") String type)
+        public TypedSymbol(Symbol symbol, Type trinoType)
         {
             this.symbol = symbol;
-            this.type = type;
+            this.trinoType = trinoType;
         }
 
         @JsonProperty
@@ -199,10 +202,16 @@ public Symbol getSymbol()
         @JsonProperty
         public String getType()
         {
-            return type;
+            return trinoType.getDisplayName();
         }
 
-        public static TypedSymbol typedSymbol(String symbol, String type)
+        @JsonIgnore
+        public Type getTrinoType()
+        {
+            return trinoType;
+        }
+
+        public static TypedSymbol typedSymbol(String symbol, Type type)
         {
             return new TypedSymbol(new Symbol(symbol), type);
         }
@@ -218,13 +227,13 @@ public boolean equals(Object o)
             }
             TypedSymbol that = (TypedSymbol) o;
             return symbol.equals(that.symbol)
-                    && type.equals(that.type);
+                    && trinoType.equals(that.trinoType);
         }
 
         @Override
         public int hashCode()
         {
-            return Objects.hash(symbol, type);
+            return Objects.hash(symbol, trinoType);
         }
     }
 }

core/trino-main/src/main/java/io/trino/sql/planner/planprinter/PlanPrinter.java

@@ -2160,7 +2160,7 @@ public NodeRepresentation addNode(
                     rootNode.getClass().getSimpleName(),
                     descriptor,
                     rootNode.getOutputSymbols().stream()
-                            .map(s -> new TypedSymbol(new Symbol(anonymizer.anonymize(s)), types.get(s).getDisplayName()))
+                            .map(s -> new TypedSymbol(new Symbol(anonymizer.anonymize(s)), types.get(s)))
                             .collect(toImmutableList()),
                     stats.map(s -> s.get(rootNode.getId())),
                     estimatedStats,

core/trino-main/src/main/java/io/trino/sql/planner/planprinter/TextRenderer.java

@@ -84,7 +84,7 @@ private String writeTextOutput(StringBuilder output, PlanRepresentation plan, In
             output.append(indentMultilineString(reorderJoinStatsAndCost, indent.detailIndent()));
         }
 
-        List<PlanNodeStatsAndCostSummary> estimates = node.getEstimates(plan.getTypes());
+        List<PlanNodeStatsAndCostSummary> estimates = node.getEstimates();
         if (!estimates.isEmpty()) {
             output.append(indentMultilineString(printEstimates(estimates), indent.detailIndent()));
         }

core/trino-main/src/test/java/io/trino/cost/TestValuesNodeStats.java

@@ -17,6 +17,7 @@
 import io.trino.sql.planner.Symbol;
 import org.junit.jupiter.api.Test;
 
+import static io.trino.cost.PlanNodeStatsEstimate.unknown;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.VarcharType.createVarcharType;
@@ -30,11 +31,11 @@ public class TestValuesNodeStats
     public void testStatsForValuesNode()
     {
         tester().assertStatsFor(pb -> pb
-                .values(ImmutableList.of(pb.symbol("a", BIGINT), pb.symbol("b", DOUBLE)),
-                        ImmutableList.of(
-                                ImmutableList.of(expression("3+3"), expression("13.5e0")),
-                                ImmutableList.of(expression("55"), expression("null")),
-                                ImmutableList.of(expression("6"), expression("13.5e0")))))
+                        .values(ImmutableList.of(pb.symbol("a", BIGINT), pb.symbol("b", DOUBLE)),
+                                ImmutableList.of(
+                                        ImmutableList.of(expression("3+3"), expression("13.5e0")),
+                                        ImmutableList.of(expression("55"), expression("null")),
+                                        ImmutableList.of(expression("6"), expression("13.5e0")))))
                 .check(outputStats -> outputStats.equalTo(
                         PlanNodeStatsEstimate.builder()
                                 .setOutputRowCount(3)
@@ -57,12 +58,12 @@ public void testStatsForValuesNode()
                                 .build()));
 
         tester().assertStatsFor(pb -> pb
-                .values(ImmutableList.of(pb.symbol("v", createVarcharType(30))),
-                        ImmutableList.of(
-                                ImmutableList.of(expression("'Alice'")),
-                                ImmutableList.of(expression("'has'")),
-                                ImmutableList.of(expression("'a cat'")),
-                                ImmutableList.of(expression("null")))))
+                        .values(ImmutableList.of(pb.symbol("v", createVarcharType(30))),
+                                ImmutableList.of(
+                                        ImmutableList.of(expression("'Alice'")),
+                                        ImmutableList.of(expression("'has'")),
+                                        ImmutableList.of(expression("'a cat'")),
+                                        ImmutableList.of(expression("null")))))
                 .check(outputStats -> outputStats.equalTo(
                         PlanNodeStatsEstimate.builder()
                                 .setOutputRowCount(4)
@@ -76,6 +77,15 @@ public void testStatsForValuesNode()
                                 .build()));
     }
 
+    @Test
+    public void testDivisionByZero()
+    {
+        tester().assertStatsFor(pb -> pb
+                        .values(ImmutableList.of(pb.symbol("a", BIGINT)),
+                                ImmutableList.of(ImmutableList.of(expression("1 / 0")))))
+                .check(outputStats -> outputStats.equalTo(unknown()));
+    }
+
     @Test
     public void testStatsForValuesNodeWithJustNulls()
     {
@@ -85,30 +95,30 @@ public void testStatsForValuesNodeWithJustNulls()
                 .build();
 
         tester().assertStatsFor(pb -> pb
-                .values(ImmutableList.of(pb.symbol("a", BIGINT)),
-                        ImmutableList.of(
-                                ImmutableList.of(expression("3 + null")))))
+                        .values(ImmutableList.of(pb.symbol("a", BIGINT)),
+                                ImmutableList.of(
+                                        ImmutableList.of(expression("3 + null")))))
                 .check(outputStats -> outputStats.equalTo(nullAStats));
 
         tester().assertStatsFor(pb -> pb
-                .values(ImmutableList.of(pb.symbol("a", BIGINT)),
-                        ImmutableList.of(
-                                ImmutableList.of(expression("null")))))
+                        .values(ImmutableList.of(pb.symbol("a", BIGINT)),
+                                ImmutableList.of(
+                                        ImmutableList.of(expression("null")))))
                 .check(outputStats -> outputStats.equalTo(nullAStats));
 
         tester().assertStatsFor(pb -> pb
-                .values(ImmutableList.of(pb.symbol("a", UNKNOWN)),
-                        ImmutableList.of(
-                                ImmutableList.of(expression("null")))))
+                        .values(ImmutableList.of(pb.symbol("a", UNKNOWN)),
+                                ImmutableList.of(
+                                        ImmutableList.of(expression("null")))))
                 .check(outputStats -> outputStats.equalTo(nullAStats));
     }
 
     @Test
     public void testStatsForEmptyValues()
     {
         tester().assertStatsFor(pb -> pb
-                .values(ImmutableList.of(pb.symbol("a", BIGINT)),
-                        ImmutableList.of()))
+                        .values(ImmutableList.of(pb.symbol("a", BIGINT)),
+                                ImmutableList.of()))
                 .check(outputStats -> outputStats.equalTo(
                         PlanNodeStatsEstimate.builder()
                                 .setOutputRowCount(0)