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

Reverts ObjectPool. As it stands, it has a few problems:

1. ObjectPool doesn't work with spilling and memory accounting.
2. I don't think in the long run the idea of an object pool is what we want to support, since it essentially goes back to unmanaged memory, and creates pressure on GC, and is hard to account for the total in memory size.
3. The ObjectPool patch removed the specialized getters for strings and binary, and as a result, actually introduced branches when reading non primitive data types.

If we do want to support arbitrary user defined types in the future, I think we can just add an object array in UnsafeRow, rather than relying on indirect memory addressing through a pool. We also need to pick execution strategies that are optimized for those, rather than keeping a lot of unserialized JVM objects in memory during aggregation.

This is probably the hardest thing I had to revert in Spark, due to recent patches that also change the same part of the code. Would be great to get a careful look.

Author: Reynold Xin <[email protected]>

Closes apache#7591 from rxin/revert-object-pool and squashes the following commits:

01db0bc [Reynold Xin] Scala style.
eda89fc [Reynold Xin] Fixed describe.
2967118 [Reynold Xin] Fixed accessor for JoinedRow.
e3294eb [Reynold Xin] Merge branch 'master' into revert-object-pool
657855f [Reynold Xin] Temp commit.
c20f2c8 [Reynold Xin] Style fix.
fe37079 [Reynold Xin] Revert "[SPARK-8579] [SQL] support arbitrary object in UnsafeRow"

Author: rxin

project/SparkBuild.scala

@@ -543,7 +543,7 @@ object TestSettings {
     javaOptions in Test += "-Dspark.ui.enabled=false",
     javaOptions in Test += "-Dspark.ui.showConsoleProgress=false",
     javaOptions in Test += "-Dspark.driver.allowMultipleContexts=true",
-    javaOptions in Test += "-Dspark.unsafe.exceptionOnMemoryLeak=true",
+    //javaOptions in Test += "-Dspark.unsafe.exceptionOnMemoryLeak=true",
     javaOptions in Test += "-Dsun.io.serialization.extendedDebugInfo=true",
     javaOptions in Test += "-Dderby.system.durability=test",
     javaOptions in Test ++= System.getProperties.filter(_._1 startsWith "spark")

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());
   }
 
 }