Add query.max-select-rows property

Author: vvvvh123

core/trino-main/src/main/java/io/trino/SystemSessionProperties.java

@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.OptionalInt;
+import java.util.OptionalLong;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
@@ -220,6 +221,7 @@ public final class SystemSessionProperties
     public static final String CLOSE_IDLE_WRITERS_TRIGGER_DURATION = "close_idle_writers_trigger_duration";
     public static final String COLUMNAR_FILTER_EVALUATION_ENABLED = "columnar_filter_evaluation_enabled";
     public static final String SPOOLING_ENABLED = "spooling_enabled";
+    public static final String QUERY_MAX_SELECT_ROWS = "query_max_select_rows";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -419,6 +421,12 @@ public SystemSessionProperties(
                         "Temporary: Maximum number of stages a query can have",
                         queryManagerConfig.getMaxStageCount(),
                         true),
+                longProperty(
+                        QUERY_MAX_SELECT_ROWS,
+                        "Maximum rows returned by a top-level SELECT query",
+                        queryManagerConfig.getQueryMaxSelectRows().isPresent() ? queryManagerConfig.getQueryMaxSelectRows().getAsLong() : null,
+                        value -> validateNonNegativeLongValue(value, QUERY_MAX_SELECT_ROWS),
+                        false),
                 booleanProperty(
                         DICTIONARY_AGGREGATION,
                         "Enable optimization for aggregations on dictionaries",
@@ -1298,6 +1306,12 @@ public static int getQueryMaxStageCount(Session session)
         return session.getSystemProperty(QUERY_MAX_STAGE_COUNT, Integer.class);
     }
 
+    public static OptionalLong getQueryMaxSelectRows(Session session)
+    {
+        Long value = session.getSystemProperty(QUERY_MAX_SELECT_ROWS, Long.class);
+        return value == null ? OptionalLong.empty() : OptionalLong.of(value);
+    }
+
     public static boolean isUseTableScanNodePartitioning(Session session)
     {
         return session.getSystemProperty(USE_TABLE_SCAN_NODE_PARTITIONING, Boolean.class);

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

@@ -28,6 +28,7 @@
 import jakarta.validation.constraints.NotNull;
 
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -99,6 +100,8 @@ public class QueryManagerConfig
     private Optional<DataSize> queryMaxWritePhysicalSize = Optional.empty();
     private int queryReportedRuleStatsLimit = 10;
     private int dispatcherQueryPoolSize = DISPATCHER_THREADPOOL_MAX_SIZE;
+    @Min(0)
+    private Long queryMaxSelectRows;
 
     private int requiredWorkers = 1;
     private Duration requiredWorkersMaxWait = new Duration(5, TimeUnit.MINUTES);
@@ -1228,4 +1231,17 @@ public QueryManagerConfig setFaultTolerantExecutionAdaptiveJoinReorderingMinSize
         this.faultTolerantExecutionAdaptiveJoinReorderingMinSizeThreshold = faultTolerantExecutionAdaptiveJoinReorderingMinSizeThreshold;
         return this;
     }
+
+    public OptionalLong getQueryMaxSelectRows()
+    {
+        return queryMaxSelectRows == null ? OptionalLong.empty() : OptionalLong.of(queryMaxSelectRows);
+    }
+
+    @Config("query.max-select-rows")
+    @ConfigDescription("Maximum rows returned by a top-level SELECT query")
+    public QueryManagerConfig setQueryMaxSelectRows(Long queryMaxSelectRows)
+    {
+        this.queryMaxSelectRows = queryMaxSelectRows;
+        return this;
+    }
 }

core/trino-main/src/main/java/io/trino/sql/analyzer/Analyzer.java

@@ -31,6 +31,7 @@
 import io.trino.sql.tree.NodeRef;
 import io.trino.sql.tree.Parameter;
 import io.trino.sql.tree.Statement;
+import io.trino.util.StatementUtils;
 
 import java.util.List;
 import java.util.Map;
@@ -39,6 +40,8 @@
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractAggregateFunctions;
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractExpressions;
 import static io.trino.sql.analyzer.ExpressionTreeUtils.extractWindowExpressions;
+import static io.trino.sql.analyzer.QueryType.DESCRIBE;
+import static io.trino.sql.analyzer.QueryType.EXPLAIN;
 import static io.trino.sql.analyzer.QueryType.OTHERS;
 import static io.trino.sql.analyzer.SemanticExceptions.semanticException;
 import static io.trino.tracing.ScopedSpan.scopedSpan;
@@ -84,7 +87,16 @@ public Analysis analyze(Statement statement)
                 .setParent(Context.current().with(session.getQuerySpan()))
                 .startSpan();
         try (var _ = scopedSpan(span)) {
-            return analyze(statement, OTHERS);
+            // derive high-level type of the *original* statement before it gets rewritten
+            QueryType queryType = StatementUtils.getQueryType(statement)
+                    .map(original -> switch (original) {
+                        case DESCRIBE -> DESCRIBE;
+                        case EXPLAIN -> EXPLAIN;
+                        default -> OTHERS;
+                    })
+                    .orElse(OTHERS);
+
+            return analyze(statement, queryType);
         }
     }
 

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import io.trino.Session;
+import io.trino.SystemSessionProperties;
 import io.trino.metadata.Metadata;
 import io.trino.metadata.ResolvedFunction;
 import io.trino.metadata.TableHandle;
@@ -29,6 +30,7 @@
 import io.trino.spi.connector.ColumnSchema;
 import io.trino.spi.connector.RowChangeParadigm;
 import io.trino.spi.connector.SortOrder;
+import io.trino.spi.resourcegroups.QueryType;
 import io.trino.spi.type.DecimalType;
 import io.trino.spi.type.Int128;
 import io.trino.spi.type.RowType;
@@ -110,6 +112,7 @@
 import io.trino.sql.tree.WindowFrame;
 import io.trino.sql.tree.WindowOperation;
 import io.trino.type.Reals;
+import io.trino.util.StatementUtils;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -121,6 +124,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -239,7 +243,12 @@ public RelationPlan plan(Query query)
         Optional<OrderingScheme> orderingScheme = orderingScheme(builder, query.getOrderBy(), analysis.getOrderByExpressions(query));
         builder = sort(builder, orderingScheme);
         builder = offset(builder, query.getOffset());
-        builder = limit(builder, query.getLimit(), orderingScheme);
+        if (analysis.getStatement() == query) {
+            builder = limitWithTopLevelSelectCheck(builder, query, orderingScheme);
+        }
+        else {
+            builder = limit(builder, query.getLimit(), orderingScheme);
+        }
         builder = builder.appendProjections(outputs, symbolAllocator, idAllocator);
 
         return new RelationPlan(
@@ -2247,6 +2256,41 @@ private PlanBuilder offset(PlanBuilder subPlan, Optional<Offset> offset)
                         analysis.getOffset(offset.get())));
     }
 
+    private PlanBuilder limitWithTopLevelSelectCheck(PlanBuilder subPlan, Query query, Optional<OrderingScheme> orderingScheme)
+    {
+        Optional<Node> limit = query.getLimit();
+        if (limit.isPresent() && analysis.getLimit(limit.get()).isPresent()) {
+            Optional<OrderingScheme> tiesResolvingScheme = Optional.empty();
+            if (limit.get() instanceof FetchFirst && ((FetchFirst) limit.get()).isWithTies()) {
+                tiesResolvingScheme = orderingScheme;
+            }
+            return subPlan.withNewRoot(
+                    new LimitNode(
+                            idAllocator.getNextId(),
+                            subPlan.getRoot(),
+                            analysis.getLimit(limit.get()).getAsLong(),
+                            tiesResolvingScheme,
+                            false,
+                            ImmutableList.of()));
+        }
+
+        OptionalLong sessionLimit = SystemSessionProperties.getQueryMaxSelectRows(session);
+        if (sessionLimit.isPresent() && !analysis.isDescribe() &&
+                StatementUtils.getQueryType(analysis.getStatement())
+                        .map(type -> type == QueryType.SELECT)
+                        .orElse(false)) {
+            return subPlan.withNewRoot(
+                    new LimitNode(
+                            idAllocator.getNextId(),
+                            subPlan.getRoot(),
+                            sessionLimit.getAsLong(),
+                            Optional.empty(),
+                            false,
+                            ImmutableList.of()));
+        }
+        return subPlan;
+    }
+
     private PlanBuilder limit(PlanBuilder subPlan, Optional<Node> limit, Optional<OrderingScheme> orderingScheme)
     {
         if (limit.isPresent() && analysis.getLimit(limit.get()).isPresent()) {

core/trino-main/src/test/java/io/trino/execution/TestQueryManagerConfig.java

@@ -73,6 +73,7 @@ public void testDefaults()
                 .setDispatcherQueryPoolSize(Integer.toString(max(50, Runtime.getRuntime().availableProcessors() * 10)))
                 .setQueryMaxScanPhysicalBytes(null)
                 .setQueryMaxWritePhysicalSize(null)
+                .setQueryMaxSelectRows(null)
                 .setRequiredWorkers(1)
                 .setRequiredWorkersMaxWait(new Duration(5, MINUTES))
                 .setRetryPolicy(RetryPolicy.NONE)
@@ -157,6 +158,7 @@ public void testExplicitPropertyMappings()
                 .put("query.dispatcher-query-pool-size", "151")
                 .put("query.max-scan-physical-bytes", "1kB")
                 .put("query.max-write-physical-size", "1TB")
+                .put("query.max-select-rows", "10")
                 .put("query-manager.required-workers", "333")
                 .put("query-manager.required-workers-max-wait", "33m")
                 .put("retry-policy", "QUERY")
@@ -238,6 +240,7 @@ public void testExplicitPropertyMappings()
                 .setDispatcherQueryPoolSize("151")
                 .setQueryMaxScanPhysicalBytes(DataSize.of(1, KILOBYTE))
                 .setQueryMaxWritePhysicalSize(DataSize.of(1, TERABYTE))
+                .setQueryMaxSelectRows(10L)
                 .setRequiredWorkers(333)
                 .setRequiredWorkersMaxWait(new Duration(33, MINUTES))
                 .setRetryPolicy(RetryPolicy.QUERY)

docs/src/main/sphinx/admin/properties-query-management.md

@@ -283,3 +283,12 @@ The {ref}`retry policy <fte-retry-policy>` to use for
 - `TASK` - Retry individual tasks within a query in the event of failure.
   Requires configuration of an {ref}`exchange manager <fte-exchange-manager>`.
 - `QUERY` - Retry the whole query in the event of failure.
+
+## `query.max-select-rows`
+
+- **Type:** {ref}`prop-type-long`
+- **Session property:** `query_max_select_rows`
+
+The maximum number of rows that can be returned for a top-level `SELECT` query.
+If this property is set and a query also includes an explicit `LIMIT` clause, the more restrictive
+(the minimum) of the two limits is enforced.

docs/src/main/sphinx/admin/properties.md

@@ -119,3 +119,11 @@ Some `integer` type properties enforce their own minimum and maximum values.
 The properties of type `string` support a set of values that consist of a
 sequence of characters. Allowed values are defined on a property-by-property
 basis, refer to the specific property for its supported and default values.
+
+(prop-type-long)=
+### `long`
+
+The properties of type `long` support whole numeric values, similar to the `integer` type,
+but can accommodate a significantly larger range of numbers. For example, in addition
+to values like `100`, it supports very large numbers such as `4000000000`. Negative values are also supported.
+Like `integer`, `long` values must be whole numbers; decimal values such as `2.5` are not supported.