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/guice/MSQIndexingModule.java

@@ -84,6 +84,7 @@
 import org.apache.druid.msq.input.table.TableInputSpec;
 import org.apache.druid.msq.kernel.NilExtraInfoHolder;
 import org.apache.druid.msq.querykit.InputNumberDataSource;
+import org.apache.druid.msq.querykit.WindowOperatorQueryFrameProcessorFactory;
 import org.apache.druid.msq.querykit.common.OffsetLimitFrameProcessorFactory;
 import org.apache.druid.msq.querykit.common.SortMergeJoinFrameProcessorFactory;
 import org.apache.druid.msq.querykit.groupby.GroupByPostShuffleFrameProcessorFactory;
@@ -158,7 +159,7 @@ public List<? extends Module> getJacksonModules()
         NilExtraInfoHolder.class,
         SortMergeJoinFrameProcessorFactory.class,
         QueryResultFrameProcessorFactory.class,
-
+        WindowOperatorQueryFrameProcessorFactory.class,
         // DataSource classes (note: ExternalDataSource is in MSQSqlModule)
         InputNumberDataSource.class,
 

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

@@ -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,222 @@
+/*
+ * 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.querykit;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Iterables;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.channel.FrameWithPartition;
+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.util.SettableLongVirtualColumn;
+import org.apache.druid.frame.write.FrameWriter;
+import org.apache.druid.frame.write.FrameWriterFactory;
+import org.apache.druid.java.util.common.Pair;
+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.WindowOperatorQuery;
+import org.apache.druid.query.rowsandcols.LazilyDecoratedRowsAndColumns;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowAndColumns;
+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.VirtualColumn;
+import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.RowSignature;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+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();
+
+      // 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
+
+      // Create a frame rows and columns what would
+      RowBasedFrameRowAndColumns frameRowsAndColumns = new RowBasedFrameRowAndColumns(f, inputFrameReader.signature());
+      Operator op = getOperator(frameRowsAndColumns);
+      //
+      //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());
+          LazilyDecoratedRowsAndColumns tmpldrc = new LazilyDecoratedRowsAndColumns(
+              rac,
+              null,
+              null,
+              null,
+              OffsetLimit.limit(Integer.MAX_VALUE),
+              null,
+              null
+          );
+          Pair<byte[], RowSignature> pairFrames = tmpldrc.naiveMaterializeToRowFrames(rac);
+          Frame f = Frame.wrap(pairFrames.lhs);
+          try {
+            Iterables.getOnlyElement(outputChannels())
+                     .write(new FrameWithPartition(f, FrameWithPartition.NO_PARTITION));
+          }
+          catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+          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();
+  }
+
+  private static Operator getOperator(RowBasedFrameRowAndColumns frameRowsAndColumns)
+  {
+    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;
+      }
+    };
+    return op;
+  }
+}

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

@@ -0,0 +1,73 @@
+/*
+ * 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.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
+    );
+  }
+}