[SPARK-21492][SQL][WIP] Memory leak in SortMergeJoin

sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowIterator.java

@@ -0,0 +1,26 @@
+/*
+ * 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.spark.sql.execution;
+
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import scala.collection.AbstractIterator;
+
+public abstract class UnsafeExternalRowIterator extends AbstractIterator<UnsafeRow> {
+
+    public abstract void close();
+}

sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java

@@ -20,7 +20,6 @@
 import java.io.IOException;
 import java.util.function.Supplier;
 
-import scala.collection.AbstractIterator;
 import scala.collection.Iterator;
 import scala.math.Ordering;
 
@@ -169,7 +168,7 @@ public Iterator<UnsafeRow> sort() throws IOException {
         // here in order to prevent memory leaks.
         cleanupResources();
       }
-      return new AbstractIterator<UnsafeRow>() {
+      return new UnsafeExternalRowIterator() {
 
         private final int numFields = schema.length();
         private UnsafeRow row = new UnsafeRow(numFields);
@@ -203,6 +202,11 @@ public UnsafeRow next() {
           }
           throw new RuntimeException("Exception should have been re-thrown in next()");
         }
+
+        @Override
+        public void close() {
+          cleanupResources();
+        }
       };
     } catch (IOException e) {
       cleanupResources();

sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala

@@ -195,6 +195,7 @@ case class SortMergeJoinExec(
                     currentRightMatches = null
                     currentLeftRow = null
                     rightMatchesIterator = null
+                    smjScanner.close()
                     return false
                   }
                 }
@@ -204,6 +205,7 @@ case class SortMergeJoinExec(
                   return true
                 }
               }
+              smjScanner.close()
               false
             }
 
@@ -285,6 +287,7 @@ case class SortMergeJoinExec(
                   }
                 }
               }
+              smjScanner.close()
               false
             }
 
@@ -326,6 +329,7 @@ case class SortMergeJoinExec(
                   return true
                 }
               }
+              smjScanner.close()
               false
             }
 
@@ -365,6 +369,7 @@ case class SortMergeJoinExec(
                 numOutputRows += 1
                 return true
               }
+              smjScanner.close()
               false
             }
 
@@ -669,7 +674,7 @@ private[joins] class SortMergeJoinScanner(
     streamedIter: RowIterator,
     bufferedIter: RowIterator,
     inMemoryThreshold: Int,
-    spillThreshold: Int) {
+    spillThreshold: Int) extends CloseableScanner {
   private[this] var streamedRow: InternalRow = _
   private[this] var streamedRowKey: InternalRow = _
   private[this] var bufferedRow: InternalRow = _
@@ -785,6 +790,15 @@ private[joins] class SortMergeJoinScanner(
     }
   }
 
+  /**
+   * Once the join has been completed, the iterators for both relations
+   * should be closed, so that acquired memory can be released.
+   */
+  def close(): Unit = {
+    closeIterator(streamedIter)
+    closeIterator(bufferedIter)
+  }
+
   // --- Private methods --------------------------------------------------------------------------
 
   /**
@@ -952,7 +966,11 @@ private abstract class OneSideOuterIterator(
 
   override def advanceNext(): Boolean = {
     val r = advanceBufferUntilBoundConditionSatisfied() || advanceStream()
-    if (r) numOutputRows += 1
+    if (r) {
+      numOutputRows += 1
+    } else {
+      smjScanner.close()
+    }
     r
   }
 
@@ -967,7 +985,7 @@ private class SortMergeFullOuterJoinScanner(
     rightIter: RowIterator,
     boundCondition: InternalRow => Boolean,
     leftNullRow: InternalRow,
-    rightNullRow: InternalRow)  {
+    rightNullRow: InternalRow) extends CloseableScanner {
   private[this] val joinedRow: JoinedRow = new JoinedRow()
   private[this] var leftRow: InternalRow = _
   private[this] var leftRowKey: InternalRow = _
@@ -1130,6 +1148,15 @@ private class SortMergeFullOuterJoinScanner(
       false
     }
   }
+
+  /**
+   * Once the join has been completed, the iterators for both relations
+   * should be closed, so that acquired memory can be released.
+   */
+  def close(): Unit = {
+    closeIterator(leftIter)
+    closeIterator(rightIter)
+  }
 }
 
 private class FullOuterIterator(
@@ -1140,9 +1167,26 @@ private class FullOuterIterator(
 
   override def advanceNext(): Boolean = {
     val r = smjScanner.advanceNext()
-    if (r) numRows += 1
+    if (r) {
+      numRows += 1
+    } else {
+      smjScanner.close()
+    }
     r
   }
 
   override def getRow: InternalRow = resultProj(joinedRow)
 }
+
+trait CloseableScanner {
+  def closeIterator(iter: RowIterator): Unit = {
+    if (iter.isInstanceOf[RowIteratorFromScala]) {
+      val rowIter = iter.asInstanceOf[RowIteratorFromScala]
+      val underlyingIter = rowIter.toScala
+      if (underlyingIter.isInstanceOf[UnsafeExternalRowIterator]) {
+        val toClose = underlyingIter.asInstanceOf[UnsafeExternalRowIterator]
+        toClose.close()
+      }
+    }
+  }
+}