Limit the subquery results by memory usage

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/inline/InlineInputSliceReader.java

@@ -38,7 +38,7 @@
 
 /**
  * Reads {@link InlineInputSlice} using {@link SegmentWrangler} (which is expected to contain an
- * {@link org.apache.druid.segment.InlineSegmentWrangler}).
+ * {@link org.apache.druid.segment.IterableBasedInlineSegmentWrangler}).
  */
 public class InlineInputSliceReader implements InputSliceReader
 {

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

@@ -150,13 +150,6 @@ public void testUnplannableQueries()
 
   }
 
-  @Ignore
-  @Override
-  public void testMaxSubqueryRows()
-  {
-
-  }
-
   @Ignore
   @Override
   public void testQueryWithMoreThanMaxNumericInFilter()

processing/src/main/java/org/apache/druid/frame/field/FieldWriters.java

@@ -28,6 +28,7 @@
 import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.DimensionSelector;
 import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.nested.NestedDataComplexTypeSerde;
 import org.apache.druid.segment.serde.ComplexMetricSerde;
 import org.apache.druid.segment.serde.ComplexMetrics;
 
@@ -57,10 +58,15 @@ private FieldWriters()
   public static FieldWriter create(
       final ColumnSelectorFactory columnSelectorFactory,
       final String columnName,
-      final ColumnType columnType
+      final ColumnType columnType,
+      final boolean allowNullColumnType
   )
   {
     if (columnType == null) {
+      // Returning Complex<Json> writer since we do not know the type of column.
+      if (allowNullColumnType) {
+        return makeComplexWriter(columnSelectorFactory, columnName, NestedDataComplexTypeSerde.TYPE_NAME);
+      }
       throw new UnsupportedColumnTypeException(columnName, null);
     }
 

processing/src/main/java/org/apache/druid/frame/read/FrameReader.java

@@ -71,8 +71,14 @@ private FrameReader(
   /**
    * Create a reader for frames with a given {@link RowSignature}. The signature must exactly match the frames to be
    * read, or else behavior is undefined.
+   * If the columnType is null, we store the data as {@link ColumnType#NESTED_DATA}. This can be done if we know that
+   * the data that we receive can be serded generically using the nested data. It is currently used in the brokers to
+   * store the data with unknown types into frames.
+   * @param signature signature used to generate the reader
+   * @param allowNullTypes if ColumnTypes can be null. The column types that are null would be interpreted
+   *                       as Complex JSON types
    */
-  public static FrameReader create(final RowSignature signature)
+  public static FrameReader create(final RowSignature signature, final boolean allowNullTypes)
   {
     // Double-check that the frame does not have any disallowed field names. Generally, we expect this to be
     // caught on the write side, but we do it again here for safety.
@@ -85,12 +91,17 @@ public static FrameReader create(final RowSignature signature)
     final List<FieldReader> fieldReaders = new ArrayList<>(signature.size());
 
     for (int columnNumber = 0; columnNumber < signature.size(); columnNumber++) {
-      final ColumnType columnType =
-          Preconditions.checkNotNull(
-              signature.getColumnType(columnNumber).orElse(null),
-              "Type for column [%s]",
-              signature.getColumnName(columnNumber)
-          );
+      ColumnType columnType;
+      if (!allowNullTypes) {
+        columnType =
+            Preconditions.checkNotNull(
+                signature.getColumnType(columnNumber).orElse(null),
+                "Type for column [%s]",
+                signature.getColumnName(columnNumber)
+            );
+      } else {
+        columnType = signature.getColumnType(columnNumber).orElse(ColumnType.NESTED_DATA);
+      }
 
       columnReaders.add(FrameColumnReaders.create(columnNumber, columnType));
       fieldReaders.add(FieldReaders.create(signature.getColumnName(columnNumber), columnType));
@@ -99,6 +110,11 @@ public static FrameReader create(final RowSignature signature)
     return new FrameReader(signature, columnReaders, fieldReaders);
   }
 
+  public static FrameReader create(final RowSignature signature)
+  {
+    return create(signature, false);
+  }
+
   public RowSignature signature()
   {
     return signature;

processing/src/main/java/org/apache/druid/frame/segment/FrameCursorUtils.java

@@ -19,10 +19,18 @@
 
 package org.apache.druid.frame.segment;
 
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.write.FrameWriter;
+import org.apache.druid.frame.write.FrameWriterFactory;
+import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.ResourceLimitExceededException;
 import org.apache.druid.query.filter.BoundDimFilter;
 import org.apache.druid.query.filter.Filter;
 import org.apache.druid.query.ordering.StringComparators;
+import org.apache.druid.segment.Cursor;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.filter.BoundFilter;
 import org.apache.druid.segment.filter.Filters;
@@ -67,4 +75,54 @@ public static Filter buildFilter(@Nullable Filter filter, Interval interval)
       );
     }
   }
+
+  /**
+   * Writes a {@link Cursor} to a {@link Frame}. This method iterates over the rows of the cursor, and writes the columns
+   * to the cursor
+   * @param cursor Cursor to write to the frame
+   * @param frameWriterFactory Frame writer factory to write to the frame.
+   *                           Determines the signature of the rows that are written to the frames
+   * @param memoryLimitBytes Limit in bytes, if needs to be enforced while converting the cursor to the frame. If adding
+   *                         a row causes the frame size to exceed this limit, we throw an {@link ResourceLimitExceededException}
+   */
+  public static Frame cursorToFrame(
+      Cursor cursor,
+      FrameWriterFactory frameWriterFactory,
+      @Nullable Long memoryLimitBytes
+  )
+  {
+    Frame frame;
+
+    try (final FrameWriter frameWriter = frameWriterFactory.newFrameWriter(cursor.getColumnSelectorFactory())) {
+      while (!cursor.isDone()) {
+        if (!frameWriter.addSelection()) {
+          throw new ISE(
+              "Unable to append row to the frame. Allocator capacity: [%s]",
+              frameWriterFactory.allocatorCapacity()
+          );
+        }
+
+        if (memoryLimitBytes != null && memoryLimitBytes < frameWriter.getTotalSize()) {
+          throw new ResourceLimitExceededException(
+              StringUtils.format(
+                  "Exceeded total bytes allocated for this subquery. "
+                  + "Current size [%d], total row count [%d], allocated size [%d]. "
+                  + "Please limit the amount of data that the results occupy, or increase the limit using the context "
+                  + "parameter [%s]",
+                  frameWriter.getTotalSize(),
+                  frameWriter.getNumRows(),
+                  memoryLimitBytes,
+                  QueryContexts.MAX_SUBQUERY_MEMORY_BYTES_KEY
+              )
+          );
+        }
+
+        cursor.advance();
+      }
+
+      frame = Frame.wrap(frameWriter.toByteArray());
+    }
+
+    return frame;
+  }
 }

processing/src/main/java/org/apache/druid/frame/write/FrameWriters.java

@@ -29,6 +29,7 @@
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 
+import javax.annotation.Nonnull;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -53,24 +54,38 @@ private FrameWriters()
    * @param signature        signature of the frames
    * @param sortColumns      sort columns for the frames. If nonempty, {@link FrameSort#sort} is used to sort the
    *                         resulting frames.
+   * @param allowNullColumnTypes to allow null ColumnType in the signature. This should only be enabled when the user
+   *                             knows that the column objects exist as native Java POJOs (LinkedList, Maps etc), which
+   *                             can be serded using the Druid's nested columns
    */
   public static FrameWriterFactory makeFrameWriterFactory(
       final FrameType frameType,
       final MemoryAllocatorFactory allocatorFactory,
-      final RowSignature signature,
-      final List<KeyColumn> sortColumns
+      @Nonnull final RowSignature signature,
+      final List<KeyColumn> sortColumns,
+      final boolean allowNullColumnTypes
   )
   {
     switch (Preconditions.checkNotNull(frameType, "frameType")) {
       case COLUMNAR:
-        return new ColumnarFrameWriterFactory(allocatorFactory, signature, sortColumns);
+        return new ColumnarFrameWriterFactory(allocatorFactory, signature, sortColumns, allowNullColumnTypes);
       case ROW_BASED:
-        return new RowBasedFrameWriterFactory(allocatorFactory, signature, sortColumns);
+        return new RowBasedFrameWriterFactory(allocatorFactory, signature, sortColumns, allowNullColumnTypes);
       default:
         throw new ISE("Unrecognized frame type [%s]", frameType);
     }
   }
 
+  public static FrameWriterFactory makeFrameWriterFactory(
+      final FrameType frameType,
+      final MemoryAllocatorFactory allocatorFactory,
+      final RowSignature signature,
+      final List<KeyColumn> sortColumns
+  )
+  {
+    return makeFrameWriterFactory(frameType, allocatorFactory, signature, sortColumns, false);
+  }
+
   /**
    * Returns a copy of "signature" with columns rearranged so the provided sortColumns appear as a prefix.
    * Throws an error if any of the sortColumns are not present in the input signature, or if any of their

processing/src/main/java/org/apache/druid/frame/write/RowBasedFrameWriterFactory.java

@@ -40,16 +40,19 @@ public class RowBasedFrameWriterFactory implements FrameWriterFactory
   private final MemoryAllocatorFactory allocatorFactory;
   private final RowSignature signature;
   private final List<KeyColumn> sortColumns;
+  private final boolean allowNullColumnTypes;
 
   public RowBasedFrameWriterFactory(
       final MemoryAllocatorFactory allocatorFactory,
       final RowSignature signature,
-      final List<KeyColumn> sortColumns
+      final List<KeyColumn> sortColumns,
+      final boolean allowNullColumnTypes
   )
   {
     this.allocatorFactory = allocatorFactory;
     this.signature = signature;
     this.sortColumns = sortColumns;
+    this.allowNullColumnTypes = allowNullColumnTypes;
 
     FrameWriterUtils.verifySortColumns(sortColumns, signature);
   }
@@ -70,7 +73,7 @@ public FrameWriter newFrameWriter(final ColumnSelectorFactory columnSelectorFact
     return new RowBasedFrameWriter(
         signature,
         sortColumns,
-        makeFieldWriters(columnSelectorFactory),
+        makeFieldWriters(columnSelectorFactory, allowNullColumnTypes),
         FrameReaderUtils.makeRowMemorySupplier(columnSelectorFactory, signature),
         rowOrderMemory,
         rowOffsetMemory,
@@ -102,7 +105,7 @@ public FrameType frameType()
    * The returned {@link FieldWriter} objects are not thread-safe, and should only be used with a
    * single frame writer.
    */
-  private List<FieldWriter> makeFieldWriters(final ColumnSelectorFactory columnSelectorFactory)
+  private List<FieldWriter> makeFieldWriters(final ColumnSelectorFactory columnSelectorFactory, final boolean allowNullColumnTypes)
   {
     final List<FieldWriter> fieldWriters = new ArrayList<>();
 
@@ -111,7 +114,7 @@ private List<FieldWriter> makeFieldWriters(final ColumnSelectorFactory columnSel
         final String column = signature.getColumnName(i);
         // note: null type won't work, but we'll get a nice error from FrameColumnWriters.create
         final ColumnType columnType = signature.getColumnType(i).orElse(null);
-        fieldWriters.add(FieldWriters.create(columnSelectorFactory, column, columnType));
+        fieldWriters.add(FieldWriters.create(columnSelectorFactory, column, columnType, allowNullColumnTypes));
       }
     }
     catch (Throwable e) {

processing/src/main/java/org/apache/druid/frame/write/columnar/ColumnarFrameWriterFactory.java

@@ -44,6 +44,7 @@ public class ColumnarFrameWriterFactory implements FrameWriterFactory
   private final MemoryAllocatorFactory allocatorFactory;
   private final RowSignature signature;
   private final List<KeyColumn> keyColumns;
+  private final boolean allowNullTypes;
 
   /**
    * Create a ColumnarFrameWriterFactory.
@@ -57,12 +58,14 @@ public class ColumnarFrameWriterFactory implements FrameWriterFactory
   public ColumnarFrameWriterFactory(
       final MemoryAllocatorFactory allocatorFactory,
       final RowSignature signature,
-      final List<KeyColumn> keyColumns
+      final List<KeyColumn> keyColumns,
+      final boolean allowNullTypes
   )
   {
     this.allocatorFactory = Preconditions.checkNotNull(allocatorFactory, "allocatorFactory");
     this.signature = signature;
     this.keyColumns = Preconditions.checkNotNull(keyColumns, "sortColumns");
+    this.allowNullTypes = allowNullTypes;
 
     if (!keyColumns.isEmpty()) {
       throw new IAE("Columnar frames cannot be sorted");
@@ -86,7 +89,7 @@ public FrameWriter newFrameWriter(final ColumnSelectorFactory columnSelectorFact
         final String column = signature.getColumnName(i);
         // note: null type won't work, but we'll get a nice error from FrameColumnWriters.create
         final ColumnType columnType = signature.getColumnType(i).orElse(null);
-        columnWriters.add(FrameColumnWriters.create(columnSelectorFactory, allocator, column, columnType));
+        columnWriters.add(FrameColumnWriters.create(columnSelectorFactory, allocator, column, columnType, allowNullTypes));
       }
     }
     catch (Throwable e) {

processing/src/main/java/org/apache/druid/frame/write/columnar/FrameColumnWriters.java

@@ -29,6 +29,7 @@
 import org.apache.druid.segment.DimensionSelector;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.nested.NestedDataComplexTypeSerde;
 import org.apache.druid.segment.serde.ComplexMetricSerde;
 import org.apache.druid.segment.serde.ComplexMetrics;
 
@@ -56,10 +57,15 @@ static FrameColumnWriter create(
       final ColumnSelectorFactory columnSelectorFactory,
       final MemoryAllocator allocator,
       final String column,
-      final ColumnType type
+      final ColumnType type,
+      final boolean allowNullType
   )
   {
     if (type == null) {
+      if (allowNullType) {
+        // Serde unknown types as complex
+        return makeComplexWriter(columnSelectorFactory, allocator, column, NestedDataComplexTypeSerde.TYPE_NAME);
+      }
       throw new UnsupportedColumnTypeException(column, null);
     }
 

processing/src/main/java/org/apache/druid/jackson/DruidDefaultSerializersModule.java

@@ -32,6 +32,8 @@
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Yielder;
 import org.apache.druid.java.util.common.jackson.JacksonUtils;
+import org.apache.druid.query.FramesBackedInlineDataSource;
+import org.apache.druid.query.FramesBackedInlineDataSourceSerializer;
 import org.apache.druid.query.context.ResponseContext;
 import org.apache.druid.query.context.ResponseContextDeserializer;
 import org.joda.time.DateTimeZone;
@@ -52,6 +54,8 @@ public DruidDefaultSerializersModule()
 
     JodaStuff.register(this);
 
+    addSerializer(FramesBackedInlineDataSource.class, new FramesBackedInlineDataSourceSerializer());
+
     addDeserializer(
         DateTimeZone.class,
         new JsonDeserializer<DateTimeZone>()