SPARK-24717 Split out min retain version of state for memory in HDFSBackedStateStoreProvider

* introduce BoundedSortedMap which implements bounded size of sorted map
  * only first N elements will be retained
* replace loadedMaps to BoundedSortedMap to retain only N versions of states
  * no need to cleanup in maintenance phase
* introduce new configuration: spark.sql.streaming.minBatchesToRetainInMemory

Author: HeartSaVioR

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -825,6 +825,13 @@ object SQLConf {
     .intConf
     .createWithDefault(100)
 
+  val MAX_BATCHES_TO_RETAIN_IN_MEMORY = buildConf("spark.sql.streaming.maxBatchesToRetainInMemory")
+    .internal()
+    .doc("The maximum number of batches which will be retained in memory to avoid " +
+      "loading from files.")
+    .intConf
+    .createWithDefault(2)
+
   val UNSUPPORTED_OPERATION_CHECK_ENABLED =
     buildConf("spark.sql.streaming.unsupportedOperationCheck")
       .internal()
@@ -1463,6 +1470,8 @@ class SQLConf extends Serializable with Logging {
 
   def minBatchesToRetain: Int = getConf(MIN_BATCHES_TO_RETAIN)
 
+  def maxBatchesToRetainInMemory: Int = getConf(MAX_BATCHES_TO_RETAIN_IN_MEMORY)
+
   def parquetFilterPushDown: Boolean = getConf(PARQUET_FILTER_PUSHDOWN_ENABLED)
 
   def parquetFilterPushDownDate: Boolean = getConf(PARQUET_FILTER_PUSHDOWN_DATE_ENABLED)

sql/core/src/main/java/org/apache/spark/sql/streaming/state/BoundedSortedMap.java

@@ -0,0 +1,79 @@
+/*
+ * 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.streaming.state;
+
+import java.util.Comparator;
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * This class implements bounded {@link java.util.SortedMap} based on {@link java.util.TreeMap}.
+ * <p/>
+ * As TreeMap does, this implementation sorts elements in natural order, and cuts off
+ * smaller elements to retain at most bigger N elements.
+ * <p/>
+ * You can provide reversed order of comparator to retain smaller elements instead.
+ * <p/>
+ * This class is not thread-safe, so synchronization would be needed to use this concurrently.
+ *
+ * @param <K> key type
+ * @param <V> value type
+ */
+public class BoundedSortedMap<K, V> extends TreeMap<K, V> {
+
+  private final int limit;
+
+  /**
+   * Constructor
+   *
+   * @param comparator comparator instance to compare between keys
+   * @param limit      bounded size
+   */
+  public BoundedSortedMap(Comparator<K> comparator, int limit) {
+    super(comparator);
+    this.limit = limit;
+  }
+
+  @Override
+  public void putAll(Map<? extends K, ? extends V> map) {
+    for (Map.Entry<? extends K, ? extends V> entry : map.entrySet()) {
+      put(entry.getKey(), entry.getValue());
+    }
+  }
+
+  @Override
+  public V put(K key, V value) {
+    // This method doesn't guarantee thread-safety, so it must be guarded with synchronization
+
+    while (size() > limit) {
+      remove(lastKey());
+    }
+
+    if (size() == limit) {
+      K lk = lastKey();
+      if (comparator().compare(lk, key) > 0) {
+        remove(lk);
+        return super.put(key, value);
+      }
+
+      // unchanged
+      return null;
+    }
+
+    return super.put(key, value);
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala

@@ -35,6 +35,7 @@ import org.apache.spark.io.LZ4CompressionCodec
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.execution.streaming.CheckpointFileManager
 import org.apache.spark.sql.execution.streaming.CheckpointFileManager.CancellableFSDataOutputStream
+import org.apache.spark.sql.streaming.state.BoundedSortedMap
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.util.{SizeEstimator, Utils}
 
@@ -203,6 +204,9 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
     this.valueSchema = valueSchema
     this.storeConf = storeConf
     this.hadoopConf = hadoopConf
+    this.numberOfVersionsRetainInMemory = storeConf.maxVersionsToRetainInMemory
+    this.loadedMaps = new BoundedSortedMap[Long, MapType](Ordering[Long].reverse,
+      numberOfVersionsRetainInMemory)
     fm.mkdirs(baseDir)
   }
 
@@ -220,7 +224,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
   }
 
   override def close(): Unit = {
-    loadedMaps.values.foreach(_.clear())
+    loadedMaps.values.asScala.foreach(_.clear())
   }
 
   override def supportedCustomMetrics: Seq[StateStoreCustomMetric] = {
@@ -240,7 +244,11 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
   @volatile private var storeConf: StateStoreConf = _
   @volatile private var hadoopConf: Configuration = _
 
-  private lazy val loadedMaps = new mutable.HashMap[Long, MapType]
+  // taking default value first: this will be updated by init method with configuration
+  @volatile private var numberOfVersionsRetainInMemory: Int = 2
+  @volatile private var loadedMaps = new BoundedSortedMap[Long, MapType](Ordering[Long].reverse,
+    numberOfVersionsRetainInMemory)
+
   private lazy val baseDir = stateStoreId.storeCheckpointLocation()
   private lazy val fm = CheckpointFileManager.create(baseDir, hadoopConf)
   private lazy val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf)
@@ -260,7 +268,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
    */
   private[state] def latestIterator(): Iterator[UnsafeRowPair] = synchronized {
     val versionsInFiles = fetchFiles().map(_.version).toSet
-    val versionsLoaded = loadedMaps.keySet
+    val versionsLoaded = loadedMaps.keySet.asScala
     val allKnownVersions = versionsInFiles ++ versionsLoaded
     val unsafeRowTuple = new UnsafeRowPair()
     if (allKnownVersions.nonEmpty) {
@@ -274,7 +282,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
   private def loadMap(version: Long): MapType = {
 
     // Shortcut if the map for this version is already there to avoid a redundant put.
-    val loadedCurrentVersionMap = synchronized { loadedMaps.get(version) }
+    val loadedCurrentVersionMap = synchronized { Option(loadedMaps.get(version)) }
     if (loadedCurrentVersionMap.isDefined) {
       return loadedCurrentVersionMap.get
     }
@@ -302,7 +310,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
           lastAvailableMap = Some(new MapType)
         } else {
           lastAvailableMap =
-            synchronized { loadedMaps.get(lastAvailableVersion) }
+            synchronized { Option(loadedMaps.get(lastAvailableVersion)) }
               .orElse(readSnapshotFile(lastAvailableVersion))
         }
       }
@@ -506,7 +514,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
         val lastVersion = files.last.version
         val deltaFilesForLastVersion =
           filesForVersion(files, lastVersion).filter(_.isSnapshot == false)
-        synchronized { loadedMaps.get(lastVersion) } match {
+        synchronized { Option(loadedMaps.get(lastVersion)) } match {
           case Some(map) =>
             if (deltaFilesForLastVersion.size > storeConf.minDeltasForSnapshot) {
               val (_, e2) = Utils.timeTakenMs(writeSnapshotFile(lastVersion, map))
@@ -536,10 +544,6 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
         val earliestVersionToRetain = files.last.version - storeConf.minVersionsToRetain
         if (earliestVersionToRetain > 0) {
           val earliestFileToRetain = filesForVersion(files, earliestVersionToRetain).head
-          synchronized {
-            val mapsToRemove = loadedMaps.keys.filter(_ < earliestVersionToRetain).toSeq
-            mapsToRemove.foreach(loadedMaps.remove)
-          }
           val filesToDelete = files.filter(_.version < earliestFileToRetain.version)
           val (_, e2) = Utils.timeTakenMs {
             filesToDelete.foreach { f =>

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala

@@ -34,6 +34,9 @@ class StateStoreConf(@transient private val sqlConf: SQLConf)
   /** Minimum versions a State Store implementation should retain to allow rollbacks */
   val minVersionsToRetain: Int = sqlConf.minBatchesToRetain
 
+  /** Maximum count of versions a State Store implementation should retain in memory */
+  val maxVersionsToRetainInMemory: Int = sqlConf.maxBatchesToRetainInMemory
+
   /**
    * Optional fully qualified name of the subclass of [[StateStoreProvider]]
    * managing state data. That is, the implementation of the State Store to use.

sql/core/src/test/java/test/org/apache/spark/sql/streaming/state/BoundedSortedMapSuite.java

@@ -0,0 +1,147 @@
+/*
+ * 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 test.org.apache.spark.sql.streaming.state;
+
+import org.apache.spark.sql.streaming.state.BoundedSortedMap;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Comparator;
+
+public class BoundedSortedMapSuite {
+
+  @Test
+  public void testAddElementBelowBoundedCount() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.naturalOrder(), 2);
+
+    map.put(2, 2);
+    map.put(1, 1);
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(1), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(2), map.lastKey());
+  }
+
+  @Test
+  public void testAddSmallestElementToFullOfMap() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.naturalOrder(), 2);
+
+    map.put(2, 2);
+    map.put(1, 1);
+
+    map.put(0, 0);
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(0), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(1), map.lastKey());
+  }
+
+  @Test
+  public void testAddMiddleOfElementToFullOfMap() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.naturalOrder(), 2);
+
+    map.put(3, 3);
+    map.put(1, 1);
+
+    // 3 is being cut off
+    map.put(2, 2);
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(1), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(2), map.lastKey());
+  }
+
+  @Test
+  public void testAddBiggestOfElementToFullOfMap() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.naturalOrder(), 2);
+
+    map.put(2, 2);
+    map.put(1, 1);
+
+    // 3 is not added
+    Assert.assertNull(map.put(3, 3));
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(1), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(2), map.lastKey());
+  }
+
+  @Test
+  public void testReversedOrderAddElementBelowBoundedCount() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.reverseOrder(), 2);
+
+    map.put(2, 2);
+    map.put(1, 1);
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(2), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(1), map.lastKey());
+  }
+
+  @Test
+  public void testReversedOrderAddBiggestElementToFullOfMap() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.reverseOrder(), 2);
+
+    map.put(2, 2);
+    map.put(1, 1);
+
+    // 1 is being cut off
+    map.put(3, 3);
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(3), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(2), map.lastKey());
+  }
+
+  @Test
+  public void testReversedOrderAddMiddleOfElementToFullOfMap() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.reverseOrder(), 2);
+
+    map.put(3, 3);
+    map.put(1, 1);
+
+    // 1 is being cut off
+    map.put(2, 2);
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(3), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(2), map.lastKey());
+  }
+
+  @Test
+  public void testReversedOrderAddSmallestOfElementToFullOfMap() {
+    BoundedSortedMap<Integer, Integer> map = new BoundedSortedMap<Integer, Integer>(
+        Comparator.reverseOrder(), 2);
+
+    map.put(3, 3);
+    map.put(2, 2);
+
+    // 1 is not added
+    Assert.assertNull(map.put(1, 1));
+
+    Assert.assertEquals(2, map.size());
+    Assert.assertEquals(Integer.valueOf(3), map.firstKey());
+    Assert.assertEquals(Integer.valueOf(2), map.lastKey());
+  }
+}