Add a limit to the number of columns in the CLUSTERED BY clause

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Limits.java

@@ -29,6 +29,16 @@ public class Limits
    */
   public static final int MAX_FRAME_COLUMNS = 2000;
 
+  /**
+   * Maximum number of columns that can appear in the clustered by clause
+   *
+   * There is some arbitrariness in the limit, but it is chosen such that the datasketches sketches do not blow up in
+   * memory while computing the partitions for the clustered by keys.
+   * This limit along sequential merge of the sketches will help prevent OOMs in both the workers and the controller
+   * tasks
+   */
+  public static final int MAX_CLUSTERED_BY_COLUMNS = (int) (MAX_FRAME_COLUMNS * 0.75);
+
   /**
    * Maximum number of workers that can be used in a stage, regardless of available memory.
    */

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryValidator.java

@@ -21,8 +21,10 @@
 
 import com.google.common.math.IntMath;
 import com.google.common.primitives.Ints;
+import org.apache.druid.frame.key.ClusterBy;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.msq.indexing.error.MSQException;
+import org.apache.druid.msq.indexing.error.TooManyClusteredByColumnsFault;
 import org.apache.druid.msq.indexing.error.TooManyColumnsFault;
 import org.apache.druid.msq.indexing.error.TooManyInputFilesFault;
 import org.apache.druid.msq.indexing.error.TooManyWorkersFault;
@@ -55,6 +57,15 @@ public static void validateQueryDef(final QueryDefinition queryDef)
         throw new ISE("Number of workers must be greater than 0");
       }
     }
+
+    // Check if the number of columns in the query's CLUSTERED BY clause donot exceed the limit
+    ClusterBy queryClusteredBy = queryDef.getFinalStageDefinition().getClusterBy();
+    int queryClusteredByColumnsSize = queryClusteredBy.getColumns().size();
+    if (queryClusteredByColumnsSize > Limits.MAX_CLUSTERED_BY_COLUMNS) {
+      throw new MSQException(
+          new TooManyClusteredByColumnsFault(queryClusteredByColumnsSize, Limits.MAX_CLUSTERED_BY_COLUMNS)
+      );
+    }
   }
 
   /**

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/error/TooManyClusteredByColumnsFault.java

@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.indexing.error;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonTypeName;
+
+import java.util.Objects;
+
+@JsonTypeName(TooManyClusteredByColumnsFault.CODE)
+public class TooManyClusteredByColumnsFault extends BaseMSQFault
+{
+  static final String CODE = "TooManyClusteredByColumns";
+
+  private final int numColumns;
+  private final int maxColumns;
+
+  @JsonCreator
+  public TooManyClusteredByColumnsFault(
+      @JsonProperty("numColumns") final int numColumns,
+      @JsonProperty("maxColumns") final int maxColumns
+  )
+  {
+    super(CODE, "Too many clustered by columns (requested = %d, max = %d)", numColumns, maxColumns);
+    this.numColumns = numColumns;
+    this.maxColumns = maxColumns;
+  }
+
+  @JsonProperty
+  public int getNumColumns()
+  {
+    return numColumns;
+  }
+
+  @JsonProperty
+  public int getMaxColumns()
+  {
+    return maxColumns;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    if (!super.equals(o)) {
+      return false;
+    }
+    TooManyClusteredByColumnsFault that = (TooManyClusteredByColumnsFault) o;
+    return numColumns == that.numColumns && maxColumns == that.maxColumns;
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(super.hashCode(), numColumns, maxColumns);
+  }
+}

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/error/MSQFaultSerdeTest.java

@@ -70,6 +70,7 @@ public void testFaultSerde() throws IOException
     assertFaultSerde(new TaskStartTimeoutFault(10));
     assertFaultSerde(new TooManyBucketsFault(10));
     assertFaultSerde(new TooManyColumnsFault(10, 8));
+    assertFaultSerde(new TooManyClusteredByColumnsFault(10, 8));
     assertFaultSerde(new TooManyInputFilesFault(15, 10, 5));
     assertFaultSerde(new TooManyPartitionsFault(10));
     assertFaultSerde(new TooManyWarningsFault(10, "the error"));