Revert "[SPARK-8579] [SQL] support arbitrary object in UnsafeRow"

This reverts commit ed359de.

Conflicts:
	sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeFixedWidthAggregationMap.java
	sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
	sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnsafeRowConverter.scala
	sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/UnsafeFixedWidthAggregationMapSuite.scala
	sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/UnsafeRowConverterSuite.scala
	sql/core/src/main/scala/org/apache/spark/sql/execution/GeneratedAggregate.scala

Author: rxin

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeFixedWidthAggregationMap.java

@@ -19,11 +19,9 @@
 
 import java.util.Iterator;
 
-import scala.Function1;
-
 import org.apache.spark.sql.catalyst.InternalRow;
-import org.apache.spark.sql.catalyst.util.ObjectPool;
-import org.apache.spark.sql.catalyst.util.UniqueObjectPool;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
 import org.apache.spark.unsafe.PlatformDependent;
 import org.apache.spark.unsafe.map.BytesToBytesMap;
 import org.apache.spark.unsafe.memory.MemoryLocation;
@@ -40,48 +38,26 @@ public final class UnsafeFixedWidthAggregationMap {
    * An empty aggregation buffer, encoded in UnsafeRow format. When inserting a new key into the
    * map, we copy this buffer and use it as the value.
    */
-  private final byte[] emptyBuffer;
+  private final byte[] emptyAggregationBuffer;
 
-  /**
-   * An empty row used by `initProjection`
-   */
-  private static final InternalRow emptyRow = new GenericInternalRow();
+  private final StructType aggregationBufferSchema;
 
-  /**
-   * Whether can the empty aggregation buffer be reuse without calling `initProjection` or not.
-   */
-  private final boolean reuseEmptyBuffer;
+  private final StructType groupingKeySchema;
 
   /**
-   * The projection used to initialize the emptyBuffer
+   * Encodes grouping keys as UnsafeRows.
    */
-  private final Function1<InternalRow, InternalRow> initProjection;
-
-  /**
-   * Encodes grouping keys or buffers as UnsafeRows.
-   */
-  private final UnsafeRowConverter keyConverter;
-  private final UnsafeRowConverter bufferConverter;
+  private final UnsafeRowConverter groupingKeyToUnsafeRowConverter;
 
   /**
    * A hashmap which maps from opaque bytearray keys to bytearray values.
    */
   private final BytesToBytesMap map;
 
-  /**
-   * An object pool for objects that are used in grouping keys.
-   */
-  private final UniqueObjectPool keyPool;
-
-  /**
-   * An object pool for objects that are used in aggregation buffers.
-   */
-  private final ObjectPool bufferPool;
-
   /**
    * Re-used pointer to the current aggregation buffer
    */
-  private final UnsafeRow currentBuffer = new UnsafeRow();
+  private final UnsafeRow currentAggregationBuffer = new UnsafeRow();
 
   /**
    * Scratch space that is used when encoding grouping keys into UnsafeRow format.
@@ -93,59 +69,86 @@ public final class UnsafeFixedWidthAggregationMap {
 
   private final boolean enablePerfMetrics;
 
+  /**
+   * @return true if UnsafeFixedWidthAggregationMap supports grouping keys with the given schema,
+   *         false otherwise.
+   */
+  public static boolean supportsGroupKeySchema(StructType schema) {
+    for (StructField field: schema.fields()) {
+      if (!UnsafeRow.readableFieldTypes.contains(field.dataType())) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /**
+   * @return true if UnsafeFixedWidthAggregationMap supports aggregation buffers with the given
+   *         schema, false otherwise.
+   */
+  public static boolean supportsAggregationBufferSchema(StructType schema) {
+    for (StructField field: schema.fields()) {
+      if (!UnsafeRow.settableFieldTypes.contains(field.dataType())) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   /**
    * Create a new UnsafeFixedWidthAggregationMap.
    *
-   * @param initProjection the default value for new keys (a "zero" of the agg. function)
-   * @param keyConverter the converter of the grouping key, used for row conversion.
-   * @param bufferConverter the converter of the aggregation buffer, used for row conversion.
+   * @param emptyAggregationBuffer the default value for new keys (a "zero" of the agg. function)
+   * @param aggregationBufferSchema the schema of the aggregation buffer, used for row conversion.
+   * @param groupingKeySchema the schema of the grouping key, used for row conversion.
    * @param memoryManager the memory manager used to allocate our Unsafe memory structures.
    * @param initialCapacity the initial capacity of the map (a sizing hint to avoid re-hashing).
    * @param enablePerfMetrics if true, performance metrics will be recorded (has minor perf impact)
    */
   public UnsafeFixedWidthAggregationMap(
-      Function1<InternalRow, InternalRow> initProjection,
-      UnsafeRowConverter keyConverter,
-      UnsafeRowConverter bufferConverter,
+      InternalRow emptyAggregationBuffer,
+      StructType aggregationBufferSchema,
+      StructType groupingKeySchema,
       TaskMemoryManager memoryManager,
       int initialCapacity,
       boolean enablePerfMetrics) {
-    this.initProjection = initProjection;
-    this.keyConverter = keyConverter;
-    this.bufferConverter = bufferConverter;
-    this.enablePerfMetrics = enablePerfMetrics;
-
+    this.emptyAggregationBuffer =
+      convertToUnsafeRow(emptyAggregationBuffer, aggregationBufferSchema);
+    this.aggregationBufferSchema = aggregationBufferSchema;
+    this.groupingKeyToUnsafeRowConverter = new UnsafeRowConverter(groupingKeySchema);
+    this.groupingKeySchema = groupingKeySchema;
     this.map = new BytesToBytesMap(memoryManager, initialCapacity, enablePerfMetrics);
-    this.keyPool = new UniqueObjectPool(100);
-    this.bufferPool = new ObjectPool(initialCapacity);
+    this.enablePerfMetrics = enablePerfMetrics;
+  }
 
-    InternalRow initRow = initProjection.apply(emptyRow);
-    int emptyBufferSize = bufferConverter.getSizeRequirement(initRow);
-    this.emptyBuffer = new byte[emptyBufferSize];
-    int writtenLength = bufferConverter.writeRow(
-      initRow, emptyBuffer, PlatformDependent.BYTE_ARRAY_OFFSET, emptyBufferSize,
-      bufferPool);
-    assert (writtenLength == emptyBuffer.length): "Size requirement calculation was wrong!";
-    // re-use the empty buffer only when there is no object saved in pool.
-    reuseEmptyBuffer = bufferPool.size() == 0;
+  /**
+   * Convert a Java object row into an UnsafeRow, allocating it into a new byte array.
+   */
+  private static byte[] convertToUnsafeRow(InternalRow javaRow, StructType schema) {
+    final UnsafeRowConverter converter = new UnsafeRowConverter(schema);
+    final int size = converter.getSizeRequirement(javaRow);
+    final byte[] unsafeRow = new byte[size];
+    final int writtenLength =
+      converter.writeRow(javaRow, unsafeRow, PlatformDependent.BYTE_ARRAY_OFFSET, size);
+    assert (writtenLength == unsafeRow.length): "Size requirement calculation was wrong!";
+    return unsafeRow;
   }
 
   /**
    * Return the aggregation buffer for the current group. For efficiency, all calls to this method
    * return the same object.
    */
   public UnsafeRow getAggregationBuffer(InternalRow groupingKey) {
-    final int groupingKeySize = keyConverter.getSizeRequirement(groupingKey);
+    final int groupingKeySize = groupingKeyToUnsafeRowConverter.getSizeRequirement(groupingKey);
     // Make sure that the buffer is large enough to hold the key. If it's not, grow it:
     if (groupingKeySize > groupingKeyConversionScratchSpace.length) {
       groupingKeyConversionScratchSpace = new byte[groupingKeySize];
     }
-    final int actualGroupingKeySize = keyConverter.writeRow(
+    final int actualGroupingKeySize = groupingKeyToUnsafeRowConverter.writeRow(
       groupingKey,
       groupingKeyConversionScratchSpace,
       PlatformDependent.BYTE_ARRAY_OFFSET,
-      groupingKeySize,
-      keyPool);
+      groupingKeySize);
     assert (groupingKeySize == actualGroupingKeySize) : "Size requirement calculation was wrong!";
 
     // Probe our map using the serialized key
@@ -156,32 +159,25 @@ public UnsafeRow getAggregationBuffer(InternalRow groupingKey) {
     if (!loc.isDefined()) {
       // This is the first time that we've seen this grouping key, so we'll insert a copy of the
       // empty aggregation buffer into the map:
-      if (!reuseEmptyBuffer) {
-        // There is some objects referenced by emptyBuffer, so generate a new one
-        InternalRow initRow = initProjection.apply(emptyRow);
-        bufferConverter.writeRow(initRow, emptyBuffer, PlatformDependent.BYTE_ARRAY_OFFSET,
-          groupingKeySize, bufferPool);
-      }
       loc.putNewKey(
         groupingKeyConversionScratchSpace,
         PlatformDependent.BYTE_ARRAY_OFFSET,
         groupingKeySize,
-        emptyBuffer,
+        emptyAggregationBuffer,
         PlatformDependent.BYTE_ARRAY_OFFSET,
-        emptyBuffer.length
+        emptyAggregationBuffer.length
       );
     }
 
     // Reset the pointer to point to the value that we just stored or looked up:
     final MemoryLocation address = loc.getValueAddress();
-    currentBuffer.pointTo(
+    currentAggregationBuffer.pointTo(
       address.getBaseObject(),
       address.getBaseOffset(),
-      bufferConverter.numFields(),
-      loc.getValueLength(),
-      bufferPool
+      aggregationBufferSchema.length(),
+      loc.getValueLength()
     );
-    return currentBuffer;
+    return currentAggregationBuffer;
   }
 
   /**
@@ -217,16 +213,14 @@ public MapEntry next() {
         entry.key.pointTo(
           keyAddress.getBaseObject(),
           keyAddress.getBaseOffset(),
-          keyConverter.numFields(),
-          loc.getKeyLength(),
-          keyPool
+          groupingKeySchema.length(),
+          loc.getKeyLength()
         );
         entry.value.pointTo(
           valueAddress.getBaseObject(),
           valueAddress.getBaseOffset(),
-          bufferConverter.numFields(),
-          loc.getValueLength(),
-          bufferPool
+          aggregationBufferSchema.length(),
+          loc.getValueLength()
         );
         return entry;
       }
@@ -254,8 +248,6 @@ public void printPerfMetrics() {
     System.out.println("Number of hash collisions: " + map.getNumHashCollisions());
     System.out.println("Time spent resizing (ns): " + map.getTimeSpentResizingNs());
     System.out.println("Total memory consumption (bytes): " + map.getTotalMemoryConsumption());
-    System.out.println("Number of unique objects in keys: " + keyPool.size());
-    System.out.println("Number of objects in buffers: " + bufferPool.size());
   }
 
 }