Clean up final row copying code.

JoshRosen · JoshRosen · commit c56ec18c7106 · 2015-07-08T00:09:55.000-07:00
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql.catalyst.expressions;
 
-import javax.annotation.Nullable;
-
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.util.ObjectPool;
 import org.apache.spark.unsafe.PlatformDependent;
@@ -57,17 +55,15 @@
  */
 public final class UnsafeRow extends MutableRow {
 
-  /** Hack for if we want to pass around an UnsafeRow which also carries around its backing data */
-  @Nullable public byte[] backingArray;
   private Object baseObject;
   private long baseOffset;
 
   /** A pool to hold non-primitive objects */
   private ObjectPool pool;
 
-  Object getBaseObject() { return baseObject; }
-  long getBaseOffset() { return baseOffset; }
-  ObjectPool getPool() { return pool; }
+  public Object getBaseObject() { return baseObject; }
+  public long getBaseOffset() { return baseOffset; }
+  public ObjectPool getPool() { return pool; }
 
   /** The number of fields in this row, used for calculating the bitset width (and in assertions) */
   private int numFields;
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 
+import org.apache.spark.sql.Row;
 import scala.collection.Iterator;
 import scala.math.Ordering;
 
@@ -150,28 +151,34 @@ public boolean hasNext() {
           return sortedIterator.hasNext();
         }
 
+        /**
+         * Called prior to returning this iterator's last row. This copies the row's data into an
+         * on-heap byte array so that the pointer to the row data will not be dangling after the
+         * sorter's memory pages are freed.
+         */
+        private void detachRowFromPage(UnsafeRow row, int rowLength) {
+          final byte[] rowDataCopy = new byte[rowLength];
+          PlatformDependent.copyMemory(
+            row.getBaseObject(),
+            row.getBaseOffset(),
+            rowDataCopy,
+            PlatformDependent.BYTE_ARRAY_OFFSET,
+            rowLength
+          );
+          row.pointTo(rowDataCopy, PlatformDependent.BYTE_ARRAY_OFFSET, numFields, row.getPool());
+        }
+
         @Override
         public InternalRow next() {
           try {
             sortedIterator.loadNext();
-            if (hasNext()) {
-              row.pointTo(
-                sortedIterator.getBaseObject(), sortedIterator.getBaseOffset(), numFields, objPool);
-              return row;
-            } else {
-              final byte[] rowDataCopy = new byte[sortedIterator.getRecordLength()];
-              PlatformDependent.copyMemory(
-                sortedIterator.getBaseObject(),
-                sortedIterator.getBaseOffset(),
-                rowDataCopy,
-                PlatformDependent.BYTE_ARRAY_OFFSET,
-                sortedIterator.getRecordLength()
-              );
-              row.backingArray = rowDataCopy;
-              row.pointTo(rowDataCopy, PlatformDependent.BYTE_ARRAY_OFFSET, numFields, objPool);
+            row.pointTo(
+              sortedIterator.getBaseObject(), sortedIterator.getBaseOffset(), numFields, objPool);
+            if (!hasNext()) {
+              detachRowFromPage(row, sortedIterator.getRecordLength());
               cleanupResources();
-              return row;
             }
+            return row;
           } catch (IOException e) {
             cleanupResources();
             // Scala iterators don't declare any checked exceptions, so we need to use this hack
