Window function on msq

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

@@ -166,6 +166,7 @@
 import org.apache.druid.msq.querykit.QueryKitUtils;
 import org.apache.druid.msq.querykit.ShuffleSpecFactories;
 import org.apache.druid.msq.querykit.ShuffleSpecFactory;
+import org.apache.druid.msq.querykit.WindowOperatorQueryKit;
 import org.apache.druid.msq.querykit.groupby.GroupByQueryKit;
 import org.apache.druid.msq.querykit.results.QueryResultFrameProcessorFactory;
 import org.apache.druid.msq.querykit.scan.ScanQueryKit;
@@ -184,6 +185,7 @@
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.groupby.GroupByQuery;
 import org.apache.druid.query.groupby.GroupByQueryConfig;
+import org.apache.druid.query.operator.WindowOperatorQuery;
 import org.apache.druid.query.scan.ScanQuery;
 import org.apache.druid.segment.DimensionHandlerUtils;
 import org.apache.druid.segment.column.ColumnHolder;
@@ -1172,6 +1174,7 @@ private QueryKit makeQueryControllerToolKit()
         ImmutableMap.<Class<? extends Query>, QueryKit>builder()
                     .put(ScanQuery.class, new ScanQueryKit(context.jsonMapper()))
                     .put(GroupByQuery.class, new GroupByQueryKit(context.jsonMapper()))
+                    .put(WindowOperatorQuery.class, new WindowOperatorQueryKit(context.jsonMapper()))
                     .build();
 
     return new MultiQueryKit(kitMap);

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java

@@ -40,7 +40,7 @@
 import java.util.List;
 import java.util.function.Function;
 
-public abstract class BaseLeafFrameProcessor implements FrameProcessor<Object>
+public abstract class  BaseLeafFrameProcessor implements FrameProcessor<Object>
 {
   private final ReadableInput baseInput;
   private final ResourceHolder<WritableFrameChannel> outputChannelHolder;
@@ -63,6 +63,7 @@ protected BaseLeafFrameProcessor(
   @Override
   public List<ReadableFrameChannel> inputChannels()
   {
+    // somu: need to clarify if the data is in broker only
     if (baseInput.hasSegment()) {
       return Collections.emptyList();
     } else {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java

@@ -0,0 +1,172 @@
+package org.apache.druid.msq.querykit;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.channel.ReadableFrameChannel;
+import org.apache.druid.frame.channel.WritableFrameChannel;
+import org.apache.druid.frame.processor.ReturnOrAwait;
+import org.apache.druid.frame.read.FrameReader;
+import org.apache.druid.frame.segment.FrameSegment;
+import org.apache.druid.frame.util.SettableLongVirtualColumn;
+import org.apache.druid.frame.write.FrameWriter;
+import org.apache.druid.frame.write.FrameWriterFactory;
+import org.apache.druid.java.util.common.Unit;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.input.ReadableInput;
+import org.apache.druid.msq.input.table.SegmentWithDescriptor;
+import org.apache.druid.msq.querykit.scan.ScanQueryFrameProcessor;
+import org.apache.druid.query.operator.OffsetLimit;
+import org.apache.druid.query.operator.Operator;
+import org.apache.druid.query.operator.OperatorFactory;
+import org.apache.druid.query.operator.SegmentToRowsAndColumnsOperator;
+import org.apache.druid.query.operator.WindowOperatorQuery;
+import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowAndColumns;
+import org.apache.druid.query.scan.ScanQuery;
+import org.apache.druid.segment.Cursor;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.SegmentReference;
+import org.apache.druid.segment.SimpleAscendingOffset;
+import org.apache.druid.segment.SimpleSettableOffset;
+import org.apache.druid.segment.StorageAdapter;
+import org.apache.druid.segment.VirtualColumn;
+import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.timeline.SegmentId;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
+
+public class WindowOperatorQueryFrameProcessor extends BaseLeafFrameProcessor
+{
+
+  private static final Logger log = new Logger(ScanQueryFrameProcessor.class);
+  private final WindowOperatorQuery query;
+  private final ObjectMapper jsonMapper;
+  private final SettableLongVirtualColumn partitionBoostVirtualColumn;
+  private final VirtualColumns frameWriterVirtualColumns;
+  private final Closer closer = Closer.create();
+
+  private Cursor cursor;
+  private Segment segment;
+  private final SimpleSettableOffset cursorOffset = new SimpleAscendingOffset(Integer.MAX_VALUE);
+  private FrameWriter frameWriter;
+  private long currentAllocatorCapacity; // Used for generating FrameRowTooLargeException if needed
+
+  public WindowOperatorQueryFrameProcessor(
+      final WindowOperatorQuery query,
+      final ObjectMapper jsonMapper,
+      final ReadableInput baseInput,
+      final Function<SegmentReference, SegmentReference> segmentMapFn,
+      final ResourceHolder<WritableFrameChannel> outputChannelHolder,
+      final ResourceHolder<FrameWriterFactory> frameWriterFactoryHolder
+  )
+  {
+    super(
+        baseInput,
+        segmentMapFn,
+        outputChannelHolder,
+        frameWriterFactoryHolder
+    );
+    this.query = query;
+    this.jsonMapper = jsonMapper;
+    this.partitionBoostVirtualColumn = new SettableLongVirtualColumn(QueryKitUtils.PARTITION_BOOST_COLUMN);
+
+    final List<VirtualColumn> frameWriterVirtualColumns = new ArrayList<>();
+    frameWriterVirtualColumns.add(partitionBoostVirtualColumn);
+
+    final VirtualColumn segmentGranularityVirtualColumn =
+        QueryKitUtils.makeSegmentGranularityVirtualColumn(jsonMapper, query);
+
+    if (segmentGranularityVirtualColumn != null) {
+      frameWriterVirtualColumns.add(segmentGranularityVirtualColumn);
+    }
+
+    this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns);
+  }
+  // deep storage
+  @Override
+  protected ReturnOrAwait<Unit> runWithSegment(SegmentWithDescriptor segment) throws IOException
+  {
+    return null;
+  }
+
+  // realtime
+  @Override
+  protected ReturnOrAwait<Unit> runWithLoadedSegment(SegmentWithDescriptor segment) throws IOException
+  {
+    return null;
+  }
+
+  // previous stage output
+  @Override
+  protected ReturnOrAwait<Unit> runWithInputChannel(ReadableFrameChannel inputChannel, FrameReader inputFrameReader)
+      throws IOException
+  {
+    // Read the frames from the channel
+    // convert to FrameRowsAndColumns
+
+   if (inputChannel.canRead()) {
+     Frame f = inputChannel.read();
+     final FrameSegment frameSegment = new FrameSegment(f, inputFrameReader, SegmentId.dummy("x"));
+
+     // the frame here is row based
+     // frame rows and columns need columnar. discuss with Eric
+     // Action item: need to implement a new rows and columns that accept a row-based frame
+     RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature());
+     LazilyDecoratedRowsAndColumns ldrc = new LazilyDecoratedRowsAndColumns(frameRowsAndColumns, null, null, null, OffsetLimit.limit(Integer.MAX_VALUE), null, null);
+     // Create an operator on top of the created rows and columns
+     Operator op = new Operator()
+     {
+       @Nullable
+       @Override
+       public Closeable goOrContinue(Closeable continuationObject, Receiver receiver)
+       {
+         receiver.push(ldrc);
+         receiver.completed();
+         return continuationObject;
+       }
+     };
+     //
+     //Operator op = new SegmentToRowsAndColumnsOperator(frameSegment);
+     // On the operator created above add the operators present in the query that we want to chain
+
+     for ( OperatorFactory of : query.getOperators()) {
+       op = of.wrap(op);
+     }
+
+     // Let the operator run
+     // the results that come in the receiver must be pushed to the outout channel
+     // need to transform the output rows and columns back to frame
+     Operator.go(op, new Operator.Receiver()
+     {
+       @Override
+       public Operator.Signal push(RowsAndColumns rac)
+       {
+         //outputFrameChannel.output(rac.toFrame());
+         return Operator.Signal.GO;
+       }
+
+       @Override
+       public void completed()
+       {
+
+       }
+     });
+
+   } else if (inputChannel.isFinished()) {
+     return ReturnOrAwait.returnObject(Unit.instance());
+   } else {
+     return ReturnOrAwait.awaitAll(inputChannels().size());
+   }
+    return ReturnOrAwait.runAgain();
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessorFactory.java

@@ -0,0 +1,54 @@
+package org.apache.druid.msq.querykit;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonTypeName;
+import com.google.common.base.Preconditions;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.frame.channel.WritableFrameChannel;
+import org.apache.druid.frame.processor.FrameProcessor;
+import org.apache.druid.frame.write.FrameWriterFactory;
+import org.apache.druid.msq.input.ReadableInput;
+import org.apache.druid.msq.kernel.FrameContext;
+import org.apache.druid.query.operator.WindowOperatorQuery;
+import org.apache.druid.segment.SegmentReference;
+
+import java.util.function.Function;
+
+@JsonTypeName("window")
+public class WindowOperatorQueryFrameProcessorFactory extends BaseLeafFrameProcessorFactory
+{
+  private final WindowOperatorQuery query;
+
+  @JsonCreator
+  public WindowOperatorQueryFrameProcessorFactory(@JsonProperty("query") WindowOperatorQuery query)
+  {
+    super(query);
+    this.query = Preconditions.checkNotNull(query, "query");
+  }
+
+  @JsonProperty
+  public WindowOperatorQuery getQuery()
+  {
+    return query;
+  }
+
+  @Override
+  protected FrameProcessor<Object> makeProcessor(
+      ReadableInput baseInput,
+      Function<SegmentReference, SegmentReference> segmentMapFn,
+      ResourceHolder<WritableFrameChannel> outputChannelHolder,
+      ResourceHolder<FrameWriterFactory> frameWriterFactoryHolder,
+      FrameContext frameContext
+  )
+  {
+    return new WindowOperatorQueryFrameProcessor(
+        query,
+        frameContext.jsonMapper(),
+        baseInput,
+        segmentMapFn,
+        outputChannelHolder,
+        frameWriterFactoryHolder
+    );
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java

@@ -0,0 +1,64 @@
+package org.apache.druid.msq.querykit;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.msq.input.stage.StageInputSpec;
+import org.apache.druid.msq.kernel.QueryDefinition;
+import org.apache.druid.msq.kernel.QueryDefinitionBuilder;
+import org.apache.druid.msq.kernel.StageDefinition;
+import org.apache.druid.query.Query;
+import org.apache.druid.query.operator.WindowOperatorQuery;
+import org.apache.druid.segment.column.RowSignature;
+
+public class WindowOperatorQueryKit implements QueryKit<WindowOperatorQuery>
+{
+  private final ObjectMapper jsonMapper;
+
+  public WindowOperatorQueryKit(ObjectMapper jsonMapper)
+  {
+    this.jsonMapper = jsonMapper;
+  }
+  @Override
+  public QueryDefinition makeQueryDefinition(
+      String queryId,
+      WindowOperatorQuery originalQuery,
+      QueryKit<Query<?>> queryKit,
+      ShuffleSpecFactory resultShuffleSpecFactory,
+      int maxWorkerCount,
+      int minStageNumber
+  )
+  {
+    // need to validate query first
+
+    final QueryDefinitionBuilder queryDefBuilder = QueryDefinition.builder().queryId(queryId);
+    final DataSourcePlan dataSourcePlan = DataSourcePlan.forDataSource(
+        queryKit,
+        queryId,
+        originalQuery.context(),
+        originalQuery.getDataSource(),
+        originalQuery.getQuerySegmentSpec(),
+        originalQuery.getFilter(),
+        null,
+        maxWorkerCount,
+        minStageNumber,
+        false
+    );
+
+
+    dataSourcePlan.getSubQueryDefBuilder().ifPresent(queryDefBuilder::addAll);
+
+    final int firstStageNumber = Math.max(minStageNumber, queryDefBuilder.getNextStageNumber());
+    final WindowOperatorQuery queryToRun = (WindowOperatorQuery) originalQuery.withDataSource(dataSourcePlan.getNewDataSource());
+    RowSignature rowSignature = queryToRun.getRowSignature();
+
+    // Create a new stage which takes in the subquery as an input
+    queryDefBuilder.add(
+        StageDefinition.builder(firstStageNumber)
+                       .inputs(new StageInputSpec(firstStageNumber-1))
+                       .signature(rowSignature)
+                       .maxWorkerCount(maxWorkerCount)
+                       .shuffleSpec(null)
+                       .processorFactory(new WindowOperatorQueryFrameProcessorFactory(queryToRun))
+    );
+    return queryDefBuilder.queryId(queryId).build();
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java

@@ -149,18 +149,23 @@ public ReturnOrAwait<Object> runIncrementally(final IntSet readableInputs) throw
         }
 
         final int fullRowSize = query.getResultRowSignature().size();
-        rowSupplierFromFrameCursor = () -> {
-          final ResultRow row = ResultRow.create(fullRowSize);
-          for (int i = 0; i < fieldSuppliers.length; i++) {
-            row.set(i, fieldSuppliers[i].get());
+        rowSupplierFromFrameCursor = new Supplier<ResultRow>()
+        {
+          @Override
+          public ResultRow get()
+          {
+            final ResultRow row = ResultRow.create(fullRowSize);
+            for (int i = 0; i < fieldSuppliers.length; i++) {
+              row.set(i, fieldSuppliers[i].get());
+            }
+
+            for (int i = fieldSuppliers.length; i < fullRowSize; i++) {
+              // Post-aggregators.
+              row.set(i, null);
+            }
+
+            return row;
           }
-
-          for (int i = fieldSuppliers.length; i < fullRowSize; i++) {
-            // Post-aggregators.
-            row.set(i, null);
-          }
-
-          return row;
         };
       }
     }

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java

@@ -194,6 +194,9 @@ private static ScanQuery prepareScanQueryForDataServer(@NotNull ScanQuery scanQu
   @Override
   protected ReturnOrAwait<Unit> runWithLoadedSegment(final SegmentWithDescriptor segment) throws IOException
   {
+    // Try to run with the segment on the server
+    // if segment was not found for some reason
+    // call runWithSegment which would load from deep storage
     if (cursor == null) {
       ScanQuery preparedQuery = prepareScanQueryForDataServer(query);
       final Pair<LoadedSegmentDataProvider.DataServerQueryStatus, Yielder<Object[]>> statusSequencePair =
@@ -244,7 +247,9 @@ protected ReturnOrAwait<Unit> runWithLoadedSegment(final SegmentWithDescriptor s
   @Override
   protected ReturnOrAwait<Unit> runWithSegment(final SegmentWithDescriptor segment) throws IOException
   {
+    // why not remove order by and limit here ?
     if (cursor == null) {
+      // load from deep storage
       final ResourceHolder<Segment> segmentHolder = closer.register(segment.getOrLoad());
 
       final Yielder<Cursor> cursorYielder = Yielders.each(

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

@@ -113,9 +113,9 @@ public boolean featureAvailable(EngineFeature feature, PlannerContext plannerCon
       case TOPN_QUERY:
       case TIME_BOUNDARY_QUERY:
       case GROUPING_SETS:
-      case WINDOW_FUNCTIONS:
       case ALLOW_TOP_LEVEL_UNION_ALL:
         return false;
+      case WINDOW_FUNCTIONS:
       case UNNEST:
       case CAN_SELECT:
       case CAN_INSERT: