Implement the new state format for stream-stream join

Author: HeartSaVioR

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/join/StreamingSymmetricHashJoinExec.scala

@@ -433,7 +433,7 @@ case class StreamingSymmetricHashJoinExec(
         }
 
         val initIterFn = { () =>
-          val removedRowIter = joinerManager.leftSideJoiner.removeOldState()
+          val removedRowIter = joinerManager.leftSideJoiner.removeAndReturnOldState()
           removedRowIter.filterNot { kv =>
             stateFormatVersion match {
               case 1 => matchesWithRightSideState(new UnsafeRowPair(kv.key, kv.value))
@@ -459,7 +459,7 @@ case class StreamingSymmetricHashJoinExec(
         }
 
         val initIterFn = { () =>
-          val removedRowIter = joinerManager.rightSideJoiner.removeOldState()
+          val removedRowIter = joinerManager.rightSideJoiner.removeAndReturnOldState()
           removedRowIter.filterNot { kv =>
             stateFormatVersion match {
               case 1 => matchesWithLeftSideState(new UnsafeRowPair(kv.key, kv.value))
@@ -484,13 +484,13 @@ case class StreamingSymmetricHashJoinExec(
           }
 
         val leftSideInitIterFn = { () =>
-          val removedRowIter = joinerManager.leftSideJoiner.removeOldState()
+          val removedRowIter = joinerManager.leftSideJoiner.removeAndReturnOldState()
           removedRowIter.filterNot(isKeyToValuePairMatched)
             .map(pair => joinedRow.withLeft(pair.value).withRight(nullRight))
         }
 
         val rightSideInitIterFn = { () =>
-          val removedRowIter = joinerManager.rightSideJoiner.removeOldState()
+          val removedRowIter = joinerManager.rightSideJoiner.removeAndReturnOldState()
           removedRowIter.filterNot(isKeyToValuePairMatched)
             .map(pair => joinedRow.withLeft(nullLeft).withRight(pair.value))
         }
@@ -539,22 +539,19 @@ case class StreamingSymmetricHashJoinExec(
         // the outer side (e.g., left side for left outer join) while generating the outer "null"
         // outputs. Now, we have to remove unnecessary state rows from the other side (e.g., right
         // side for the left outer join) if possible. In all cases, nothing needs to be outputted,
-        // hence the removal needs to be done greedily by immediately consuming the returned
-        // iterator.
+        // hence the removal needs to be done greedily.
         //
         // For full outer joins, we have already removed unnecessary states from both sides, so
         // nothing needs to be outputted here.
-        val cleanupIter = joinType match {
-          case Inner | LeftSemi => joinerManager.removeOldState()
-          case LeftOuter => joinerManager.rightSideJoiner.removeOldState()
-          case RightOuter => joinerManager.leftSideJoiner.removeOldState()
-          case FullOuter => Iterator.empty
-          case _ => throwBadJoinTypeException()
-        }
-        while (cleanupIter.hasNext) {
-          cleanupIter.next()
-          numRemovedStateRows += 1
-        }
+        numRemovedStateRows += (
+          joinType match {
+            case Inner | LeftSemi => joinerManager.removeOldState()
+            case LeftOuter => joinerManager.rightSideJoiner.removeOldState()
+            case RightOuter => joinerManager.leftSideJoiner.removeOldState()
+            case FullOuter => 0L
+            case _ => throwBadJoinTypeException()
+          }
+        )
       }
 
       // Commit all state changes and update state store metrics
@@ -643,7 +640,7 @@ case class StreamingSymmetricHashJoinExec(
     private[this] val keyGenerator = UnsafeProjection.create(joinKeys, inputAttributes)
 
     private[this] val stateKeyWatermarkPredicateFunc = stateWatermarkPredicate match {
-      case Some(JoinStateKeyWatermarkPredicate(expr)) =>
+      case Some(JoinStateKeyWatermarkPredicate(expr, _)) =>
         // inputSchema can be empty as expr should only have BoundReferences and does not require
         // the schema to generated predicate. See [[StreamingSymmetricHashJoinHelper]].
         Predicate.create(expr, Seq.empty).eval _
@@ -652,7 +649,7 @@ case class StreamingSymmetricHashJoinExec(
     }
 
     private[this] val stateValueWatermarkPredicateFunc = stateWatermarkPredicate match {
-      case Some(JoinStateValueWatermarkPredicate(expr)) =>
+      case Some(JoinStateValueWatermarkPredicate(expr, _)) =>
         Predicate.create(expr, inputAttributes).eval _
       case _ =>
         Predicate.create(Literal(false), Seq.empty).eval _  // false = do not remove if no predicate
@@ -792,6 +789,29 @@ case class StreamingSymmetricHashJoinExec(
       joinStateManager.get(key)
     }
 
+    // FIXME: doc!
+    def removeOldState(): Long = {
+      stateWatermarkPredicate match {
+        case Some(JoinStateKeyWatermarkPredicate(_, stateWatermark)) =>
+          joinStateManager match {
+            case s: SupportsEvictByCondition =>
+              s.evictByKeyCondition(stateKeyWatermarkPredicateFunc)
+
+            case s: SupportsEvictByTimestamp =>
+              s.evictByTimestamp(stateWatermark)
+          }
+        case Some(JoinStateValueWatermarkPredicate(_, stateWatermark)) =>
+          joinStateManager match {
+            case s: SupportsEvictByCondition =>
+              s.evictByValueCondition(stateValueWatermarkPredicateFunc)
+
+            case s: SupportsEvictByTimestamp =>
+              s.evictByTimestamp(stateWatermark)
+          }
+        case _ => 0L
+      }
+    }
+
     /**
      * Builds an iterator over old state key-value pairs, removing them lazily as they're produced.
      *
@@ -802,12 +822,24 @@ case class StreamingSymmetricHashJoinExec(
      * We do this to avoid requiring either two passes or full materialization when
      * processing the rows for outer join.
      */
-    def removeOldState(): Iterator[KeyToValuePair] = {
+    def removeAndReturnOldState(): Iterator[KeyToValuePair] = {
       stateWatermarkPredicate match {
-        case Some(JoinStateKeyWatermarkPredicate(expr)) =>
-          joinStateManager.removeByKeyCondition(stateKeyWatermarkPredicateFunc)
-        case Some(JoinStateValueWatermarkPredicate(expr)) =>
-          joinStateManager.removeByValueCondition(stateValueWatermarkPredicateFunc)
+        case Some(JoinStateKeyWatermarkPredicate(_, stateWatermark)) =>
+          joinStateManager match {
+            case s: SupportsEvictByCondition =>
+              s.evictAndReturnByKeyCondition(stateKeyWatermarkPredicateFunc)
+
+            case s: SupportsEvictByTimestamp =>
+              s.evictAndReturnByTimestamp(stateWatermark)
+          }
+        case Some(JoinStateValueWatermarkPredicate(_, stateWatermark)) =>
+          joinStateManager match {
+            case s: SupportsEvictByCondition =>
+              s.evictAndReturnByValueCondition(stateValueWatermarkPredicateFunc)
+
+            case s: SupportsEvictByTimestamp =>
+              s.evictAndReturnByTimestamp(stateWatermark)
+          }
         case _ => Iterator.empty
       }
     }
@@ -836,8 +868,12 @@ case class StreamingSymmetricHashJoinExec(
   private case class OneSideHashJoinerManager(
       leftSideJoiner: OneSideHashJoiner, rightSideJoiner: OneSideHashJoiner) {
 
-    def removeOldState(): Iterator[KeyToValuePair] = {
-      leftSideJoiner.removeOldState() ++ rightSideJoiner.removeOldState()
+    def removeOldState(): Long = {
+      leftSideJoiner.removeOldState() + rightSideJoiner.removeOldState()
+    }
+
+    def removeAndReturnOldState(): Iterator[KeyToValuePair] = {
+      leftSideJoiner.removeAndReturnOldState() ++ rightSideJoiner.removeAndReturnOldState()
     }
 
     def metrics: StateStoreMetrics = {

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/join/StreamingSymmetricHashJoinHelper.scala

@@ -46,12 +46,12 @@ object StreamingSymmetricHashJoinHelper extends Logging {
     override def toString: String = s"$desc: $expr"
   }
   /** Predicate for watermark on state keys */
-  case class JoinStateKeyWatermarkPredicate(expr: Expression)
+  case class JoinStateKeyWatermarkPredicate(expr: Expression, stateWatermark: Long)
     extends JoinStateWatermarkPredicate {
     def desc: String = "key predicate"
   }
   /** Predicate for watermark on state values */
-  case class JoinStateValueWatermarkPredicate(expr: Expression)
+  case class JoinStateValueWatermarkPredicate(expr: Expression, stateWatermark: Long)
     extends JoinStateWatermarkPredicate {
     def desc: String = "value predicate"
   }
@@ -212,8 +212,11 @@ object StreamingSymmetricHashJoinHelper extends Logging {
           oneSideJoinKeys(joinKeyOrdinalForWatermark.get).dataType,
           oneSideJoinKeys(joinKeyOrdinalForWatermark.get).nullable)
         val expr = watermarkExpression(Some(keyExprWithWatermark), eventTimeWatermarkForEviction)
-        expr.map(JoinStateKeyWatermarkPredicate.apply _)
-
+        expr.map { e =>
+          // watermarkExpression only provides the expression when eventTimeWatermarkForEviction
+          // is defined
+          JoinStateKeyWatermarkPredicate(e, eventTimeWatermarkForEviction.get)
+        }
       } else if (isWatermarkDefinedOnInput) { // case 2 in the StreamingSymmetricHashJoinExec docs
         val stateValueWatermark = StreamingJoinHelper.getStateValueWatermark(
           attributesToFindStateWatermarkFor = AttributeSet(oneSideInputAttributes),
@@ -222,8 +225,11 @@ object StreamingSymmetricHashJoinHelper extends Logging {
           eventTimeWatermarkForEviction)
         val inputAttributeWithWatermark = oneSideInputAttributes.find(_.metadata.contains(delayKey))
         val expr = watermarkExpression(inputAttributeWithWatermark, stateValueWatermark)
-        expr.map(JoinStateValueWatermarkPredicate.apply _)
-
+        expr.map { e =>
+          // watermarkExpression only provides the expression when eventTimeWatermarkForEviction
+          // is defined
+          JoinStateValueWatermarkPredicate(e, stateValueWatermark.get)
+        }
       } else {
         None
       }

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/join/StreamingSymmetricHashJoinValueRowConverter.scala

@@ -0,0 +1,104 @@
+/*
+ * 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.streaming.operators.stateful.join
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Literal, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.execution.streaming.operators.stateful.join.SymmetricHashJoinStateManager.ValueAndMatchPair
+import org.apache.spark.sql.types.BooleanType
+
+/**
+ * Converter between the value row stored in state store and the (actual value, match) pair.
+ */
+trait StreamingSymmetricHashJoinValueRowConverter {
+  /** Defines the schema of the value row (the value side of K-V in state store). */
+  def valueAttributes: Seq[Attribute]
+
+  /**
+   * Convert the value row to (actual value, match) pair.
+   *
+   * NOTE: implementations should ensure the result row is NOT reused during execution, so
+   * that caller can safely read the value in any time.
+   */
+  def convertValue(value: UnsafeRow): ValueAndMatchPair
+
+  /**
+   * Build the value row from (actual value, match) pair. This is expected to be called just
+   * before storing to the state store.
+   *
+   * NOTE: depending on the implementation, the result row "may" be reused during execution
+   * (to avoid initialization of object), so the caller should ensure that the logic doesn't
+   * affect by such behavior. Call copy() against the result row if needed.
+   */
+  def convertToValueRow(value: UnsafeRow, matched: Boolean): UnsafeRow
+}
+
+class StreamingSymmetricHashJoinValueRowConverterFormatV1(
+    inputValueAttributes: Seq[Attribute]) extends StreamingSymmetricHashJoinValueRowConverter {
+  override val valueAttributes: Seq[Attribute] = inputValueAttributes
+
+  override def convertValue(value: UnsafeRow): ValueAndMatchPair = {
+    if (value != null) ValueAndMatchPair(value, false) else null
+  }
+
+  override def convertToValueRow(value: UnsafeRow, matched: Boolean): UnsafeRow = value
+}
+
+class StreamingSymmetricHashJoinValueRowConverterFormatV2(
+    inputValueAttributes: Seq[Attribute]) extends StreamingSymmetricHashJoinValueRowConverter {
+  private val valueWithMatchedExprs = inputValueAttributes :+ Literal(true)
+  private val indexOrdinalInValueWithMatchedRow = inputValueAttributes.size
+
+  private val valueWithMatchedRowGenerator = UnsafeProjection.create(valueWithMatchedExprs,
+    inputValueAttributes)
+
+  override val valueAttributes: Seq[Attribute] = inputValueAttributes :+
+    AttributeReference("matched", BooleanType)()
+
+  // Projection to generate key row from (value + matched) row
+  private val valueRowGenerator = UnsafeProjection.create(
+    inputValueAttributes, valueAttributes)
+
+  override def convertValue(value: UnsafeRow): ValueAndMatchPair = {
+    if (value != null) {
+      ValueAndMatchPair(valueRowGenerator(value).copy(),
+        value.getBoolean(indexOrdinalInValueWithMatchedRow))
+    } else {
+      null
+    }
+  }
+
+  override def convertToValueRow(value: UnsafeRow, matched: Boolean): UnsafeRow = {
+    val row = valueWithMatchedRowGenerator(value)
+    row.setBoolean(indexOrdinalInValueWithMatchedRow, matched)
+    row
+  }
+}
+
+object StreamingSymmetricHashJoinValueRowConverter {
+  def create(
+      inputValueAttributes: Seq[Attribute],
+      stateFormatVersion: Int): StreamingSymmetricHashJoinValueRowConverter = {
+    stateFormatVersion match {
+      case 1 => new StreamingSymmetricHashJoinValueRowConverterFormatV1(inputValueAttributes)
+      case 2 | 3 | 4 =>
+        new StreamingSymmetricHashJoinValueRowConverterFormatV2(inputValueAttributes)
+      case _ => throw new IllegalArgumentException ("Incorrect state format version! " +
+        s"version $stateFormatVersion")
+    }
+  }
+}