Add configurable final stages to MSQ ingestion queries

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQSqlModule.java

@@ -54,6 +54,8 @@ public void configure(Binder binder)
     // We want this module to bring InputSourceModule along for the ride.
     binder.install(new InputSourceModule());
 
+    binder.bind(MSQTerminalStageSpecFactory.class).toInstance(new MSQTerminalStageSpecFactory());
+
     binder.bind(MSQTaskSqlEngine.class).in(LazySingleton.class);
 
     // Set up the EXTERN macro.

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQTerminalStageSpecFactory.java

@@ -0,0 +1,37 @@
+/*
+ * 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.guice;
+
+import org.apache.druid.msq.indexing.destination.SegmentGenerationStageSpec;
+import org.apache.druid.msq.indexing.destination.TerminalStageSpec;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.sql.calcite.rel.DruidQuery;
+
+public class MSQTerminalStageSpecFactory
+{
+  /**
+   * Creates a {@link TerminalStageSpec} which determines the final of a query. Currently, always returns a segment
+   * generation spec, but this can be used to configure a wide range of behaviours.
+   */
+  public TerminalStageSpec createTerminalStageSpec(DruidQuery druidQuery, PlannerContext plannerContext)
+  {
+    return SegmentGenerationStageSpec.instance();
+  }
+}

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

@@ -243,7 +243,8 @@ private static DataSourceMSQDestination buildMSQDestination(
         dataSchema.getDimensionsSpec()
                   .getDimensions()
                   .stream()
-                  .collect(Collectors.toMap(DimensionSchema::getName, Function.identity()))
+                  .collect(Collectors.toMap(DimensionSchema::getName, Function.identity())),
+        null
     );
   }
 

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

@@ -51,6 +51,8 @@ public class DataSourceMSQDestination implements MSQDestination
   @Nullable
   private final List<Interval> replaceTimeChunks;
 
+  private final TerminalStageSpec terminalStageSpec;
+
   @Nullable
   private final Map<String, DimensionSchema> dimensionSchemas;
 
@@ -60,14 +62,16 @@ public DataSourceMSQDestination(
       @JsonProperty("segmentGranularity") Granularity segmentGranularity,
       @JsonProperty("segmentSortOrder") @Nullable List<String> segmentSortOrder,
       @JsonProperty("replaceTimeChunks") @Nullable List<Interval> replaceTimeChunks,
-      @JsonProperty("dimensionSchemas") @Nullable Map<String, DimensionSchema> dimensionSchemas
+      @JsonProperty("dimensionSchemas") @Nullable Map<String, DimensionSchema> dimensionSchemas,
+      @JsonProperty("terminalStageSpec") @Nullable TerminalStageSpec terminalStageSpec
   )
   {
     this.dataSource = Preconditions.checkNotNull(dataSource, "dataSource");
     this.segmentGranularity = Preconditions.checkNotNull(segmentGranularity, "segmentGranularity");
     this.segmentSortOrder = segmentSortOrder != null ? segmentSortOrder : Collections.emptyList();
     this.replaceTimeChunks = replaceTimeChunks;
     this.dimensionSchemas = dimensionSchemas;
+    this.terminalStageSpec = terminalStageSpec != null ? terminalStageSpec : SegmentGenerationStageSpec.instance();
 
     if (replaceTimeChunks != null) {
       // Verify that if replaceTimeChunks is provided, it is nonempty.
@@ -105,6 +109,17 @@ public String getDataSource()
     return dataSource;
   }
 
+  /**
+   * Returns the terminal stage spec.
+   * <p>
+   * The terminal stage spec, is a way to tell the MSQ task how to convert the results into segments at the final stage.
+   */
+  @JsonProperty
+  public TerminalStageSpec getTerminalStageSpec()
+  {
+    return terminalStageSpec;
+  }
+
   @JsonProperty
   public Granularity getSegmentGranularity()
   {
@@ -177,13 +192,14 @@ public boolean equals(Object o)
            && Objects.equals(segmentGranularity, that.segmentGranularity)
            && Objects.equals(segmentSortOrder, that.segmentSortOrder)
            && Objects.equals(replaceTimeChunks, that.replaceTimeChunks)
-           && Objects.equals(dimensionSchemas, that.dimensionSchemas);
+           && Objects.equals(dimensionSchemas, that.dimensionSchemas)
+           && Objects.equals(terminalStageSpec, that.terminalStageSpec);
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(dataSource, segmentGranularity, segmentSortOrder, replaceTimeChunks, dimensionSchemas);
+    return Objects.hash(dataSource, segmentGranularity, segmentSortOrder, replaceTimeChunks, dimensionSchemas, terminalStageSpec);
   }
 
   @Override
@@ -195,6 +211,7 @@ public String toString()
            ", segmentSortOrder=" + segmentSortOrder +
            ", replaceTimeChunks=" + replaceTimeChunks +
            ", dimensionSchemas=" + dimensionSchemas +
+           ", terminalStageSpec=" + terminalStageSpec +
            '}';
   }
 

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

@@ -0,0 +1,115 @@
+/*
+ * 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.destination;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Iterables;
+import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import org.apache.druid.frame.key.ClusterBy;
+import org.apache.druid.msq.indexing.MSQSpec;
+import org.apache.druid.msq.indexing.MSQTuningConfig;
+import org.apache.druid.msq.indexing.processor.SegmentGeneratorFrameProcessorFactory;
+import org.apache.druid.msq.input.stage.ReadablePartition;
+import org.apache.druid.msq.input.stage.StageInputSlice;
+import org.apache.druid.msq.input.stage.StageInputSpec;
+import org.apache.druid.msq.kernel.QueryDefinition;
+import org.apache.druid.msq.kernel.StageDefinition;
+import org.apache.druid.msq.kernel.StageDefinitionBuilder;
+import org.apache.druid.msq.kernel.controller.WorkerInputs;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.segment.indexing.DataSchema;
+import org.apache.druid.segment.realtime.appenderator.SegmentIdWithShardSpec;
+import org.apache.druid.sql.calcite.planner.ColumnMappings;
+
+import javax.annotation.Nullable;
+import java.util.ArrayList;
+import java.util.List;
+
+public class SegmentGenerationStageSpec implements TerminalStageSpec
+{
+  public static final String TYPE = "segmentGeneration";
+
+  private static final SegmentGenerationStageSpec INSTANCE = new SegmentGenerationStageSpec();
+
+  private SegmentGenerationStageSpec()
+  {
+  }
+
+  @JsonCreator
+  public static SegmentGenerationStageSpec instance()
+  {
+    return INSTANCE;
+  }
+
+  @Override
+  public StageDefinitionBuilder constructFinalStage(QueryDefinition queryDef, MSQSpec querySpec, ObjectMapper jsonMapper)
+  {
+    final MSQTuningConfig tuningConfig = querySpec.getTuningConfig();
+    final ColumnMappings columnMappings = querySpec.getColumnMappings();
+    final RowSignature querySignature = queryDef.getFinalStageDefinition().getSignature();
+    final ClusterBy queryClusterBy = queryDef.getFinalStageDefinition().getClusterBy();
+
+    // Add a segment-generation stage.
+    final DataSchema dataSchema =
+        SegmentGenerationUtils.makeDataSchemaForIngestion(querySpec, querySignature, queryClusterBy, columnMappings, jsonMapper);
+
+    return StageDefinition.builder(queryDef.getNextStageNumber())
+                       .inputs(new StageInputSpec(queryDef.getFinalStageDefinition().getStageNumber()))
+                       .maxWorkerCount(tuningConfig.getMaxNumWorkers())
+                       .processorFactory(
+                           new SegmentGeneratorFrameProcessorFactory(
+                               dataSchema,
+                               columnMappings,
+                               tuningConfig
+                           )
+    );
+  }
+
+  public Int2ObjectMap<List<SegmentIdWithShardSpec>> getWorkerInfo(
+      final WorkerInputs workerInputs,
+      @Nullable final List<SegmentIdWithShardSpec> segmentsToGenerate
+  )
+  {
+    final Int2ObjectMap<List<SegmentIdWithShardSpec>> retVal = new Int2ObjectAVLTreeMap<>();
+
+    // Empty segments validation already happens when the stages are started -- so we cannot have both
+    // isFailOnEmptyInsertEnabled and segmentsToGenerate.isEmpty() be true here.
+    if (segmentsToGenerate == null || segmentsToGenerate.isEmpty()) {
+      return retVal;
+    }
+
+    for (final int workerNumber : workerInputs.workers()) {
+      // SegmentGenerator stage has a single input from another stage.
+      final StageInputSlice stageInputSlice =
+          (StageInputSlice) Iterables.getOnlyElement(workerInputs.inputsForWorker(workerNumber));
+
+      final List<SegmentIdWithShardSpec> workerSegments = new ArrayList<>();
+      retVal.put(workerNumber, workerSegments);
+
+      for (final ReadablePartition partition : stageInputSlice.getPartitions()) {
+        workerSegments.add(segmentsToGenerate.get(partition.getPartitionNumber()));
+      }
+    }
+
+    return retVal;
+  }
+}