WIP fixed the issue with benchmark run

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.execution.streaming
 
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
 import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.execution.streaming.state.SessionWindowLinkedListState
@@ -115,30 +115,13 @@ class MergingSortWithSessionWindowLinkedListStateIterator(
     val stateSessionsEnclosingCurrentRow = findSessionPointerEnclosingEvent(currentRow,
       startPointer = lastCheckpointOnStateRows)
 
-    // FIXME: debugging...
-    def loadSession(
-        keys: UnsafeRow,
-        sessionStart: Long,
-        stateSessionsEnclosingCurrentRow: Option[(Option[Long], Option[Long])])
-      : UnsafeRow = {
-      val sessionState = state.get(currentRow.keys, sessionStart)
-      require(sessionState != null,
-        s"Session must be presented in state: it may represent dangling pointer - " +
-          s"$sessionStart / key: $keys / currentRow: $currentRow" +
-          s"sessionsEnclosingCurrentRow: $stateSessionsEnclosingCurrentRow" +
-          s"Pointers: ${state.iteratePointers(currentRow.keys).toList}" +
-          s"Values: ${state.get(currentRow.keys).toList}")
-      sessionState
-    }
-
     var prevSessionToEmit: Option[SessionRowInformation] = None
     stateSessionsEnclosingCurrentRow match {
       case None =>
       case Some(x) =>
         x._1 match {
           case Some(prev) =>
-            val prevSession = SessionRowInformation.of(
-              loadSession(currentRow.keys, prev, stateSessionsEnclosingCurrentRow))
+            val prevSession = SessionRowInformation.of(state.get(currentRow.keys, prev))
 
             val sessionLaterThanCheckpoint = lastCheckpointOnStateRows match {
               case Some(lastCheckpoint) => lastCheckpoint < prevSession.sessionStart
@@ -161,8 +144,7 @@ class MergingSortWithSessionWindowLinkedListStateIterator(
 
         x._2 match {
           case Some(next) =>
-            val nextSession = SessionRowInformation.of(
-              loadSession(currentRow.keys, next, stateSessionsEnclosingCurrentRow))
+            val nextSession = SessionRowInformation.of(state.get(currentRow.keys, next))
 
             val sessionLaterThanCheckpoint = lastCheckpointOnStateRows match {
               case Some(lastCheckpoint) => lastCheckpoint < nextSession.sessionStart

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

@@ -157,7 +157,6 @@ class SessionWindowLinkedListState(
   def update(key: UnsafeRow, sessionStart: Long, newValue: UnsafeRow): Unit = {
     val targetPointer = keyAndSessionStartToPointerStore.get(key, sessionStart)
     assertValidPointer(targetPointer)
-
     keyAndSessionStartToValueStore.put(key, sessionStart, newValue)
   }
 
@@ -243,6 +242,9 @@ class SessionWindowLinkedListState(
     val prevOption = targetPointer._1
     val nextOption = targetPointer._2
 
+    keyAndSessionStartToPointerStore.remove(key, sessionStart)
+    keyAndSessionStartToValueStore.remove(key, sessionStart)
+
     targetPointer match {
       case (Some(prev), Some(next)) =>
         keyAndSessionStartToPointerStore.updateNext(key, prev, nextOption)
@@ -260,11 +262,10 @@ class SessionWindowLinkedListState(
           throw new IllegalStateException("The element has pointer information for head, " +
             "but the list has different head.")
         }
-        keyAndSessionStartToPointerStore.remove(key, sessionStart)
+
         keyToHeadSessionStartStore.remove(key)
     }
 
-    keyAndSessionStartToValueStore.remove(key, sessionStart)
   }
 
   def removeByValueCondition(removalCondition: UnsafeRow => Boolean,
@@ -454,15 +455,15 @@ class SessionWindowLinkedListState(
     )
   }
 
-  private[state] def getIteratorOfHeadPointers: Iterator[KeyAndHeadSessionStart] = {
+  private[sql] def getIteratorOfHeadPointers: Iterator[KeyAndHeadSessionStart] = {
     keyToHeadSessionStartStore.iterator
   }
 
-  private[state] def getIteratorOfRawPointers: Iterator[KeyWithSessionStartAndPointers] = {
+  private[sql] def getIteratorOfRawPointers: Iterator[KeyWithSessionStartAndPointers] = {
     keyAndSessionStartToPointerStore.iterator
   }
 
-  private[state] def getIteratorOfRawValues: Iterator[KeyWithSessionStartAndValue] = {
+  private[sql] def getIteratorOfRawValues: Iterator[KeyWithSessionStartAndValue] = {
     keyAndSessionStartToValueStore.iterator
   }
 
@@ -646,7 +647,7 @@ class SessionWindowLinkedListState(
 
     def updatePrev(key: UnsafeRow, sessionStart: Long, prevSessionStart: Option[Long]): Unit = {
       val actualKeyRow = keyWithSessionStartRow(key, sessionStart)
-      val row = stateStore.get(actualKeyRow)
+      val row = stateStore.get(actualKeyRow).copy()
       setPrevSessionStart(row, prevSessionStart)
       stateStore.put(actualKeyRow, row)
     }

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/MergingSortWithSessionWindowLinkedListStateIteratorSuite.scala

@@ -395,7 +395,6 @@ class MergingSortWithSessionWindowLinkedListStateIteratorSuite extends SharedSQL
   }
 
   private def assertRowsEquals(expectedRow: InternalRow, retRow: InternalRow): Unit = {
-
     val tupleFromExpectedRow = getTupleFromRow(expectedRow)
     val tupleFromInternalRow = getTupleFromRow(retRow)
     try {

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/SessionWindowLinkedListStateSuite.scala

@@ -217,200 +217,6 @@ class SessionWindowLinkedListStateSuite extends StreamTest {
     removeByWatermarkTest(stopOnConditionMismatch = false)
   }
 
-  test("run chaos monkey") {
-
-    // FIXME: too many args
-    def printFailureInformation(
-        ex: TestFailedException,
-        state: SessionWindowLinkedListState,
-        operation: Int,
-        addBefore: Boolean,
-        opIdx: Int,
-        targetIdx: Int,
-        key: UnsafeRow,
-        headPointersBeforeOp: List[(Int, Long)],
-        rawPointersBeforeOp: List[(Int, Long, Option[Long], Option[Long])],
-        pointersBeforeOp: List[Long],
-        valuesBeforeOp: List[(Int, Int)],
-        refListBeforeOp: java.util.LinkedList[String],
-        refList: java.util.LinkedList[String]): Unit = {
-      logError("Assertion failure!", ex)
-
-      logError("===== Operation information =====")
-      val opString = operation match {
-        case 0 => "Append"
-        case 1 => "Remove"
-        case 2 => "RemoveValuesByCondition"
-        case _ => throw new IllegalStateException(s"Unknown operation $operation")
-      }
-      val addPositionString = if (addBefore) "AddBefore" else "AddAfter"
-      logError(s"Operation Index: $opIdx")
-      logError(s"Operation: $opString")
-      logError(s"Position to add: $addPositionString")
-      logError(s"Target index: $targetIdx")
-
-      logError("===== Before applying operation =====")
-      logError(s"Head pointers in state: $headPointersBeforeOp")
-      logError(s"Raw pointers in state: $rawPointersBeforeOp")
-      logError(s"Pointers in state via iteratePointers: $pointersBeforeOp")
-      logError(s"Values in state: $valuesBeforeOp")
-      logError(s"Values in reference list: $refListBeforeOp")
-
-      logError("===== After applying operation =====")
-
-      val headPointers = state.getIteratorOfHeadPointers.map { pair =>
-        (toKeyInt(pair.key), pair.sessionStart)
-      }.toList
-      val pointers = state.iteratePointers(key).map(_._1).toList
-      val rawPointers = state.getIteratorOfRawPointers.map { pointer =>
-        (toKeyInt(pointer.key), pointer.sessionStart, pointer.prevSessionStart,
-          pointer.nextSessionStart)
-      }.toList
-      val values = state.getIteratorOfRawValues.map { value =>
-        (toKeyInt(value.key), toValueInt(value.value))
-      }.toList
-
-      logError(s"Head pointers in state: $headPointers")
-      logError(s"Raw pointers in state: $rawPointers")
-      logError(s"Pointers in state via iteratePointers: $pointers")
-      logError(s"Values in state: $values")
-      logError(s"Values in reference list: $refList")
-    }
-
-    withSessionWindowLinkedListState(inputValueAttribs, keyExprs) { state =>
-      implicit val st = state
-
-      assert(numRows === 0)
-
-      val rand = new Random()
-
-      val keys = (0 to 2).map(id => toKeyRow(id).copy())
-      // using String type to avoid confusion in remove(int) vs remove(Object)
-      // which LinkedList[Integer] will be remove(int) vs remove(Integer)
-      val refLists = keys.map(_ => new java.util.LinkedList[String]())
-
-      val maxOperations = 100000
-      (0 until maxOperations).foreach { opIdx =>
-
-        val selectedKeyIdx = rand.nextInt(keys.length)
-        val selectedKey = keys(selectedKeyIdx)
-        val selectedRefList = refLists(selectedKeyIdx)
-
-        // 0: append, 1: remove, 2: removeValueByCondition
-        val operation = rand.nextInt(3)
-        val addBefore = rand.nextBoolean()
-        val targetIdx = if (selectedRefList.isEmpty) -1 else rand.nextInt(selectedRefList.size())
-
-        val headPointersBeforeOp = state.getIteratorOfHeadPointers.map { pair =>
-          (toKeyInt(pair.key), pair.sessionStart)
-        }.toList
-        val pointersBeforeOp = state.iteratePointers(selectedKey).map(_._1).toList
-        val rawPointersBeforeOp = state.getIteratorOfRawPointers.map { pointer =>
-          (toKeyInt(pointer.key), pointer.sessionStart, pointer.prevSessionStart,
-            pointer.nextSessionStart)
-        }.toList
-        val valuesBeforeOp = state.getIteratorOfRawValues.map { value =>
-          (toKeyInt(value.key), toValueInt(value.value))
-        }.toList
-
-        val refListBeforeOp = new java.util.LinkedList[String](refLists(selectedKeyIdx))
-
-        operation match {
-          case 0 =>
-            if (selectedRefList.isEmpty) {
-              assert(state.isEmpty(selectedKey))
-              state.setHead(selectedKey, opIdx, toInputValue(opIdx))
-              selectedRefList.addFirst(String.valueOf(opIdx))
-            } else {
-              val addBefore = rand.nextBoolean()
-              if (addBefore) {
-                val idxToAddBefore = selectedRefList.get(targetIdx)
-                selectedRefList.add(targetIdx, String.valueOf(opIdx))
-                state.addBefore(selectedKey, opIdx, toInputValue(opIdx), idxToAddBefore.toInt)
-              } else {
-                val idxToAddAfter = selectedRefList.get(targetIdx)
-                selectedRefList.add(targetIdx + 1, String.valueOf(opIdx))
-                state.addAfter(selectedKey, opIdx, toInputValue(opIdx), idxToAddAfter.toInt)
-              }
-            }
-
-          case 1 =>
-            if (selectedRefList.isEmpty) {
-              assert(state.isEmpty(selectedKey))
-              // skip removing
-            } else {
-              val pointerToRemove = selectedRefList.get(targetIdx)
-              selectedRefList.remove(targetIdx)
-              state.remove(selectedKey, pointerToRemove.toInt)
-            }
-
-          case 2 =>
-            if (selectedRefList.isEmpty) {
-              assert(state.isEmpty(selectedKey))
-              // skip removing
-            } else {
-              val pointerToRemove = selectedRefList.get(targetIdx)
-              val removedIter = state.removeByValueCondition { r =>
-                toValueInt(r) <= pointerToRemove.toInt
-              }
-
-              val valuesFromRef = new scala.collection.mutable.MutableList[Int]()
-              refLists.foreach { refList =>
-                val refIter = refList.iterator()
-                while (refIter.hasNext) {
-                  val ref = refIter.next()
-                  if (ref.toInt <= pointerToRemove.toInt) {
-                    valuesFromRef += ref.toInt
-                    refIter.remove()
-                  }
-                }
-              }
-
-              try {
-                assert(removedIter.map(pair => toValueInt(pair.value)).toSet ==
-                  valuesFromRef.toSet)
-              } catch {
-                case ex: TestFailedException =>
-                  printFailureInformation(ex, state, operation, addBefore, opIdx, targetIdx,
-                    selectedKey, headPointersBeforeOp, rawPointersBeforeOp, pointersBeforeOp,
-                    valuesBeforeOp, refListBeforeOp, selectedRefList)
-
-                  throw ex
-              }
-            }
-        }
-
-        keys.indices.foreach { index =>
-          val key = keys(index)
-          val refList = refLists(index)
-
-          try {
-            if (refList.isEmpty) {
-              assert(state.isEmpty(key), s"Reference list is empty but " +
-                s"state list for $key is not empty")
-            } else {
-              import scala.collection.JavaConverters._
-              val statePointers = state.iteratePointers(key).map(_._1).toList
-              assert(refList.asScala.map(_.toInt) === statePointers,
-                s"State pointers for $key is expected to be $refList but $statePointers")
-
-              val stateValues = state.get(key).map(toValueInt).toList
-              assert(refList.asScala.map(_.toInt) ===
-                stateValues, s"State list for $key is expected to be $refList but $stateValues")
-            }
-          } catch {
-            case ex: TestFailedException =>
-              printFailureInformation(ex, state, operation, addBefore, opIdx, targetIdx,
-                selectedKey, headPointersBeforeOp, rawPointersBeforeOp, pointersBeforeOp,
-                valuesBeforeOp, refListBeforeOp, selectedRefList)
-
-              throw ex
-          }
-        }
-      }
-    }
-  }
-
   private def removeByWatermarkTest(stopOnConditionMismatch: Boolean): Unit = {
     withSessionWindowLinkedListState(inputValueAttribs, keyExprs) { state =>
       implicit val st = state