Add scale-writers.max-nodes-count property to limit writer scaling

.github/workflows/ci.yml

@@ -92,7 +92,7 @@ jobs:
         if: steps.cache.outputs.cache-hit != 'true'
         run: rm -rf ~/.m2/repository/io/trino/trino-*
 
-  check-commits:
+  check-commits-dispatcher:
     runs-on: ubuntu-latest
     if: github.event_name == 'pull_request'
     outputs:
@@ -101,12 +101,12 @@ jobs:
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0 # checkout all commits to be able to determine merge base
-      - name: Check Commits
+      - name: Block illegal commits
         uses: trinodb/github-actions/block-commits@c2991972560c5219d9ae5fb68c0c9d687ffcdd10
         with:
           action-merge: fail
           action-fixup: none
-      - name: Set matrix
+      - name: Set matrix (dispatch commit checks)
         id: set-matrix
         run: |
           # The output from rev-list ends with a newline, so we have to filter out index -1 in jq since it's an empty string
@@ -125,12 +125,13 @@ jobs:
           echo "Commit matrix: $(jq '.' commit-matrix.json)"
           echo "matrix=$(jq -c '.' commit-matrix.json)" >> $GITHUB_OUTPUT
 
-  check-commits-dispatcher:
+  check-commit:
     runs-on: ubuntu-latest
-    needs: check-commits
-    if: github.event_name == 'pull_request' && needs.check-commits.outputs.matrix != ''
+    needs: check-commits-dispatcher
+    if: github.event_name == 'pull_request' && needs.check-commits-dispatcher.outputs.matrix != ''
     strategy:
-      matrix: ${{ fromJson(needs.check-commits.outputs.matrix) }}
+      fail-fast: false
+      matrix: ${{ fromJson(needs.check-commits-dispatcher.outputs.matrix) }}
     steps:
       - uses: actions/checkout@v3
         with:

client/trino-cli/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <groupId>io.trino</groupId>
         <artifactId>trino-root</artifactId>
-        <version>406-SNAPSHOT</version>
+        <version>407-SNAPSHOT</version>
         <relativePath>../../pom.xml</relativePath>
     </parent>
 

client/trino-client/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <groupId>io.trino</groupId>
         <artifactId>trino-root</artifactId>
-        <version>406-SNAPSHOT</version>
+        <version>407-SNAPSHOT</version>
         <relativePath>../../pom.xml</relativePath>
     </parent>
 

client/trino-jdbc/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <groupId>io.trino</groupId>
         <artifactId>trino-root</artifactId>
-        <version>406-SNAPSHOT</version>
+        <version>407-SNAPSHOT</version>
         <relativePath>../../pom.xml</relativePath>
     </parent>
 

client/trino-jdbc/src/test/java/io/trino/jdbc/TestTrinoDatabaseMetaData.java

@@ -813,7 +813,7 @@ public void testGetColumns()
                 assertColumnSpec(rs, Types.TIME, 15L, null, 6L, null, createTimeType(6));
                 assertColumnSpec(rs, Types.TIME, 18L, null, 9L, null, createTimeType(9));
                 assertColumnSpec(rs, Types.TIME, 21L, null, 12L, null, createTimeType(12));
-                assertColumnSpec(rs, Types.TIME_WITH_TIMEZONE, 18L, null, 3L, null, TimeWithTimeZoneType.TIME_WITH_TIME_ZONE);
+                assertColumnSpec(rs, Types.TIME_WITH_TIMEZONE, 18L, null, 3L, null, TimeWithTimeZoneType.TIME_TZ_MILLIS);
                 assertColumnSpec(rs, Types.TIME_WITH_TIMEZONE, 14L, null, 0L, null, createTimeWithTimeZoneType(0));
                 assertColumnSpec(rs, Types.TIME_WITH_TIMEZONE, 18L, null, 3L, null, createTimeWithTimeZoneType(3));
                 assertColumnSpec(rs, Types.TIME_WITH_TIMEZONE, 21L, null, 6L, null, createTimeWithTimeZoneType(6));

core/trino-main/pom.xml

@@ -5,7 +5,7 @@
     <parent>
         <groupId>io.trino</groupId>
         <artifactId>trino-root</artifactId>
-        <version>406-SNAPSHOT</version>
+        <version>407-SNAPSHOT</version>
         <relativePath>../../pom.xml</relativePath>
     </parent>
 

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

@@ -60,7 +60,8 @@ public final class SystemSessionProperties
     public static final String JOIN_MAX_BROADCAST_TABLE_SIZE = "join_max_broadcast_table_size";
     public static final String JOIN_MULTI_CLAUSE_INDEPENDENCE_FACTOR = "join_multi_clause_independence_factor";
     public static final String DISTRIBUTED_INDEX_JOIN = "distributed_index_join";
-    public static final String HASH_PARTITION_COUNT = "hash_partition_count";
+    public static final String MAX_HASH_PARTITION_COUNT = "max_hash_partition_count";
+    public static final String MIN_HASH_PARTITION_COUNT = "min_hash_partition_count";
     public static final String PREFER_STREAMING_OPERATORS = "prefer_streaming_operators";
     public static final String TASK_WRITER_COUNT = "task_writer_count";
     public static final String TASK_PARTITIONED_WRITER_COUNT = "task_partitioned_writer_count";
@@ -80,6 +81,7 @@ public final class SystemSessionProperties
     public static final String PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS = "preferred_write_partitioning_min_number_of_partitions";
     public static final String SCALE_WRITERS = "scale_writers";
     public static final String TASK_SCALE_WRITERS_ENABLED = "task_scale_writers_enabled";
+    public static final String MAX_WRITERS_NODES_COUNT = "max_writers_nodes_count";
     public static final String TASK_SCALE_WRITERS_MAX_WRITER_COUNT = "task_scale_writers_max_writer_count";
     public static final String WRITER_MIN_SIZE = "writer_min_size";
     public static final String PUSH_TABLE_WRITE_THROUGH_UNION = "push_table_write_through_union";
@@ -173,10 +175,13 @@ public final class SystemSessionProperties
     public static final String ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS = "adaptive_partial_aggregation_min_rows";
     public static final String ADAPTIVE_PARTIAL_AGGREGATION_UNIQUE_ROWS_RATIO_THRESHOLD = "adaptive_partial_aggregation_unique_rows_ratio_threshold";
     public static final String JOIN_PARTITIONED_BUILD_MIN_ROW_COUNT = "join_partitioned_build_min_row_count";
+    public static final String MIN_INPUT_SIZE_PER_TASK = "min_input_size_per_task";
+    public static final String MIN_INPUT_ROWS_PER_TASK = "min_input_rows_per_task";
     public static final String USE_EXACT_PARTITIONING = "use_exact_partitioning";
     public static final String FORCE_SPILLING_JOIN = "force_spilling_join";
     public static final String FAULT_TOLERANT_EXECUTION_EVENT_DRIVEN_SCHEDULER_ENABLED = "fault_tolerant_execution_event_driven_scheduler_enabled";
     public static final String FAULT_TOLERANT_EXECUTION_FORCE_PREFERRED_WRITE_PARTITIONING_ENABLED = "fault_tolerant_execution_force_preferred_write_partitioning_enabled";
+    public static final String PAGE_PARTITIONING_BUFFER_POOL_SIZE = "page_partitioning_buffer_pool_size";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -241,9 +246,14 @@ public SystemSessionProperties(
                         optimizerConfig.isDistributedIndexJoinsEnabled(),
                         false),
                 integerProperty(
-                        HASH_PARTITION_COUNT,
-                        "Number of partitions for distributed joins and aggregations",
-                        queryManagerConfig.getHashPartitionCount(),
+                        MAX_HASH_PARTITION_COUNT,
+                        "Maximum number of partitions for distributed joins and aggregations",
+                        queryManagerConfig.getMaxHashPartitionCount(),
+                        false),
+                integerProperty(
+                        MIN_HASH_PARTITION_COUNT,
+                        "Minimum number of partitions for distributed joins and aggregations",
+                        queryManagerConfig.getMinHashPartitionCount(),
                         false),
                 booleanProperty(
                         PREFER_STREAMING_OPERATORS,
@@ -286,6 +296,11 @@ public SystemSessionProperties(
                         "Scale out writers based on throughput (use minimum necessary)",
                         featuresConfig.isScaleWriters(),
                         false),
+                integerProperty(
+                        MAX_WRITERS_NODES_COUNT,
+                        "Set upper limit on number of nodes that take part in writing if task.scale-writers.enabled is set",
+                        queryManagerConfig.getMaxWritersNodesCount(),
+                        false),
                 booleanProperty(
                         TASK_SCALE_WRITERS_ENABLED,
                         "Scale the number of concurrent table writers per task based on throughput",
@@ -860,6 +875,16 @@ public SystemSessionProperties(
                         optimizerConfig.getJoinPartitionedBuildMinRowCount(),
                         value -> validateNonNegativeLongValue(value, JOIN_PARTITIONED_BUILD_MIN_ROW_COUNT),
                         false),
+                dataSizeProperty(
+                        MIN_INPUT_SIZE_PER_TASK,
+                        "Minimum input data size required per task. This will help optimizer determine hash partition count for joins and aggregations",
+                        optimizerConfig.getMinInputSizePerTask(),
+                        false),
+                longProperty(
+                        MIN_INPUT_ROWS_PER_TASK,
+                        "Minimum input rows required per task. This will help optimizer determine hash partition count for joins and aggregations",
+                        optimizerConfig.getMinInputRowsPerTask(),
+                        false),
                 booleanProperty(
                         USE_EXACT_PARTITIONING,
                         "When enabled this forces data repartitioning unless the partitioning of upstream stage matches exactly what downstream stage expects",
@@ -879,6 +904,10 @@ public SystemSessionProperties(
                         FAULT_TOLERANT_EXECUTION_FORCE_PREFERRED_WRITE_PARTITIONING_ENABLED,
                         "Force preferred write partitioning for fault tolerant execution",
                         queryManagerConfig.isFaultTolerantExecutionForcePreferredWritePartitioningEnabled(),
+                        true),
+                integerProperty(PAGE_PARTITIONING_BUFFER_POOL_SIZE,
+                        "Maximum number of free buffers in the per task partitioned page buffer pool. Setting this to zero effectively disables the pool",
+                        taskManagerConfig.getPagePartitioningBufferPoolSize(),
                         true));
     }
 
@@ -918,9 +947,14 @@ public static boolean isDistributedIndexJoinEnabled(Session session)
         return session.getSystemProperty(DISTRIBUTED_INDEX_JOIN, Boolean.class);
     }
 
-    public static int getHashPartitionCount(Session session)
+    public static int getMaxHashPartitionCount(Session session)
+    {
+        return session.getSystemProperty(MAX_HASH_PARTITION_COUNT, Integer.class);
+    }
+
+    public static int getMinHashPartitionCount(Session session)
     {
-        return session.getSystemProperty(HASH_PARTITION_COUNT, Integer.class);
+        return session.getSystemProperty(MIN_HASH_PARTITION_COUNT, Integer.class);
     }
 
     public static boolean preferStreamingOperators(Session session)
@@ -968,6 +1002,11 @@ public static int getTaskScaleWritersMaxWriterCount(Session session)
         return session.getSystemProperty(TASK_SCALE_WRITERS_MAX_WRITER_COUNT, Integer.class);
     }
 
+    public static int getMaxWritersNodesCount(Session session)
+    {
+        return session.getSystemProperty(MAX_WRITERS_NODES_COUNT, Integer.class);
+    }
+
     public static DataSize getWriterMinSize(Session session)
     {
         return session.getSystemProperty(WRITER_MIN_SIZE, DataSize.class);
@@ -1548,6 +1587,16 @@ public static long getJoinPartitionedBuildMinRowCount(Session session)
         return session.getSystemProperty(JOIN_PARTITIONED_BUILD_MIN_ROW_COUNT, Long.class);
     }
 
+    public static DataSize getMinInputSizePerTask(Session session)
+    {
+        return session.getSystemProperty(MIN_INPUT_SIZE_PER_TASK, DataSize.class);
+    }
+
+    public static long getMinInputRowsPerTask(Session session)
+    {
+        return session.getSystemProperty(MIN_INPUT_ROWS_PER_TASK, Long.class);
+    }
+
     public static boolean isUseExactPartitioning(Session session)
     {
         return session.getSystemProperty(USE_EXACT_PARTITIONING, Boolean.class);
@@ -1571,4 +1620,9 @@ public static boolean isFaultTolerantExecutionForcePreferredWritePartitioningEna
         }
         return session.getSystemProperty(FAULT_TOLERANT_EXECUTION_FORCE_PREFERRED_WRITE_PARTITIONING_ENABLED, Boolean.class);
     }
+
+    public static int getPagePartitioningBufferPoolSize(Session session)
+    {
+        return session.getSystemProperty(PAGE_PARTITIONING_BUFFER_POOL_SIZE, Integer.class);
+    }
 }

core/trino-main/src/main/java/io/trino/connector/ConnectorContextInstance.java

@@ -17,6 +17,7 @@
 import io.trino.spi.PageIndexerFactory;
 import io.trino.spi.PageSorter;
 import io.trino.spi.VersionEmbedder;
+import io.trino.spi.connector.CatalogHandle;
 import io.trino.spi.connector.ConnectorContext;
 import io.trino.spi.connector.MetadataProvider;
 import io.trino.spi.type.TypeManager;
@@ -38,8 +39,10 @@ public class ConnectorContextInstance
     private final PageIndexerFactory pageIndexerFactory;
     private final Supplier<ClassLoader> duplicatePluginClassLoaderFactory;
     private final AtomicBoolean pluginClassLoaderDuplicated = new AtomicBoolean();
+    private final CatalogHandle catalogHandle;
 
     public ConnectorContextInstance(
+            CatalogHandle catalogHandle,
             NodeManager nodeManager,
             VersionEmbedder versionEmbedder,
             TypeManager typeManager,
@@ -55,6 +58,13 @@ public ConnectorContextInstance(
         this.pageSorter = requireNonNull(pageSorter, "pageSorter is null");
         this.pageIndexerFactory = requireNonNull(pageIndexerFactory, "pageIndexerFactory is null");
         this.duplicatePluginClassLoaderFactory = requireNonNull(duplicatePluginClassLoaderFactory, "duplicatePluginClassLoaderFactory is null");
+        this.catalogHandle = requireNonNull(catalogHandle, "catalogHandle is null");
+    }
+
+    @Override
+    public CatalogHandle getCatalogHandle()
+    {
+        return catalogHandle;
     }
 
     @Override

core/trino-main/src/main/java/io/trino/connector/DefaultCatalogFactory.java

@@ -190,6 +190,7 @@ private Connector createConnector(
             Map<String, String> properties)
     {
         ConnectorContext context = new ConnectorContextInstance(
+                catalogHandle,
                 new ConnectorAwareNodeManager(nodeManager, nodeInfo.getEnvironment(), catalogHandle, schedulerIncludeCoordinator),
                 versionEmbedder,
                 typeManager,

core/trino-main/src/main/java/io/trino/connector/InternalMetadataProvider.java

@@ -13,6 +13,7 @@
  */
 package io.trino.connector;
 
+import com.google.common.collect.ImmutableList;
 import io.trino.FullConnectorSession;
 import io.trino.Session;
 import io.trino.metadata.MaterializedViewDefinition;
@@ -54,12 +55,12 @@ public Optional<ConnectorTableSchema> getRelationMetadata(ConnectorSession conne
 
         Optional<MaterializedViewDefinition> materializedView = metadata.getMaterializedView(session, qualifiedName);
         if (materializedView.isPresent()) {
-            return Optional.of(new ConnectorTableSchema(tableName.getSchemaTableName(), toColumnSchema(materializedView.get().getColumns())));
+            return Optional.of(new ConnectorTableSchema(tableName.getSchemaTableName(), toColumnSchema(materializedView.get().getColumns()), ImmutableList.of()));
         }
 
         Optional<ViewDefinition> view = metadata.getView(session, qualifiedName);
         if (view.isPresent()) {
-            return Optional.of(new ConnectorTableSchema(tableName.getSchemaTableName(), toColumnSchema(view.get().getColumns())));
+            return Optional.of(new ConnectorTableSchema(tableName.getSchemaTableName(), toColumnSchema(view.get().getColumns()), ImmutableList.of()));
         }
 
         Optional<TableHandle> tableHandle = metadata.getTableHandle(session, qualifiedName);

core/trino-main/src/main/java/io/trino/connector/informationschema/InformationSchemaPageSource.java

@@ -54,6 +54,7 @@
 import static io.trino.connector.informationschema.InformationSchemaMetadata.defaultPrefixes;
 import static io.trino.connector.informationschema.InformationSchemaMetadata.isTablesEnumeratingTable;
 import static io.trino.metadata.MetadataListing.getViews;
+import static io.trino.metadata.MetadataListing.listMaterializedViews;
 import static io.trino.metadata.MetadataListing.listSchemas;
 import static io.trino.metadata.MetadataListing.listTableColumns;
 import static io.trino.metadata.MetadataListing.listTablePrivileges;
@@ -271,12 +272,18 @@ private void addColumnsRecords(QualifiedTablePrefix prefix)
     private void addTablesRecords(QualifiedTablePrefix prefix)
     {
         Set<SchemaTableName> tables = listTables(session, metadata, accessControl, prefix);
+        Set<SchemaTableName> materializedViews = listMaterializedViews(session, metadata, accessControl, prefix);
         Set<SchemaTableName> views = listViews(session, metadata, accessControl, prefix);
-        // TODO (https://github.com/trinodb/trino/issues/8207) define a type for materialized views
 
-        for (SchemaTableName name : union(tables, views)) {
+        for (SchemaTableName name : union(union(tables, materializedViews), views)) {
             // if table and view names overlap, the view wins
-            String type = views.contains(name) ? "VIEW" : "BASE TABLE";
+            String type = "BASE TABLE";
+            if (materializedViews.contains(name)) {
+                type = "MATERIALIZED VIEW";
+            }
+            else if (views.contains(name)) {
+                type = "VIEW";
+            }
             addRecord(
                     prefix.getCatalogName(),
                     name.getSchemaName(),

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

@@ -26,7 +26,7 @@
 
 import static io.trino.SystemSessionProperties.getCostEstimationWorkerCount;
 import static io.trino.SystemSessionProperties.getFaultTolerantExecutionPartitionCount;
-import static io.trino.SystemSessionProperties.getHashPartitionCount;
+import static io.trino.SystemSessionProperties.getMaxHashPartitionCount;
 import static io.trino.SystemSessionProperties.getRetryPolicy;
 import static java.lang.Math.min;
 import static java.lang.Math.toIntExact;
@@ -73,7 +73,7 @@ public int estimateHashedTaskCount(Session session)
             partitionCount = getFaultTolerantExecutionPartitionCount(session);
         }
         else {
-            partitionCount = getHashPartitionCount(session);
+            partitionCount = getMaxHashPartitionCount(session);
         }
         return min(estimateSourceDistributedTaskCount(session), partitionCount);
     }