Add configurable final stages to MSQ ingestion queries

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

@@ -146,6 +146,7 @@
 import org.apache.druid.msq.input.stage.StageInputSpec;
 import org.apache.druid.msq.input.stage.StageInputSpecSlicer;
 import org.apache.druid.msq.input.table.TableInputSpec;
+import org.apache.druid.msq.kernel.FrameProcessorFactory;
 import org.apache.druid.msq.kernel.QueryDefinition;
 import org.apache.druid.msq.kernel.QueryDefinitionBuilder;
 import org.apache.druid.msq.kernel.StageDefinition;
@@ -1193,8 +1194,35 @@ private Int2ObjectMap<Object> makeWorkerFactoryInfosForStage(
   {
     if (MSQControllerTask.isIngestion(querySpec) &&
         stageNumber == queryDef.getFinalStageDefinition().getStageNumber()) {
-      // noinspection unchecked,rawtypes
-      return (Int2ObjectMap) makeSegmentGeneratorWorkerFactoryInfos(workerInputs, segmentsToGenerate);
+      final DataSourceMSQDestination destination = (DataSourceMSQDestination) querySpec.getDestination();
+      if (!destination.doesSegmentMorphing()) {
+        // noinspection unchecked,rawtypes
+        return (Int2ObjectMap) makeSegmentGeneratorWorkerFactoryInfos(workerInputs, segmentsToGenerate);
+      } else {
+        // worker info is the new lock version
+        if (destination.getReplaceTimeChunks().size() != 1) {
+          throw new ISE(
+              "Must have single interval in replaceTimeChunks, but got[%s]",
+              destination.getReplaceTimeChunks()
+          );
+        }
+        try {
+          final List<TaskLock> locks;
+          locks = context.taskActionClient().submit(new LockListAction());
+          if (locks.size() == 1) {
+            final Int2ObjectMap<Object> retVal = new Int2ObjectAVLTreeMap<>();
+            for (int worker : workerInputs.workers()) {
+              retVal.put(worker, locks.get(0).getVersion());
+            }
+            return retVal;
+          } else {
+            throw new ISE("Got number of locks other than one: [%s]", locks);
+          }
+        }
+        catch (IOException e) {
+          throw new RuntimeException(e);
+        }
+      }
     } else {
       return null;
     }
@@ -1810,6 +1838,35 @@ private static QueryDefinition makeQueryDefinition(
         }
       }
 
+      // Possibly add a segment morpher stage.
+      if (((DataSourceMSQDestination) querySpec.getDestination()).doesSegmentMorphing()) {
+        final DataSourceMSQDestination destination = (DataSourceMSQDestination) querySpec.getDestination();
+        final FrameProcessorFactory segmentMorphFactory = destination.getSegmentMorphFactory();
+
+        if (!destination.isReplaceTimeChunks()) {
+          throw new MSQException(UnknownFault.forMessage("segmentMorphFactory requires replaceTimeChunks"));
+        }
+
+        builder.add(
+            StageDefinition.builder(queryDef.getNextStageNumber())
+                           .inputs(
+                               new TableInputSpec(
+                                   destination.getDataSource(),
+                                   destination.getReplaceTimeChunks(),
+                                   null,
+                                   null
+                               ),
+                               new StageInputSpec(queryDef.getFinalStageDefinition().getStageNumber())
+                           )
+                           .broadcastInputs(IntSet.of(1))
+                           .maxWorkerCount(tuningConfig.getMaxNumWorkers())
+                           .processorFactory(segmentMorphFactory)
+        );
+
+        // If there was a segment morpher, return immediately; don't add a segment-generation stage.
+        return builder.build();
+      }
+
       // Then, add a segment-generation stage.
       final DataSchema dataSchema =
           makeDataSchemaForIngestion(querySpec, querySignature, queryClusterBy, columnMappings, jsonMapper);
@@ -2723,7 +2780,8 @@ private void startStages() throws IOException, InterruptedException
         for (final StageId stageId : newStageIds) {
           // Allocate segments, if this is the final stage of an ingestion.
           if (MSQControllerTask.isIngestion(querySpec)
-              && stageId.getStageNumber() == queryDef.getFinalStageDefinition().getStageNumber()) {
+              && stageId.getStageNumber() == queryDef.getFinalStageDefinition().getStageNumber()
+              && !((DataSourceMSQDestination) querySpec.getDestination()).doesSegmentMorphing()) {
             populateSegmentsToGenerate();
           }
 

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

@@ -20,12 +20,14 @@
 package org.apache.druid.msq.indexing.destination;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.granularity.Granularity;
+import org.apache.druid.msq.kernel.FrameProcessorFactory;
 import org.apache.druid.msq.querykit.ShuffleSpecFactories;
 import org.apache.druid.msq.querykit.ShuffleSpecFactory;
 import org.apache.druid.server.security.Resource;
@@ -49,18 +51,24 @@ public class DataSourceMSQDestination implements MSQDestination
   @Nullable
   private final List<Interval> replaceTimeChunks;
 
+  @Nullable
+  @SuppressWarnings("rawtypes")
+  private final FrameProcessorFactory segmentMorphFactory;
+
   @JsonCreator
   public DataSourceMSQDestination(
       @JsonProperty("dataSource") String dataSource,
       @JsonProperty("segmentGranularity") Granularity segmentGranularity,
       @JsonProperty("segmentSortOrder") @Nullable List<String> segmentSortOrder,
-      @JsonProperty("replaceTimeChunks") @Nullable List<Interval> replaceTimeChunks
+      @JsonProperty("replaceTimeChunks") @Nullable List<Interval> replaceTimeChunks,
+      @JsonProperty("segmentMorphFactory") @Nullable FrameProcessorFactory segmentMorphFactory
   )
   {
     this.dataSource = Preconditions.checkNotNull(dataSource, "dataSource");
     this.segmentGranularity = Preconditions.checkNotNull(segmentGranularity, "segmentGranularity");
     this.segmentSortOrder = segmentSortOrder != null ? segmentSortOrder : Collections.emptyList();
     this.replaceTimeChunks = replaceTimeChunks;
+    this.segmentMorphFactory = segmentMorphFactory;
 
     if (replaceTimeChunks != null) {
       // Verify that if replaceTimeChunks is provided, it is nonempty.
@@ -98,6 +106,30 @@ public String getDataSource()
     return dataSource;
   }
 
+  /**
+   * Returns the segment morph factory, if one is present, else null.
+   * <p>
+   * The segment morph factory if present, is a way to tell the MSQ task to funnel the results at the final stage to
+   * the {@link FrameProcessorFactory} instead of a segment generation stage.
+   */
+  @Nullable
+  @JsonProperty
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public FrameProcessorFactory getSegmentMorphFactory()
+  {
+    return segmentMorphFactory;
+  }
+
+  /**
+   * Checks if the destination uses a segmentMorphFactory. If one is present, that means that the query would modify
+   * existing segments instead of generating new ones.
+   */
+  @JsonIgnore
+  public boolean doesSegmentMorphing()
+  {
+    return segmentMorphFactory != null;
+  }
+
   @JsonProperty
   public Granularity getSegmentGranularity()
   {
@@ -158,13 +190,14 @@ public boolean equals(Object o)
     return Objects.equals(dataSource, that.dataSource)
            && Objects.equals(segmentGranularity, that.segmentGranularity)
            && Objects.equals(segmentSortOrder, that.segmentSortOrder)
-           && Objects.equals(replaceTimeChunks, that.replaceTimeChunks);
+           && Objects.equals(replaceTimeChunks, that.replaceTimeChunks)
+           && Objects.equals(segmentMorphFactory, that.segmentMorphFactory);
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(dataSource, segmentGranularity, segmentSortOrder, replaceTimeChunks);
+    return Objects.hash(dataSource, segmentGranularity, segmentSortOrder, replaceTimeChunks, segmentMorphFactory);
   }
 
   @Override
@@ -175,6 +208,7 @@ public String toString()
            ", segmentGranularity=" + segmentGranularity +
            ", segmentSortOrder=" + segmentSortOrder +
            ", replaceTimeChunks=" + replaceTimeChunks +
+           (segmentMorphFactory != null ? ", segmentMorphFactory=" + segmentMorphFactory : "") +
            '}';
   }
 

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskQueryMaker.java

@@ -246,7 +246,8 @@ public QueryResponse<Object[]> runQuery(final DruidQuery druidQuery)
           targetDataSource.getDestinationName(),
           segmentGranularityObject,
           segmentSortOrder,
-          replaceTimeChunks
+          replaceTimeChunks,
+          null
       );
       MultiStageQueryContext.validateAndGetTaskLockType(sqlQueryContext,
                                                         dataSourceMSQDestination.isReplaceTimeChunks());

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQParseExceptionsTest.java

@@ -225,7 +225,7 @@ public void testIngestWithSanitizedNullByte() throws IOException
                         new ColumnMapping("v1", "agent_category")
                     )
                 ))
-                .destination(new DataSourceMSQDestination("foo1", Granularities.ALL, null, null))
+                .destination(new DataSourceMSQDestination("foo1", Granularities.ALL, null, null, null))
                 .tuningConfig(MSQTuningConfig.defaultConfig())
                 .build())
         .setQueryContext(DEFAULT_MSQ_CONTEXT)
@@ -318,7 +318,7 @@ public void testIngestWithSanitizedNullByteUsingContextParameter() throws IOExce
                         new ColumnMapping("agent_category", "agent_category")
                     )
                 ))
-                .destination(new DataSourceMSQDestination("foo1", Granularities.ALL, null, null))
+                .destination(new DataSourceMSQDestination("foo1", Granularities.ALL, null, null, null))
                 .tuningConfig(MSQTuningConfig.defaultConfig())
                 .build())
         .setQueryContext(runtimeContext)

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

@@ -58,7 +58,8 @@ public class MSQControllerTaskTest
           "target",
           Granularities.DAY,
           null,
-          INTERVALS
+          INTERVALS,
+          null
       ))
       .query(new Druids.ScanQueryBuilder()
                  .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java

@@ -199,6 +199,7 @@ public class SqlStatementResourceTest extends MSQTestBase
                  "test",
                  Granularities.DAY,
                  null,
+                 null,
                  null
              ))
              .tuningConfig(

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java

@@ -375,6 +375,7 @@ public void testEmptyCountersForDataSourceDestination()
             "test",
             Granularities.DAY,
             null,
+            null,
             null
         )
     );