/*
 * 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 java.util.Locale

import scala.annotation.tailrec

import org.apache.hadoop.conf.Configuration

import org.apache.spark.TaskContext
import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.{END_INDEX, INDEX, JOIN_SIDE, KEY, NUM_VALUES, PARTITION_ID, START_INDEX, STATE_STORE_ID, STATE_STORE_VERSION}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression, JoinedRow, Literal, SafeProjection, SpecificInternalRow, UnsafeProjection, UnsafeRow}
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorStateInfo, StatefulOpStateStoreCheckpointInfo, WatermarkSupport}
import org.apache.spark.sql.execution.streaming.operators.stateful.join.StreamingSymmetricHashJoinHelper._
import org.apache.spark.sql.execution.streaming.state.{DropLastNFieldsStatePartitionKeyExtractor, KeyStateEncoderSpec, NoopStatePartitionKeyExtractor, NoPrefixKeyStateEncoderSpec, RangeScanBoundaryUtils, StatePartitionKeyExtractor, StateSchemaBroadcast, StateStore, StateStoreCheckpointInfo, StateStoreColFamilySchema, StateStoreConf, StateStoreErrors, StateStoreId, StateStoreMetrics, StateStoreProvider, StateStoreProviderId, SupportsFineGrainedReplay, TimestampAsPostfixKeyStateEncoderSpec, TimestampAsPrefixKeyStateEncoderSpec, TimestampKeyStateEncoder}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{BooleanType, DataType, LongType, NullType, StructField, StructType}
import org.apache.spark.util.NextIterator

/**
 * Base trait of the state manager for stream-stream symmetric hash join operator.
 *
 * This defines the basic APIs for the state manager, except the methods for eviction which are
 * defined in separate traits - See [[SupportsEvictByCondition]] and [[SupportsEvictByTimestamp]].
 *
 * Implementation classes are expected to inherit those traits as needed, depending on the eviction
 * strategy they support.
 */
trait SymmetricHashJoinStateManager {
  import SymmetricHashJoinStateManager._

  /** Append a new value to the given key, with the flag of matched or not. */
  def append(key: UnsafeRow, value: UnsafeRow, matched: Boolean): Unit

  /**
   * Retrieve all matched values from given key. This doesn't update the matched flag for the
   * values being returned, hence should be only used if appropriate for the join type and the
   * side.
   */
  def get(key: UnsafeRow): Iterator[UnsafeRow]

  /**
   * Retrieve all joined rows for the given key. The joined rows are generated with the provided
   * generateJoinedRow function and filtered with the provided predicate.
   *
   * The matched flag will be updated to true for the values being returned, if it is semantically
   * required to do so.
   *
   * For skipUpdatingMatchedFlag = true, the method will skip updating the matched flag for the
   * values being returned. This is useful for the non-outer side of stream-stream join where
   * the matched flag is never checked during state eviction.
   *
   * It is caller's responsibility to consume the whole iterator.
   *
   * @param timestampRange Optional optimization hint as (minTimestamp, maxTimestamp), both
   *   inclusive. Derived classes may use it to reduce scan scope but are free to ignore it.
   *   The predicate must produce correct output regardless of whether this hint is leveraged.
   */
  def getJoinedRows(
      key: UnsafeRow,
      generateJoinedRow: InternalRow => JoinedRow,
      predicate: JoinedRow => Boolean,
      timestampRange: Option[(Long, Long)] = None,
      skipUpdatingMatchedFlag: Boolean = false): Iterator[JoinedRow]

  /**
   * Retrieve all joined rows for the given key and remove the matched rows from state. The joined
   * rows are generated with the provided generateJoinedRow function and filtered with the provided
   * predicate.
   *
   * The implementation of this method should consider the performance, for example, try making
   * removal of the matched rows and retrieval of the joined rows happen in a single pass if
   * possible; avoid using mark-and-sweep where mark requires write on the state store.
   *
   * It is caller's responsibility to consume the whole iterator.
   *
   * NOTE: For the rows which already have been marked as matched in the state, this method removes
   * them from the state without returning them. Under normal operation, this should not happen and
   * this should not be an issue, but it can occur if the join type was changed during query
   * restart. We do not define an expected behavior for such changes, so we just optimize it rather
   * than trying to provide some best-effort results.
   *
   * NOTE2: There is a further optimization opportunity -- if a row does not pass the predicate
   * (postJoinFilter), it may never match in future batches as long as expressions are
   * deterministic and also not changeable over batches (e.g. current batch timestamp). Eagerly
   * removing such rows would reduce state size further, but we need to weigh the value of applying
   * this optimization. This can be also impacted by the query changes during restart, as we
   * mentioned in NOTE in above, we do not define expected behavior for such changes.
   */
  def getJoinedRowsAndRemoveMatched(
      key: UnsafeRow,
      generateJoinedRow: InternalRow => JoinedRow,
      predicate: JoinedRow => Boolean): Iterator[JoinedRow]

  /**
   * Provide all key-value pairs in the state manager.
   *
   * It is caller's responsibility to consume the whole iterator.
   */
  def iterator: Iterator[KeyToValuePair]

  /** Commit all the changes to all the state stores */
  def commit(): Unit

  /** Abort any changes to the state stores if needed */
  def abortIfNeeded(): Unit

  /** Provide the metrics. */
  def metrics: StateStoreMetrics

  /**
   * Get state store checkpoint information of the two state stores for this joiner, after
   * they finished data processing.
   */
  def getLatestCheckpointInfo(): JoinerStateStoreCkptInfo
}

/**
 * This trait is specific to help the old version of state manager implementation (v1-v3) to work
 * with existing tests which look up the state store with key with index.
 */
trait SupportsIndexedKeys {
  def getInternalRowOfKeyWithIndex(currentKey: UnsafeRow): InternalRow

  protected[streaming] def updateNumValuesTestOnly(key: UnsafeRow, numValues: Long): Unit
}

/**
 * This trait is for state manager implementations that support eviction by condition.
 * This is for the state manager implementations which have to perform full scan
 * for eviction.
 */
trait SupportsEvictByCondition { self: SymmetricHashJoinStateManager =>
  import SymmetricHashJoinStateManager._

  /** Evict the state via condition on the key. Returns the number of values evicted. */
  def evictByKeyCondition(removalCondition: UnsafeRow => Boolean): Long

  /**
   * Evict the state via condition on the key, and return the evicted key-value pairs.
   *
   * It is caller's responsibility to consume the whole iterator.
   */
  def evictAndReturnByKeyCondition(
      removalCondition: UnsafeRow => Boolean): Iterator[KeyToValuePair]

  /** Evict the state via condition on the value. Returns the number of values evicted. */
  def evictByValueCondition(removalCondition: UnsafeRow => Boolean): Long

  /**
   * Evict the state via condition on the value, and return the evicted key-value pairs.
   *
   * It is caller's responsibility to consume the whole iterator.
   */
  def evictAndReturnByValueCondition(
      removalCondition: UnsafeRow => Boolean): Iterator[KeyToValuePair]
}

/**
 * This trait is for state manager implementations that support eviction by timestamp. This is for
 * the state manager implementations which maintain the state with event time and can efficiently
 * scan the keys with event time smaller than the given timestamp for eviction.
 */
trait SupportsEvictByTimestamp { self: SymmetricHashJoinStateManager =>
  import SymmetricHashJoinStateManager._

  /**
   * Evict the state by timestamp. Returns the number of values evicted.
   *
   * @param endTimestamp Inclusive upper bound: evicts entries with timestamp <= endTimestamp.
   * @param startTimestamp Exclusive lower bound: entries with timestamp <= startTimestamp are
   *   assumed to have been evicted already (e.g. from the previous batch). When provided,
   *   the scan starts from startTimestamp + 1.
   */
  def evictByTimestamp(endTimestamp: Long, startTimestamp: Option[Long] = None): Long

  /**
   * Evict the state by timestamp and return the evicted key-value pairs.
   *
   * It is caller's responsibility to consume the whole iterator.
   *
   * @param endTimestamp Inclusive upper bound: evicts entries with timestamp <= endTimestamp.
   * @param startTimestamp Exclusive lower bound: entries with timestamp <= startTimestamp are
   *   assumed to have been evicted already (e.g. from the previous batch). When provided,
   *   the scan starts from startTimestamp + 1.
   */
  def evictAndReturnByTimestamp(
      endTimestamp: Long, startTimestamp: Option[Long] = None): Iterator[KeyToValuePair]
}

/**
 * The version 4 of stream-stream join state manager implementation, which is designed to optimize
 * the eviction with watermark. Previous versions require full scan to find the keys to evict,
 * while this version only scans the keys with event time smaller than the watermark.
 *
 * In this implementation, we no longer build a logical array of values; instead, we store the
 * (key, timestamp) -> values in the primary store, and maintain a secondary index of
 * (timestamp, key) to scan the keys to evict for each watermark. To retrieve the values for a key,
 * we perform prefix scan with the key to get all the (key, timestamp) -> values.
 *
 * This implementation leverages (virtual) column family, and uses features which are only
 * available in RocksDB state store provider. Please make sure to set the state store provider to
 * RocksDBStateStoreProvider.
 *
 * Refer to the [[KeyWithTsToValuesStore]] and [[TsWithKeyTypeStore]] for more details.
 */
class SymmetricHashJoinStateManagerV4(
    joinSide: JoinSide,
    inputValueAttributes: Seq[Attribute],
    joinKeys: Seq[Expression],
    stateInfo: Option[StatefulOperatorStateInfo],
    storeConf: StateStoreConf,
    hadoopConf: Configuration,
    partitionId: Int,
    keyToNumValuesStateStoreCkptId: Option[String],
    keyWithIndexToValueStateStoreCkptId: Option[String],
    stateFormatVersion: Int,
    skippedNullValueCount: Option[SQLMetric] = None,
    useStateStoreCoordinator: Boolean = true,
    snapshotOptions: Option[SnapshotOptions] = None,
    joinStoreGenerator: JoinStateManagerStoreGenerator,
    joinKeyOrdinalForWatermark: Option[Int] = None)
  extends SymmetricHashJoinStateManager with SupportsEvictByTimestamp with Logging {

  // TODO: [SPARK-55729] Once the new state manager is integrated to stream-stream join operator,
  //  we should support state data source to understand the state for state format version v4.

  import SymmetricHashJoinStateManager._

  protected val keySchema = StructType(
    joinKeys.zipWithIndex.map { case (k, i) => StructField(s"field$i", k.dataType, k.nullable) })
  protected val keyAttributes = toAttributes(keySchema)
  private val eventTimeColIdxOpt = WatermarkSupport.findEventTimeColumnIndex(
    inputValueAttributes,
    // NOTE: This does not accept multiple event time columns. This is not the same with the
    // operator which we offer the backward compatibility, but it involves too many layers to
    // pass the information. The information is in SQLConf.
    allowMultipleEventTimeColumns = false)

  // When there is no event time column in the value and no watermark ordinal in the key,
  // the secondary index (TsWithKey) will never be used for eviction. Skip writing to it
  // to avoid unnecessary RocksDB merge overhead.
  // TODO(SPARK-56536): This could be further optimized by also considering whether the state
  //   watermark predicate is defined. Even when an event time column exists, the secondary index
  //   is unused if eviction is not possible (e.g., only one side defines a watermark in a time
  //   interval join). That would require propagating the predicate information here.
  private val hasEventTime: Boolean =
    eventTimeColIdxOpt.isDefined || joinKeyOrdinalForWatermark.isDefined

  private val random = new scala.util.Random(System.currentTimeMillis())
  private val bucketCountForNoEventTime = 1024
  private val extractEventTimeFn: UnsafeRow => Long = { row =>
    eventTimeColIdxOpt match {
      case Some(idx) =>
        val attr = inputValueAttributes(idx)

        if (attr.dataType.isInstanceOf[StructType]) {
          // NOTE: We assume this is window struct, as same as WatermarkSupport.watermarkExpression
          row.getStruct(idx, 2).getLong(1)
        } else {
          row.getLong(idx)
        }

      case _ =>
        // When event time column is not available, we will use random bucketing strategy to decide
        // where the new value will be stored. There is a trade-off between the bucket size and the
        // number of values in each bucket; we can tune the bucket size with the configuration if
        // we figure out the magic number to not work well.
        random.nextInt(bucketCountForNoEventTime)
    }
  }

  private val extractEventTimeFnFromKey: UnsafeRow => Option[Long] = { row =>
    joinKeyOrdinalForWatermark.map { idx =>
      val attr = keyAttributes(idx)
      if (attr.dataType.isInstanceOf[StructType]) {
        // NOTE: We assume this is window struct, as same as WatermarkSupport.watermarkExpression
        row.getStruct(idx, 2).getLong(1)
      } else {
        row.getLong(idx)
      }
    }
  }

  private val dummySchema = StructType(
    Seq(StructField("dummy", NullType, nullable = true))
  )

  // V4 uses a single store with VCFs (not separate keyToNumValues/keyWithIndexToValue stores).
  // Use the keyToNumValues checkpoint ID for loading the correct committed version.
  private val stateStoreCkptId: Option[String] = keyToNumValuesStateStoreCkptId
  private val handlerSnapshotOptions: Option[HandlerSnapshotOptions] = snapshotOptions.map { opts =>
    HandlerSnapshotOptions(
      snapshotVersion = opts.snapshotVersion,
      endVersion = opts.endVersion,
      startStateStoreCkptId = opts.startKeyToNumValuesStateStoreCkptId,
      endStateStoreCkptId = opts.endKeyToNumValuesStateStoreCkptId)
  }

  private var stateStoreProvider: StateStoreProvider = _

  // We will use the dummy schema for the default CF since we will register CF separately.
  private val stateStore = getStateStore(
    dummySchema, dummySchema, useVirtualColumnFamilies = true,
    NoPrefixKeyStateEncoderSpec(dummySchema), useMultipleValuesPerKey = false
  )

  private def getStateStore(
      keySchema: StructType,
      valueSchema: StructType,
      useVirtualColumnFamilies: Boolean,
      keyStateEncoderSpec: KeyStateEncoderSpec,
      useMultipleValuesPerKey: Boolean): StateStore = {
    val storeName = StateStoreId.DEFAULT_STORE_NAME
    val storeProviderId = StateStoreProviderId(stateInfo.get, partitionId, storeName)
    val store = if (useStateStoreCoordinator) {
      assert(handlerSnapshotOptions.isEmpty, "Should not use state store coordinator " +
        "when reading state as data source.")
      joinStoreGenerator.getStore(
        storeProviderId, keySchema, valueSchema, keyStateEncoderSpec,
        stateInfo.get.storeVersion, stateStoreCkptId, None, useVirtualColumnFamilies,
        useMultipleValuesPerKey, storeConf, hadoopConf)
    } else {
      // This class will manage the state store provider by itself.
      stateStoreProvider = StateStoreProvider.createAndInit(
        storeProviderId, keySchema, valueSchema, keyStateEncoderSpec,
        useColumnFamilies = useVirtualColumnFamilies,
        storeConf, hadoopConf, useMultipleValuesPerKey = useMultipleValuesPerKey,
        stateSchemaProvider = None)
      if (handlerSnapshotOptions.isDefined) {
        if (!stateStoreProvider.isInstanceOf[SupportsFineGrainedReplay]) {
          throw StateStoreErrors.stateStoreProviderDoesNotSupportFineGrainedReplay(
            stateStoreProvider.getClass.toString)
        }
        val opts = handlerSnapshotOptions.get
        stateStoreProvider.asInstanceOf[SupportsFineGrainedReplay]
          .replayStateFromSnapshot(
            opts.snapshotVersion,
            opts.endVersion,
            readOnly = true,
            opts.startStateStoreCkptId,
            opts.endStateStoreCkptId)
      } else {
        stateStoreProvider.getStore(stateInfo.get.storeVersion, stateStoreCkptId)
      }
    }
    logInfo(log"Loaded store ${MDC(STATE_STORE_ID, store.id)}")
    store
  }

  private val keyWithTsToValues = new KeyWithTsToValuesStore

  private val tsWithKey = new TsWithKeyTypeStore

  override def append(key: UnsafeRow, value: UnsafeRow, matched: Boolean): Unit = {
    val eventTime = extractEventTimeFnFromKey(key).getOrElse(extractEventTimeFn(value))
    // We always do blind merge for appending new value.
    keyWithTsToValues.append(key, eventTime, value, matched)
    if (hasEventTime) {
      tsWithKey.add(eventTime, key)
    }
  }

  override def getJoinedRows(
      key: UnsafeRow,
      generateJoinedRow: InternalRow => JoinedRow,
      predicate: JoinedRow => Boolean,
      timestampRange: Option[(Long, Long)] = None,
      skipUpdatingMatchedFlag: Boolean = false): Iterator[JoinedRow] = {
    def getJoinedRowsFromTsAndValues(
        ts: Long,
        valuesAndMatched: Array[ValueAndMatchPair]): Iterator[JoinedRow] = {
      new NextIterator[JoinedRow] {
        private var currentIndex = 0

        private var shouldUpdateValuesIntoStateStore = false

        override protected def getNext(): JoinedRow = {
          var ret: JoinedRow = null
          while (ret == null && currentIndex < valuesAndMatched.length) {
            val vmp = valuesAndMatched(currentIndex)

            val joinedRow = generateJoinedRow(vmp.value)
            if (predicate(joinedRow)) {
              if (!skipUpdatingMatchedFlag && !vmp.matched) {
                valuesAndMatched(currentIndex) = vmp.copy(matched = true)
                shouldUpdateValuesIntoStateStore = true
              }

              ret = joinedRow
            }

            currentIndex += 1
          }

          if (ret == null) {
            assert(currentIndex == valuesAndMatched.length)
            finished = true
            null
          } else {
            ret
          }
        }

        override protected def close(): Unit = {
          if (shouldUpdateValuesIntoStateStore) {
            val updatedValuesWithMatched = valuesAndMatched.map { vmp =>
              (vmp.value, vmp.matched)
            }.toSeq
            keyWithTsToValues.put(key, ts, updatedValuesWithMatched)
          }
        }
      }
    }

    val ret = extractEventTimeFnFromKey(key) match {
      case Some(ts) =>
        val valuesAndMatchedIter = keyWithTsToValues.get(key, ts)
        getJoinedRowsFromTsAndValues(ts, valuesAndMatchedIter.toArray)

      case _ =>
        val (minTs, maxTs) = timestampRange.getOrElse((Long.MinValue, Long.MaxValue))
        keyWithTsToValues.getValuesInRange(key, minTs, maxTs).flatMap { result =>
          val ts = result.timestamp
          val valuesAndMatched = result.values.toArray
          getJoinedRowsFromTsAndValues(ts, valuesAndMatched)
        }
    }
    ret.filter(_ != null)
  }

  override def getJoinedRowsAndRemoveMatched(
      key: UnsafeRow,
      generateJoinedRow: InternalRow => JoinedRow,
      predicate: JoinedRow => Boolean): Iterator[JoinedRow] = {
    def getJoinedRowsFromTsAndValues(
        ts: Long,
        valuesAndMatched: Array[ValueAndMatchPair]): Iterator[JoinedRow] = {
      new NextIterator[JoinedRow] {
        private var currentIndex = 0

        private val toKeep = new scala.collection.mutable.ArrayBuffer[ValueAndMatchPair]()
        private var hasRemovals = false

        override protected def getNext(): JoinedRow = {
          var ret: JoinedRow = null
          while (ret == null && currentIndex < valuesAndMatched.length) {
            val vmp = valuesAndMatched(currentIndex)

            if (vmp.matched) {
              // See the NOTE in the method doc about rationale.
              hasRemovals = true
            } else {
              val joinedRow = generateJoinedRow(vmp.value)
              if (predicate(joinedRow)) {
                hasRemovals = true
                ret = joinedRow
              } else {
                toKeep += vmp
              }
            }

            currentIndex += 1
          }

          if (ret == null) {
            assert(currentIndex == valuesAndMatched.length)
            finished = true
            null
          } else {
            ret
          }
        }

        override protected def close(): Unit = {
          if (hasRemovals) {
            if (toKeep.isEmpty) {
              keyWithTsToValues.remove(key, ts)
              tsWithKey.remove(key, ts)
            } else {
              val valuesToPut = toKeep.map(vmp => (vmp.value, vmp.matched)).toSeq
              keyWithTsToValues.put(key, ts, valuesToPut)
            }
          }
        }
      }
    }

    val ret = extractEventTimeFnFromKey(key) match {
      case Some(ts) =>
        val valuesAndMatchedIter = keyWithTsToValues.get(key, ts)
        getJoinedRowsFromTsAndValues(ts, valuesAndMatchedIter.toArray)

      case _ =>
        keyWithTsToValues.getValues(key).flatMap { result =>
          val ts = result.timestamp
          val valuesAndMatched = result.values.toArray
          getJoinedRowsFromTsAndValues(ts, valuesAndMatched)
        }
    }
    ret.filter(_ != null)
  }

  /**
   * NOTE: The entry provided by Iterator.next() will be reused. It is a caller's responsibility
   * to copy it properly if caller needs to keep the reference after next() is called again.
   */
  override def iterator: Iterator[KeyToValuePair] = {
    val reusableKeyToValuePair = KeyToValuePair()
    keyWithTsToValues.iterator().map { kv =>
      reusableKeyToValuePair.withNew(kv.key, kv.value, kv.matched)
    }
  }

  override def evictByTimestamp(endTimestamp: Long, startTimestamp: Option[Long] = None): Long = {
    require(hasEventTime,
      "evictByTimestamp requires event time; secondary index was not populated")
    var removed = 0L
    tsWithKey.scanEvictedKeys(endTimestamp, startTimestamp).foreach { evicted =>
      val key = evicted.key
      val timestamp = evicted.timestamp
      val numValues = evicted.numValues

      // Remove from both primary and secondary stores
      keyWithTsToValues.remove(key, timestamp)
      tsWithKey.remove(key, timestamp)

      removed += numValues
    }
    removed
  }

  override def evictAndReturnByTimestamp(
      endTimestamp: Long, startTimestamp: Option[Long] = None): Iterator[KeyToValuePair] = {
    require(hasEventTime,
      "evictAndReturnByTimestamp requires event time; secondary index was not populated")
    val reusableKeyToValuePair = KeyToValuePair()

    tsWithKey.scanEvictedKeys(endTimestamp, startTimestamp).flatMap { evicted =>
      val key = evicted.key
      val timestamp = evicted.timestamp
      val values = keyWithTsToValues.get(key, timestamp)

      // Remove from both primary and secondary stores
      keyWithTsToValues.remove(key, timestamp)
      tsWithKey.remove(key, timestamp)

      values.map { value =>
        reusableKeyToValuePair.withNew(key, value)
      }
    }
  }

  override def commit(): Unit = {
    stateStore.commit()
    logDebug("Committed, metrics = " + stateStore.metrics)
  }

  override def abortIfNeeded(): Unit = {
    if (!stateStore.hasCommitted) {
      logInfo(log"Aborted store ${MDC(STATE_STORE_ID, stateStore.id)}")
      stateStore.abort()
    }
    // If this class manages a state store provider by itself, it should take care of closing
    // provider instance as well.
    if (stateStoreProvider != null) {
      stateStoreProvider.close()
    }
  }

  // Clean up any state store resources if necessary at the end of the task
  Option(TaskContext.get()).foreach { _.addTaskCompletionListener[Unit] { _ => abortIfNeeded() } }

  class GetValuesResult(var timestamp: Long = -1, var values: Seq[ValueAndMatchPair] = Seq.empty) {
    def withNew(newTimestamp: Long, newValues: Seq[ValueAndMatchPair]): GetValuesResult = {
      this.timestamp = newTimestamp
      this.values = newValues
      this
    }
  }

  /**
   * The primary store to store the key-value pairs.
   *
   * The state format of the primary store is following:
   * [key][timestamp (event time)] -> [(value, matched), (value, matched), ...]
   *
   * The values are bucketed by event time to facilitate efficient eviction by watermark; the
   * secondary index will provide the way to scan the key + timestamp pairs for the eviction, and
   * it will be easy to perform retrieval/removal of the values based on key + timestamp pairs.
   * There is no case where we evict only part of the values for the same key + timestamp.
   *
   * The matched flag is used to indicate whether the value has been matched with any row from the
   * other side.
   */
  private class KeyWithTsToValuesStore {
    private val valueRowConverter = StreamingSymmetricHashJoinValueRowConverter.create(
      inputValueAttributes, stateFormatVersion = 4)

    // Set up virtual column family name in the store if it is being used
    private val colFamilyName = getStateStoreName(joinSide, KeyWithTsToValuesType)

    private val keySchemaWithTimestamp = TimestampKeyStateEncoder.keySchemaWithTimestamp(keySchema)
    private val detachTimestampProjection: UnsafeProjection =
      TimestampKeyStateEncoder.getDetachTimestampProjection(keySchemaWithTimestamp)
    private val attachTimestampProjection: UnsafeProjection =
      TimestampKeyStateEncoder.getAttachTimestampProjection(keySchema)

    // Create the specific column family in the store for this join side's KeyWithTsToValuesStore.
    stateStore.createColFamilyIfAbsent(
      colFamilyName,
      keySchema,
      valueRowConverter.valueAttributes.toStructType,
      TimestampAsPostfixKeyStateEncoderSpec(keySchemaWithTimestamp),
      useMultipleValuesPerKey = true
    )

    private def createKeyRow(key: UnsafeRow, timestamp: Long): UnsafeRow = {
      TimestampKeyStateEncoder.attachTimestamp(
        attachTimestampProjection, keySchemaWithTimestamp, key, timestamp)
    }

    def append(key: UnsafeRow, timestamp: Long, value: UnsafeRow, matched: Boolean): Unit = {
      val valueWithMatched = valueRowConverter.convertToValueRow(value, matched)
      stateStore.merge(createKeyRow(key, timestamp), valueWithMatched, colFamilyName)
    }

    def put(
        key: UnsafeRow,
        timestamp: Long,
        valuesWithMatched: Seq[(UnsafeRow, Boolean)]): Unit = {
      // copy() is required because convertToValueRow reuses its internal UnsafeProjection output
      // TODO: [SPARK-55732] StateStore.putList should allow iterator to be passed in, so that we
      //  don't need to materialize the array and copy the values here.
      val valuesToPut = valuesWithMatched.map { case (value, matched) =>
        valueRowConverter.convertToValueRow(value, matched).copy()
      }.toArray
      stateStore.putList(createKeyRow(key, timestamp), valuesToPut, colFamilyName)
    }

    def get(key: UnsafeRow, timestamp: Long): Iterator[ValueAndMatchPair] = {
      stateStore.valuesIterator(createKeyRow(key, timestamp), colFamilyName).map { valueRow =>
        valueRowConverter.convertValue(valueRow)
      }
    }

    // NOTE: We do not have a case where we only remove a part of values. Even if that is needed
    // we handle it via put() with writing a new array.
    def remove(key: UnsafeRow, timestamp: Long): Unit = {
      stateStore.remove(createKeyRow(key, timestamp), colFamilyName)
    }

    // NOTE: This assumes we consume the whole iterator to trigger completion.
    def getValues(key: UnsafeRow): Iterator[GetValuesResult] = {
      getValuesInRange(key, Long.MinValue, Long.MaxValue)
    }

    /**
     * Returns entries where minTs <= timestamp <= maxTs (both inclusive), grouped by timestamp.
     * When maxTs is bounded (< Long.MaxValue), uses rangeScanWithMultiValues for efficient
     * range access; falls back to prefixScan otherwise to stay within the key's scope.
     *
     * When prefixScan is used (maxTs == Long.MaxValue), entries outside [minTs, maxTs] are
     * filtered out so both code paths produce identical results.
     */
    def getValuesInRange(
        key: UnsafeRow, minTs: Long, maxTs: Long): Iterator[GetValuesResult] = {
      val reusableGetValuesResult = new GetValuesResult()
      // Only use rangeScan when maxTs < Long.MaxValue, since rangeScan requires
      // an exclusive end key (maxTs + 1) which would overflow at Long.MaxValue.
      val useRangeScan = maxTs < Long.MaxValue

      new NextIterator[GetValuesResult] {
        private val iter = if (useRangeScan) {
          // startKey must be copied because the second createKeyRow call below reuses
          // the same projection buffer and would otherwise overwrite its contents.
          // endKey does not need a copy: rangeScanWithMultiValues encodes both bounds
          // to independent byte arrays eagerly at call time, and the scope of endKey
          // ends with the call of rangeScanWithMultiValues.
          val startKey = createKeyRow(key, minTs).copy()
          // rangeScanWithMultiValues endKey is exclusive, so use maxTs + 1
          val endKey = Some(createKeyRow(key, maxTs + 1))
          stateStore.rangeScanWithMultiValues(Some(startKey), endKey, colFamilyName)
        } else {
          stateStore.prefixScanWithMultiValues(key, colFamilyName)
        }

        private var currentTs = -1L
        private var pastUpperBound = false
        private val valueAndMatchPairs = scala.collection.mutable.ArrayBuffer[ValueAndMatchPair]()

        private def flushAccumulated(): GetValuesResult = {
          if (valueAndMatchPairs.nonEmpty) {
            val result = reusableGetValuesResult.withNew(
              currentTs, valueAndMatchPairs.toList)
            currentTs = -1L
            valueAndMatchPairs.clear()
            result
          } else {
            finished = true
            null
          }
        }

        @tailrec
        override protected def getNext(): GetValuesResult = {
          if (pastUpperBound || !iter.hasNext) {
            flushAccumulated()
          } else {
            val unsafeRowPair = iter.next()
            val ts = TimestampKeyStateEncoder.extractTimestamp(unsafeRowPair.key)

            if (useRangeScan && (ts < minTs || ts > maxTs)) {
              throw StateStoreErrors.streamStreamJoinRangeScanTimestampOutOfRange(
                ts, minTs, maxTs)
            }

            if (ts > maxTs) {
              pastUpperBound = true
              getNext()
            } else if (ts < minTs) {
              getNext()
            } else if (currentTs == -1L || currentTs == ts) {
              currentTs = ts
              valueAndMatchPairs += valueRowConverter.convertValue(unsafeRowPair.value)
              getNext()
            } else {
              // Timestamp changed -- flush previous group before starting new one
              val prevTs = currentTs
              val prevValues = valueAndMatchPairs.toList

              currentTs = ts
              valueAndMatchPairs.clear()
              valueAndMatchPairs += valueRowConverter.convertValue(unsafeRowPair.value)

              reusableGetValuesResult.withNew(prevTs, prevValues)
            }
          }
        }

        override protected def close(): Unit = iter.close()
      }
    }

    def iterator(): Iterator[KeyAndTsToValuePair] = {
      val iter = stateStore.iteratorWithMultiValues(colFamilyName)
      val reusableKeyAndTsToValuePair = KeyAndTsToValuePair()
      iter.map { kv =>
        val keyRow = detachTimestampProjection(kv.key)
        val ts = TimestampKeyStateEncoder.extractTimestamp(kv.key)
        val value = valueRowConverter.convertValue(kv.value)

        reusableKeyAndTsToValuePair.withNew(keyRow, ts, value)
      }
    }
  }

  /**
   * The secondary index for efficient state removal with watermark.
   *
   * The state format of the secondary index is following:
   * [timestamp (adjusted for ordering, 8 bytes)][key] -> [list of empty values]
   *
   * The value part is used to track the number of values for the same (key, timestamp) in the
   * primary store, so that we can track the number of values being removed for each eviction and
   * update metrics accordingly. Alternatively, we can also maintain an integer count in the value
   * part, but we found blind merge and count later to be more efficient than read-modify-write.
   */
  private class TsWithKeyTypeStore {
    private val valueStructType = StructType(Array(StructField("__dummy__", NullType)))
    private val EMPTY_ROW =
      UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null))

    // Set up virtual column family name in the store if it is being used
    private val colFamilyName = getStateStoreName(joinSide, TsWithKeyType)

    private val keySchemaWithTimestamp = TimestampKeyStateEncoder.keySchemaWithTimestamp(keySchema)
    private val detachTimestampProjection: UnsafeProjection =
      TimestampKeyStateEncoder.getDetachTimestampProjection(keySchemaWithTimestamp)
    private val attachTimestampProjection: UnsafeProjection =
      TimestampKeyStateEncoder.getAttachTimestampProjection(keySchema)

    // Create the specific column family in the store for this join side's TsWithKeyStore.
    // Mark as internal so that numKeys counts only primary data, not the secondary index.
    stateStore.createColFamilyIfAbsent(
      colFamilyName,
      keySchema,
      valueStructType,
      TimestampAsPrefixKeyStateEncoderSpec(keySchemaWithTimestamp),
      useMultipleValuesPerKey = true,
      isInternal = true
    )

    // Returns an UnsafeRow backed by a reused projection buffer. Callers that need to
    // hold the row beyond the immediate state store call must invoke copy() on the result.
    private def createKeyRow(key: UnsafeRow, timestamp: Long): UnsafeRow = {
      TimestampKeyStateEncoder.attachTimestamp(
        attachTimestampProjection, keySchemaWithTimestamp, key, timestamp)
    }

    def add(timestamp: Long, key: UnsafeRow): Unit = {
      stateStore.merge(createKeyRow(key, timestamp), EMPTY_ROW, colFamilyName)
    }

    def remove(key: UnsafeRow, timestamp: Long): Unit = {
      stateStore.remove(createKeyRow(key, timestamp), colFamilyName)
    }

    case class EvictedKeysResult(key: UnsafeRow, timestamp: Long, numValues: Int)

    // Reusable default key row for scan boundary construction; see
    // [[RangeScanBoundaryUtils]] for rationale. Safe to reuse because createKeyRow
    // only reads this row (via BoundReference evaluations) and writes to the
    // projection's own internal buffer. Correctness relies on real stored entries
    // never having internally-null key fields, which is preserved by join-key
    // expressions being evaluated via the user's expression encoder. Preserve this
    // invariant if you change how entries are written.
    private lazy val defaultKey: UnsafeRow = RangeScanBoundaryUtils.defaultUnsafeRow(keySchema)

    /**
     * Build a scan boundary row for rangeScan. The TsWithKeyTypeStore uses
     * TimestampAsPrefixKeyStateEncoder, which encodes the row as [timestamp][key_fields].
     * We need a full-schema row (not just the timestamp) because the encoder expects all
     * key columns to be present. Default values are used for the key fields since only the
     * timestamp matters for ordering in the prefix encoder.
     */
    private def createScanBoundaryRow(timestamp: Long): UnsafeRow = {
      createKeyRow(defaultKey, timestamp).copy()
    }

    /**
     * Scan keys eligible for eviction within the timestamp range.
     *
     * This assumes we consume the whole iterator to trigger completion.
     *
     * @param endTimestamp Inclusive upper bound: entries with timestamp <= endTimestamp are
     *   eligible for eviction.
     * @param startTimestamp Exclusive lower bound: entries with timestamp <= startTimestamp
     *   are assumed to have been evicted already. The scan starts from startTimestamp + 1.
     */
    def scanEvictedKeys(
        endTimestamp: Long,
        startTimestamp: Option[Long] = None): Iterator[EvictedKeysResult] = {
      // If startTimestamp == Long.MaxValue, everything has already been evicted;
      // nothing can match, so return immediately.
      if (startTimestamp.contains(Long.MaxValue)) {
        return Iterator.empty
      }

      // rangeScanWithMultiValues: startKey is inclusive, endKey is exclusive.
      // startTimestamp is exclusive (already evicted), so we seek from st + 1.
      val startKeyRow = startTimestamp.map { st =>
        createScanBoundaryRow(st + 1)
      }
      // endTimestamp is inclusive, so we use endTimestamp + 1 as the exclusive upper bound.
      // When endTimestamp == Long.MaxValue we cannot add 1, so endKeyRow is None. This is
      // safe because rangeScanWithMultiValues with no end key uses the column-family prefix
      // as the upper bound, naturally scoping the scan within this column family.
      val endKeyRow = if (endTimestamp < Long.MaxValue) {
        Some(createScanBoundaryRow(endTimestamp + 1))
      } else {
        None
      }
      val evictIterator = stateStore.rangeScanWithMultiValues(startKeyRow, endKeyRow, colFamilyName)
      new NextIterator[EvictedKeysResult]() {
        var currentKeyRow: UnsafeRow = null
        var currentEventTime: Long = -1L
        var count: Int = 0
        var isBeyondUpperBound: Boolean = false

        override protected def getNext(): EvictedKeysResult = {
          var ret: EvictedKeysResult = null
          while (evictIterator.hasNext && ret == null && !isBeyondUpperBound) {
            val kv = evictIterator.next()
            val keyRow = detachTimestampProjection(kv.key)
            val ts = TimestampKeyStateEncoder.extractTimestamp(kv.key)

            if (keyRow == currentKeyRow && ts == currentEventTime) {
              // new value with same (key, ts)
              count += 1
            } else if (ts > endTimestamp) {
              // we found the timestamp beyond the range - we shouldn't continue further
              isBeyondUpperBound = true

              // We don't need to construct the last (key, ts) into EvictedKeysResult - the code
              // after loop will handle that if there is leftover. That said, we do not reset the
              // current (key, ts) info here.
            } else if (currentKeyRow == null && currentEventTime == -1L) {
              // first value to process
              currentKeyRow = keyRow.copy()
              currentEventTime = ts
              count = 1
            } else {
              // construct the last (key, ts) into EvictedKeysResult
              ret = EvictedKeysResult(currentKeyRow, currentEventTime, count)

              // register the next (key, ts) to process
              currentKeyRow = keyRow.copy()
              currentEventTime = ts
              count = 1
            }
          }

          if (ret != null) {
            ret
          } else if (count > 0) {
            // there is a final leftover (key, ts) to return
            ret = EvictedKeysResult(currentKeyRow, currentEventTime, count)

            // we shouldn't continue further
            currentKeyRow = null
            currentEventTime = -1L
            count = 0

            ret
          } else {
            finished = true
            null
          }
        }

        override protected def close(): Unit = {
          evictIterator.close()
        }
      }
    }
  }

  override def get(key: UnsafeRow): Iterator[UnsafeRow] = {
    keyWithTsToValues.getValues(key).flatMap { result =>
      result.values.map(_.value)
    }
  }

  def metrics: StateStoreMetrics = stateStore.metrics

  def getLatestCheckpointInfo(): JoinerStateStoreCkptInfo = {
    val keyToNumValuesCkptInfo = stateStore.getStateStoreCheckpointInfo()
    val keyWithIndexToValueCkptInfo = stateStore.getStateStoreCheckpointInfo()

    assert(keyToNumValuesCkptInfo == keyWithIndexToValueCkptInfo)

    JoinerStateStoreCkptInfo(keyToNumValuesCkptInfo, keyWithIndexToValueCkptInfo)
  }
}

/**
 * Helper class to manage state required by a single side of
 * [[org.apache.spark.sql.execution.streaming.StreamingSymmetricHashJoinExec]].
 * The interface of this class is basically that of a multi-map:
 * - Get: Returns an iterator of multiple values for given key
 * - Append: Append a new value to the given key
 * - Remove Data by predicate: Drop any state using a predicate condition on keys or values
 *
 * @param joinSide              Defines the join side
 * @param inputValueAttributes  Attributes of the input row which will be stored as value
 * @param joinKeys              Expressions to generate rows that will be used to key the value rows
 * @param stateInfo             Information about how to retrieve the correct version of state
 * @param storeConf             Configuration for the state store.
 * @param hadoopConf            Hadoop configuration for reading state data from storage
 * @param partitionId           A partition ID of source RDD.
 * @param stateFormatVersion    The version of format for state.
 * @param skippedNullValueCount The instance of SQLMetric tracking the number of skipped null
 *                              values.
 * @param useStateStoreCoordinator  Whether to use a state store coordinator to maintain the state
 *                                  store providers being used in this class. If true, Spark will
 *                                  take care of management for state store providers, e.g. running
 *                                  maintenance task for these providers.
 * @param snapshotOptions       Options controlling snapshot-based state replay for the state data
 *                              source reader.
 * @param joinStoreGenerator    The generator to create state store instances, re-using the same
 *                              instance when the join implementation uses virtual column families
 *                              for join version 3.
 *
 * Internally, the key -> multiple values is stored in two [[StateStore]]s.
 * - Store 1 ([[KeyToNumValuesStore]]) maintains mapping between key -> number of values
 * - Store 2 ([[KeyWithIndexToValueStore]]) maintains mapping; the mapping depends on the state
 *   format version:
 *   - version 1: [(key, index) -> value]
 *   - version 2: [(key, index) -> (value, matched)]
 * - Put:   update count in KeyToNumValuesStore,
 *          insert new (key, count) -> value in KeyWithIndexToValueStore
 * - Get:   read count from KeyToNumValuesStore,
 *          read each of the n values in KeyWithIndexToValueStore
 * - Remove state by predicate on keys:
 *          scan all keys in KeyToNumValuesStore to find keys that do match the predicate,
 *          delete from key from KeyToNumValuesStore, delete values in KeyWithIndexToValueStore
 * - Remove state by condition on values:
 *          scan all elements in KeyWithIndexToValueStore to find values that match
 *          the predicate, delete corresponding (key, indexToDelete) from KeyWithIndexToValueStore
 *          by overwriting with the value of (key, maxIndex), and removing [(key, maxIndex),
 *          decrement corresponding num values in KeyToNumValuesStore
 */
abstract class SymmetricHashJoinStateManagerBase(
    joinSide: JoinSide,
    inputValueAttributes: Seq[Attribute],
    joinKeys: Seq[Expression],
    stateInfo: Option[StatefulOperatorStateInfo],
    storeConf: StateStoreConf,
    hadoopConf: Configuration,
    partitionId: Int,
    keyToNumValuesStateStoreCkptId: Option[String],
    keyWithIndexToValueStateStoreCkptId: Option[String],
    stateFormatVersion: Int,
    skippedNullValueCount: Option[SQLMetric] = None,
    useStateStoreCoordinator: Boolean = true,
    snapshotOptions: Option[SnapshotOptions] = None,
    joinStoreGenerator: JoinStateManagerStoreGenerator)
  extends SymmetricHashJoinStateManager
  with SupportsEvictByCondition
  with SupportsIndexedKeys
  with Logging {

  import SymmetricHashJoinStateManager._

  protected val keySchema = StructType(
    joinKeys.zipWithIndex.map { case (k, i) => StructField(s"field$i", k.dataType, k.nullable) })
  protected val keyAttributes = toAttributes(keySchema)

  protected val keyToNumValues = new KeyToNumValuesStore(
    stateFormatVersion,
    snapshotOptions.map(_.getKeyToNumValuesHandlerOpts()))
  protected val keyWithIndexToValue = new KeyWithIndexToValueStore(
    stateFormatVersion,
    snapshotOptions.map(_.getKeyWithIndexToValueHandlerOpts()))

  /*
  =====================================================
                  Public methods
  =====================================================
   */

  /** Get all the values of a key */
  def get(key: UnsafeRow): Iterator[UnsafeRow] = {
    val numValues = keyToNumValues.get(key)
    keyWithIndexToValue.getAll(key, numValues).map(_.value)
  }

  /** Append a new value to the key */
  def append(key: UnsafeRow, value: UnsafeRow, matched: Boolean): Unit = {
    val numExistingValues = keyToNumValues.get(key)
    keyWithIndexToValue.put(key, numExistingValues, value, matched)
    keyToNumValues.put(key, numExistingValues + 1)
  }

  /**
   * Find the first non-null value index starting from the end and going up to stopIndex.
   * Used by swap-with-last compaction in both [[getJoinedRowsAndRemoveMatched]] and
   * [[evictAndReturnByValueCondition]].
   */
  protected def getRightMostNonNullIndex(
      key: UnsafeRow, stopIndex: Long, numValues: Long): Option[Long] = {
    (numValues - 1 to stopIndex by -1).find { idx =>
      keyWithIndexToValue.get(key, idx) != null
    }
  }

  /**
   * Remove the value at the given index using swap-with-last compaction: the last element is
   * moved into the hole, trailing nulls are cleaned up, and the logical array is shortened.
   *
   * @return the updated numValues after removal
   */
  protected def removeValueAtIndex(key: UnsafeRow, index: Long, numValues: Long): Long = {
    var currentNumValues = numValues
    if (index != currentNumValues - 1) {
      val valuePairAtMaxIndex = keyWithIndexToValue.get(key, currentNumValues - 1)
      if (valuePairAtMaxIndex != null) {
        keyWithIndexToValue.put(key, index, valuePairAtMaxIndex.value,
          valuePairAtMaxIndex.matched)
      } else {
        val nonNullIndex = getRightMostNonNullIndex(key, index + 1, currentNumValues)
          .getOrElse(index)
        if (nonNullIndex != index) {
          val valuePair = keyWithIndexToValue.get(key, nonNullIndex)
          keyWithIndexToValue.put(key, index, valuePair.value, valuePair.matched)
        }

        if (nonNullIndex != currentNumValues - 1) {
          logWarning(log"`keyWithIndexToValue` returns a null value for indices " +
            log"with range from startIndex=${MDC(START_INDEX, nonNullIndex + 1)} " +
            log"and endIndex=${MDC(END_INDEX, currentNumValues - 1)}.")
        }

        (currentNumValues - 1 to nonNullIndex + 1 by -1).foreach { removeIndex =>
          keyWithIndexToValue.remove(key, removeIndex)
          currentNumValues -= 1
        }
      }
    }
    keyWithIndexToValue.remove(key, currentNumValues - 1)
    currentNumValues - 1
  }

  /**
   * Streaming implementation that processes one value at a time, using the same swap-with-last
   * compaction strategy as [[evictAndReturnByValueCondition]] to avoid materializing all values
   * into memory (a single key can have millions of values in v1-v3).
   *
   * The iterator must be consumed fully without any other operations on this manager
   * or the underlying store being interleaved.
   */
  override def getJoinedRowsAndRemoveMatched(
      key: UnsafeRow,
      generateJoinedRow: InternalRow => JoinedRow,
      predicate: JoinedRow => Boolean): Iterator[JoinedRow] = {
    new NextIterator[JoinedRow] {
      private var numValues: Long = keyToNumValues.get(key)
      private var index: Long = 0L
      private var valueRemoved: Boolean = false

      override def getNext(): JoinedRow = {
        while (index < numValues) {
          val valuePair = keyWithIndexToValue.get(key, index)
          if (valuePair == null) {
            handleNullValuePair(key, index, numValues)
            skippedNullValueCount.foreach(_ += 1L)
            index += 1
          } else if (valuePair.matched) {
            // See the NOTE in the method doc about rationale.
            numValues = removeValueAtIndex(key, index, numValues)
            valueRemoved = true
          } else {
            val joinedRow = generateJoinedRow(valuePair.value)
            if (predicate(joinedRow)) {
              numValues = removeValueAtIndex(key, index, numValues)
              valueRemoved = true
              return joinedRow
            } else {
              index += 1
            }
          }
        }

        if (valueRemoved) {
          if (numValues >= 1) {
            keyToNumValues.put(key, numValues)
          } else {
            keyToNumValues.remove(key)
          }
        }
        finished = true
        null
      }

      override def close(): Unit = {}
    }
  }

  /**
   * Get all the matched values for given join condition, with marking matched.
   * This method is designed to mark joined rows properly without exposing internal index of row.
   *
   * @param skipUpdatingMatchedFlag If true, do not update the matched flag even when the row
   *                                matches.
   */
  def getJoinedRows(
      key: UnsafeRow,
      generateJoinedRow: InternalRow => JoinedRow,
      predicate: JoinedRow => Boolean,
      timestampRange: Option[(Long, Long)] = None,
      skipUpdatingMatchedFlag: Boolean = false): Iterator[JoinedRow] = {
    val numValues = keyToNumValues.get(key)
    keyWithIndexToValue.getAll(key, numValues).map { keyIdxToValue =>
      val joinedRow = generateJoinedRow(keyIdxToValue.value)
      if (predicate(joinedRow)) {
        if (!skipUpdatingMatchedFlag && !keyIdxToValue.matched) {
          keyWithIndexToValue.put(key, keyIdxToValue.valueIndex, keyIdxToValue.value,
            matched = true)
        }
        joinedRow
      } else {
        null
      }
    }.filter(_ != null)
  }

  /** Remove using a predicate on keys. */
  override def evictByKeyCondition(removalCondition: UnsafeRow => Boolean): Long = {
    var numRemoved = 0L
    keyToNumValues.iterator.foreach { keyAndNumValues =>
      val key = keyAndNumValues.key
      if (removalCondition(key)) {
        val numValue = keyAndNumValues.numValue

        (0L until numValue).foreach { idx =>
          keyWithIndexToValue.remove(key, idx)
        }

        numRemoved += numValue
        keyToNumValues.remove(key)
      }
    }
    numRemoved
  }

  /**
   * Remove using a predicate on keys.
   *
   * This produces an iterator over the (key, value, matched) tuples satisfying condition(key),
   * where the underlying store is updated as a side-effect of producing next.
   *
   * This implies the iterator must be consumed fully without any other operations on this manager
   * or the underlying store being interleaved.
   */
  override def evictAndReturnByKeyCondition(
      removalCondition: UnsafeRow => Boolean): Iterator[KeyToValuePair] = {
    new NextIterator[KeyToValuePair] {

      private val allKeyToNumValues = keyToNumValues.iterator

      private var currentKeyToNumValue: KeyAndNumValues = null
      private var currentValues: Iterator[KeyWithIndexAndValue] = null

      private def currentKey = currentKeyToNumValue.key

      private val reusedRet = new KeyToValuePair()

      private def getAndRemoveValue(): KeyToValuePair = {
        val keyWithIndexAndValue = currentValues.next()
        keyWithIndexToValue.remove(currentKey, keyWithIndexAndValue.valueIndex)
        reusedRet.withNew(currentKey, keyWithIndexAndValue.value, keyWithIndexAndValue.matched)
      }

      override def getNext(): KeyToValuePair = {
        // If there are more values for the current key, remove and return the next one.
        if (currentValues != null && currentValues.hasNext) {
          return getAndRemoveValue()
        }

        // If there weren't any values left, try and find the next key that satisfies the removal
        // condition and has values.
        while (allKeyToNumValues.hasNext) {
          currentKeyToNumValue = allKeyToNumValues.next()
          if (removalCondition(currentKey)) {
            currentValues = keyWithIndexToValue.getAll(currentKey, currentKeyToNumValue.numValue)
            keyToNumValues.remove(currentKey)

            if (currentValues.hasNext) {
              return getAndRemoveValue()
            }
          }
        }

        // We only reach here if there were no satisfying keys left, which means we're done.
        finished = true
        null
      }

      override def close(): Unit = {}
    }
  }

  /**
   * Perform a full scan to provide all available data.
   *
   * This produces an iterator over the (key, value, match) tuples. Callers are expected
   * to consume fully to clean up underlying iterators correctly.
   */
  def iterator: Iterator[KeyToValuePair] = {
    new NextIterator[KeyToValuePair] {
      // Reuse this object to avoid creation+GC overhead.
      private val reusedRet = new KeyToValuePair()

      private val allKeyToNumValues = keyToNumValues.iterator

      private var currentKey: UnsafeRow = null
      private var numValues: Long = 0L
      private var index: Long = 0L

      @tailrec
      override def getNext(): KeyToValuePair = {
        if (currentKey != null) {
          assert(index < numValues)

          val valueAndMatched = keyWithIndexToValue.get(currentKey, index)
          index += 1

          reusedRet.withNew(currentKey, valueAndMatched)

          if (index == numValues) {
            currentKey = null
            numValues = 0L
            index = 0L
          }

          reusedRet
        } else if (allKeyToNumValues.hasNext) {
          val newKeyToNumValues = allKeyToNumValues.next()
          currentKey = newKeyToNumValues.key
          numValues = newKeyToNumValues.numValue
          index = 0L

          getNext()
        } else {
          finished = true
          null
        }
      }

      override protected def close(): Unit = {}
    }
  }

  override def evictByValueCondition(removalCondition: UnsafeRow => Boolean): Long = {
    var numRemoved = 0L
    evictAndReturnByValueCondition(removalCondition).foreach { _ =>
      numRemoved += 1
    }
    numRemoved
  }

  /**
   * Remove using a predicate on values.
   *
   * At a high level, this produces an iterator over the (key, value, matched) tuples such that
   * value satisfies the predicate, where producing an element removes the value from the
   * state store and producing all elements with a given key updates it accordingly.
   *
   * This implies the iterator must be consumed fully without any other operations on this manager
   * or the underlying store being interleaved.
   */
  override def evictAndReturnByValueCondition(
      removalCondition: UnsafeRow => Boolean): Iterator[KeyToValuePair] = {
    new NextIterator[KeyToValuePair] {

      // Reuse this object to avoid creation+GC overhead.
      private val reusedRet = new KeyToValuePair()

      private val allKeyToNumValues = keyToNumValues.iterator

      private var currentKey: UnsafeRow = null
      private var numValues: Long = 0L
      private var index: Long = 0L
      private var valueRemoved: Boolean = false

      // Push the data for the current key to the numValues store, and reset the tracking variables
      // to their empty state.
      private def updateNumValueForCurrentKey(): Unit = {
        if (valueRemoved) {
          if (numValues >= 1) {
            keyToNumValues.put(currentKey, numValues)
          } else {
            keyToNumValues.remove(currentKey)
          }
        }

        currentKey = null
        numValues = 0
        index = 0
        valueRemoved = false
      }

      /**
       * Find the next value satisfying the condition, updating `currentKey` and `numValues` if
       * needed. Returns null when no value can be found.
       *
       * Null values from the state store are handled by [[handleNullValuePair]]:
       * skipped if STATE_STORE_SKIP_NULLS_FOR_STREAM_STREAM_JOINS is enabled,
       * otherwise an exception is thrown with diagnostic context.
       */
      private def findNextValueForIndex(): ValueAndMatchPair = {
        // Loop across all values for the current key, and then all other keys, until we find a
        // value satisfying the removal condition.
        def hasMoreValuesForCurrentKey = currentKey != null && index < numValues
        def hasMoreKeys = allKeyToNumValues.hasNext
        while (hasMoreValuesForCurrentKey || hasMoreKeys) {
          if (hasMoreValuesForCurrentKey) {
            // First search the values for the current key.
            val valuePair = keyWithIndexToValue.get(currentKey, index)
            if (valuePair == null) {
              handleNullValuePair(currentKey, index, numValues)
              skippedNullValueCount.foreach(_ += 1L)
              index += 1
            } else if (removalCondition(valuePair.value)) {
              return valuePair
            } else {
              index += 1
            }
          } else if (hasMoreKeys) {
            // If we can't find a value for the current key, cleanup and start looking at the next.
            // This will also happen the first time the iterator is called.
            updateNumValueForCurrentKey()

            val currentKeyToNumValue = allKeyToNumValues.next()
            currentKey = currentKeyToNumValue.key
            numValues = currentKeyToNumValue.numValue
          } else {
            // Should be unreachable, but in any case means a value couldn't be found.
            return null
          }
        }

        // We tried and failed to find the next value.
        null
      }

      override def getNext(): KeyToValuePair = {
        val currentValue = findNextValueForIndex()

        // If there's no value, clean up and finish. There aren't any more available.
        if (currentValue == null) {
          updateNumValueForCurrentKey()
          finished = true
          return null
        }

        numValues = removeValueAtIndex(currentKey, index, numValues)
        valueRemoved = true

        reusedRet.withNew(currentKey, currentValue.value, currentValue.matched)
      }

      override def close(): Unit = {}
    }
  }

  // Unsafe row to internal row projection for key of `keyWithIndexToValue`.
  lazy private val keyProjection = SafeProjection.create(keySchema)

  /** Projects the key of unsafe row to internal row for printable log message. */
  def getInternalRowOfKeyWithIndex(currentKey: UnsafeRow): InternalRow = keyProjection(currentKey)

  @volatile private var nullValueWarningLogged = false

  /**
   * Called when keyWithIndexToValue.get(key, index) returns null even though
   * keyToNumValues claims there should be a value at that index.
   *
   * Logs a warning (once per instance) with diagnostic context, then either
   * returns normally (if skipNullsForStreamStreamJoins is enabled, caller skips)
   * or throws a [[StreamStreamJoinInconsistentStateNullValue]].
   */
  protected def handleNullValuePair(
      key: UnsafeRow,
      nullIndex: Long,
      numValues: Long): Unit = {
    if (!nullValueWarningLogged) {
      nullValueWarningLogged = true
      logWarning(log"Null value detected in stream-stream join state: " +
        log"joinSide=${MDC(JOIN_SIDE, joinSide)}, key=${MDC(KEY, keyProjection(key))}, " +
        log"index=${MDC(INDEX, nullIndex)}, numValues=${MDC(NUM_VALUES, numValues)}, " +
        log"storeVersion=${MDC(STATE_STORE_VERSION, stateInfo.get.storeVersion)}, " +
        log"partitionId=${MDC(PARTITION_ID, partitionId)}")
    }

    if (!storeConf.skipNullsForStreamStreamJoins) {
      throw StateStoreErrors.streamStreamJoinNullValue(
        valueIndex = nullIndex,
        numValues = numValues,
        joinSide = joinSide.toString,
        storeVersion = stateInfo.get.storeVersion,
        partitionId = partitionId,
        configKey = SQLConf.STATE_STORE_SKIP_NULLS_FOR_STREAM_STREAM_JOINS.key)
    }
  }

  /** Commit all the changes to all the state stores */
  def commit(): Unit

  /** Abort any changes to the state stores if needed */
  def abortIfNeeded(): Unit

  /**
   * Get state store checkpoint information of the two state stores for this joiner, after
   * they finished data processing.
   *
   * For [[SymmetricHashJoinStateManagerV1]], this returns the information of the two stores
   * used for this joiner.
   *
   * For [[SymmetricHashJoinStateManagerV2]], this returns the information of the single store
   * used for the entire joiner operator. Both fields of JoinerStateStoreCkptInfo will
   * be identical.
   */
  def getLatestCheckpointInfo(): JoinerStateStoreCkptInfo

  /** Get the combined metrics of all the state stores */
  def metrics: StateStoreMetrics

  /**
   * Update number of values for a key.
   * NOTE: this function is only intended for use in unit tests
   * to simulate null values.
   */
  protected[streaming] def updateNumValuesTestOnly(key: UnsafeRow, numValues: Long): Unit = {
    keyToNumValues.put(key, numValues)
  }

  /*
  =====================================================
            Private methods and inner classes
  =====================================================
   */

  // Clean up any state store resources if necessary at the end of the task
  Option(TaskContext.get()).foreach { _.addTaskCompletionListener[Unit] { _ => abortIfNeeded() } }

  /** Helper trait for invoking common functionalities of a state store. */
  protected abstract class StateStoreHandler(
      stateStoreType: StateStoreType,
      stateStoreCkptId: Option[String],
      handlerSnapshotOptions: Option[HandlerSnapshotOptions] = None) extends Logging {
    private var stateStoreProvider: StateStoreProvider = _

    /** StateStore that the subclasses of this class is going to operate on */
    protected def stateStore: StateStore

    def commit(): Unit = {
      stateStore.commit()
      logDebug("Committed, metrics = " + stateStore.metrics)
    }

    def abortIfNeeded(): Unit = {
      if (!stateStore.hasCommitted) {
        logInfo(log"Aborted store ${MDC(STATE_STORE_ID, stateStore.id)}")
        stateStore.abort()
      }
      // If this class manages a state store provider by itself, it should take care of closing
      // provider instance as well.
      if (stateStoreProvider != null) {
        stateStoreProvider.close()
      }
    }

    def metrics: StateStoreMetrics = stateStore.metrics

    def getLatestCheckpointInfo(): StateStoreCheckpointInfo = {
      stateStore.getStateStoreCheckpointInfo()
    }

    /** Get the StateStore with the given schema */
    protected def getStateStore(
        keySchema: StructType,
        valueSchema: StructType,
        useVirtualColumnFamilies: Boolean): StateStore = {
      val storeName = if (useVirtualColumnFamilies) {
        StateStoreId.DEFAULT_STORE_NAME
      } else {
        getStateStoreName(joinSide, stateStoreType)
      }
      val storeProviderId = StateStoreProviderId(stateInfo.get, partitionId, storeName)
      val store = if (useStateStoreCoordinator) {
        assert(handlerSnapshotOptions.isEmpty, "Should not use state store coordinator " +
          "when reading state as data source.")
        joinStoreGenerator.getStore(
          storeProviderId, keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema),
          stateInfo.get.storeVersion, stateStoreCkptId, None, useVirtualColumnFamilies,
          useMultipleValuesPerKey = false, storeConf, hadoopConf)
      } else {
        // This class will manage the state store provider by itself.
        stateStoreProvider = StateStoreProvider.createAndInit(
          storeProviderId, keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema),
          useColumnFamilies = useVirtualColumnFamilies, storeConf, hadoopConf,
          useMultipleValuesPerKey = false, stateSchemaProvider = None)
        if (handlerSnapshotOptions.isDefined) {
          if (!stateStoreProvider.isInstanceOf[SupportsFineGrainedReplay]) {
            throw StateStoreErrors.stateStoreProviderDoesNotSupportFineGrainedReplay(
              stateStoreProvider.getClass.toString)
          }
          val opts = handlerSnapshotOptions.get
          stateStoreProvider.asInstanceOf[SupportsFineGrainedReplay]
            .replayStateFromSnapshot(
              opts.snapshotVersion,
              opts.endVersion,
              readOnly = true,
              opts.startStateStoreCkptId,
              opts.endStateStoreCkptId)
        } else {
          stateStoreProvider.getStore(stateInfo.get.storeVersion, stateStoreCkptId)
        }
      }
      logInfo(log"Loaded store ${MDC(STATE_STORE_ID, store.id)}")
      store
    }
  }

  /**
   * Helper class for representing data returned by [[KeyWithIndexToValueStore]].
   * Designed for object reuse.
   */
  private[join] class KeyAndNumValues(var key: UnsafeRow = null, var numValue: Long = 0) {
    def withNew(newKey: UnsafeRow, newNumValues: Long): this.type = {
      this.key = newKey
      this.numValue = newNumValues
      this
    }
  }

  /** A wrapper around a [[StateStore]] that stores [key -> number of values]. */
  protected class KeyToNumValuesStore(
      val stateFormatVersion: Int,
      val handlerSnapshotOptions: Option[HandlerSnapshotOptions] = None)
    extends StateStoreHandler(
      KeyToNumValuesType, keyToNumValuesStateStoreCkptId, handlerSnapshotOptions) {

    private val useVirtualColumnFamilies = stateFormatVersion >= 3
    private val longValueSchema = new StructType().add("value", "long")
    private val longToUnsafeRow = UnsafeProjection.create(longValueSchema)
    private val valueRow = longToUnsafeRow(new SpecificInternalRow(longValueSchema))
    protected val stateStore: StateStore =
      getStateStore(keySchema, longValueSchema, useVirtualColumnFamilies)

    // Set up virtual column family name in the store if it is being used
    private val colFamilyName = if (useVirtualColumnFamilies) {
      getStateStoreName(joinSide, KeyToNumValuesType)
    } else {
      StateStore.DEFAULT_COL_FAMILY_NAME
    }

    // Create the specific column family in the store for this join side's KeyToNumValuesStore
    if (useVirtualColumnFamilies) {
      stateStore.createColFamilyIfAbsent(
        colFamilyName,
        keySchema,
        longValueSchema,
        NoPrefixKeyStateEncoderSpec(keySchema),
        isInternal = true
      )
    }

    /** Get the number of values the key has */
    def get(key: UnsafeRow): Long = {
      val longValueRow = stateStore.get(key, colFamilyName)
      if (longValueRow != null) longValueRow.getLong(0) else 0L
    }

    /** Set the number of values the key has */
    def put(key: UnsafeRow, numValues: Long): Unit = {
      require(numValues > 0)
      valueRow.setLong(0, numValues)
      stateStore.put(key, valueRow, colFamilyName)
    }

    def remove(key: UnsafeRow): Unit = {
      stateStore.remove(key, colFamilyName)
    }

    def iterator: Iterator[KeyAndNumValues] = {
      val keyAndNumValues = new KeyAndNumValues()
      stateStore.iterator(colFamilyName).map { pair =>
        keyAndNumValues.withNew(pair.key, pair.value.getLong(0))
      }
    }
  }

  /**
   * Helper class for representing data returned by [[KeyWithIndexToValueStore]].
   * Designed for object reuse.
   */
  private[join] class KeyWithIndexAndValue(
    var key: UnsafeRow = null,
    var valueIndex: Long = -1,
    var value: UnsafeRow = null,
    var matched: Boolean = false) {

    def withNew(
        newKey: UnsafeRow,
        newIndex: Long,
        newValue: UnsafeRow,
        newMatched: Boolean): this.type = {
      this.key = newKey
      this.valueIndex = newIndex
      this.value = newValue
      this.matched = newMatched
      this
    }

    def withNew(
        newKey: UnsafeRow,
        newIndex: Long,
        newValue: ValueAndMatchPair): this.type = {
      this.key = newKey
      this.valueIndex = newIndex
      if (newValue != null) {
        this.value = newValue.value
        this.matched = newValue.matched
      } else {
        this.value = null
        this.matched = false
      }
      this
    }
  }

  /**
   * A wrapper around a [[StateStore]] that stores the mapping; the mapping depends on the
   * state format version - please refer implementations of [[KeyWithIndexToValueRowConverter]].
   */
  protected class KeyWithIndexToValueStore(
      stateFormatVersion: Int,
      handlerSnapshotOptions: Option[HandlerSnapshotOptions] = None)
    extends StateStoreHandler(
      KeyWithIndexToValueType, keyWithIndexToValueStateStoreCkptId, handlerSnapshotOptions) {

    private val useVirtualColumnFamilies = stateFormatVersion >= 3
    private val keyWithIndexExprs = keyAttributes :+ Literal(1L)
    private val keyWithIndexSchema = keySchema.add("index", LongType)
    private val indexOrdinalInKeyWithIndexRow = keyAttributes.size

    // Projection to generate (key + index) row from key row
    private val keyWithIndexRowGenerator = UnsafeProjection.create(keyWithIndexExprs, keyAttributes)

    // Projection to generate key row from (key + index) row
    private val keyRowGenerator = UnsafeProjection.create(
      keyAttributes, keyAttributes :+ AttributeReference("index", LongType)())

    private val valueRowConverter = StreamingSymmetricHashJoinValueRowConverter
      .create(inputValueAttributes, stateFormatVersion)

    protected val stateStore = getStateStore(keyWithIndexSchema,
      valueRowConverter.valueAttributes.toStructType, useVirtualColumnFamilies)

    // Set up virtual column family name in the store if it is being used
    private val colFamilyName = if (useVirtualColumnFamilies) {
      getStateStoreName(joinSide, KeyWithIndexToValueType)
    } else {
      StateStore.DEFAULT_COL_FAMILY_NAME
    }

    // Create the specific column family in the store for this join side's KeyWithIndexToValueStore
    if (useVirtualColumnFamilies) {
      stateStore.createColFamilyIfAbsent(
        colFamilyName,
        keyWithIndexSchema,
        valueRowConverter.valueAttributes.toStructType,
        NoPrefixKeyStateEncoderSpec(keyWithIndexSchema)
      )
    }

    def get(key: UnsafeRow, valueIndex: Long): ValueAndMatchPair = {
      valueRowConverter.convertValue(
        stateStore.get(keyWithIndexRow(key, valueIndex), colFamilyName)
      )
    }

    /**
     * Get all values and indices for the provided key.
     * Should not return null.
     */
    def getAll(key: UnsafeRow, numValues: Long): Iterator[KeyWithIndexAndValue] = {
      new NextIterator[KeyWithIndexAndValue] {
        private val keyWithIndexAndValue = new KeyWithIndexAndValue()
        private var index: Long = 0L

        private def hasMoreValues = index < numValues
        override protected def getNext(): KeyWithIndexAndValue = {
          while (hasMoreValues) {
            val keyWithIndex = keyWithIndexRow(key, index)
            val valuePair =
              valueRowConverter.convertValue(stateStore.get(keyWithIndex, colFamilyName))
            if (valuePair == null) {
              handleNullValuePair(key, index, numValues)
              skippedNullValueCount.foreach(_ += 1L)
              index += 1
            } else {
              keyWithIndexAndValue.withNew(key, index, valuePair)
              index += 1
              return keyWithIndexAndValue
            }
          }

          finished = true
          null
        }

        override protected def close(): Unit = {}
      }
    }

    /** Put new value for key at the given index */
    def put(key: UnsafeRow, valueIndex: Long, value: UnsafeRow, matched: Boolean): Unit = {
      val keyWithIndex = keyWithIndexRow(key, valueIndex)
      val valueWithMatched = valueRowConverter.convertToValueRow(value, matched)
      stateStore.put(keyWithIndex, valueWithMatched, colFamilyName)
    }

    /**
     * Remove key and value at given index. Note that this will create a hole in
     * (key, index) and it is upto the caller to deal with it.
     */
    def remove(key: UnsafeRow, valueIndex: Long): Unit = {
      stateStore.remove(keyWithIndexRow(key, valueIndex), colFamilyName)
    }

    /** Remove all values (i.e. all the indices) for the given key. */
    def removeAllValues(key: UnsafeRow, numValues: Long): Unit = {
      var index = 0
      while (index < numValues) {
        stateStore.remove(keyWithIndexRow(key, index), colFamilyName)
        index += 1
      }
    }

    def iterator: Iterator[KeyWithIndexAndValue] = {
      val keyWithIndexAndValue = new KeyWithIndexAndValue()
      stateStore.iterator(colFamilyName).map { pair =>
        val valuePair = valueRowConverter.convertValue(pair.value)
        keyWithIndexAndValue.withNew(
          keyRowGenerator(pair.key), pair.key.getLong(indexOrdinalInKeyWithIndexRow), valuePair)
        keyWithIndexAndValue
      }
    }

    /** Generated a row using the key and index */
    private def keyWithIndexRow(key: UnsafeRow, valueIndex: Long): UnsafeRow = {
      val row = keyWithIndexRowGenerator(key)
      row.setLong(indexOrdinalInKeyWithIndexRow, valueIndex)
      row
    }
  }
}

/**
 * Streaming join state manager that uses 4 state stores without virtual column families.
 * This implementation creates a state stores based on the join side and the type of state store.
 *
 * The keyToNumValues store tracks the number of rows for each key, and the keyWithIndexToValue
 * store contains the actual entries with an additional index column.
 */
class SymmetricHashJoinStateManagerV1(
    joinSide: JoinSide,
    inputValueAttributes: Seq[Attribute],
    joinKeys: Seq[Expression],
    stateInfo: Option[StatefulOperatorStateInfo],
    storeConf: StateStoreConf,
    hadoopConf: Configuration,
    partitionId: Int,
    keyToNumValuesStateStoreCkptId: Option[String],
    keyWithIndexToValueStateStoreCkptId: Option[String],
    stateFormatVersion: Int,
    skippedNullValueCount: Option[SQLMetric] = None,
    useStateStoreCoordinator: Boolean = true,
    snapshotOptions: Option[SnapshotOptions] = None,
    joinStoreGenerator: JoinStateManagerStoreGenerator)
  extends SymmetricHashJoinStateManagerBase(
    joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
    partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
    stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
    joinStoreGenerator) {

  /** Commit all the changes to all the state stores */
  override def commit(): Unit = {
    keyToNumValues.commit()
    keyWithIndexToValue.commit()
  }

  /** Abort any changes to the state stores if needed */
  override def abortIfNeeded(): Unit = {
    keyToNumValues.abortIfNeeded()
    keyWithIndexToValue.abortIfNeeded()
  }

  /**
   * Get state store checkpoint information of the two state stores for this joiner, after
   * they finished data processing.
   *
   * For [[SymmetricHashJoinStateManagerV1]], this returns the information of the two stores
   * used for this joiner.
   */
  override def getLatestCheckpointInfo(): JoinerStateStoreCkptInfo = {
    val keyToNumValuesCkptInfo = keyToNumValues.getLatestCheckpointInfo()
    val keyWithIndexToValueCkptInfo = keyWithIndexToValue.getLatestCheckpointInfo()

    assert(
      keyToNumValuesCkptInfo.partitionId == keyWithIndexToValueCkptInfo.partitionId,
      "two state stores in a stream-stream joiner don't return the same partition ID")
    assert(
      keyToNumValuesCkptInfo.batchVersion == keyWithIndexToValueCkptInfo.batchVersion,
      "two state stores in a stream-stream joiner don't return the same batch version")
    assert(
      keyToNumValuesCkptInfo.stateStoreCkptId.isDefined ==
        keyWithIndexToValueCkptInfo.stateStoreCkptId.isDefined,
      "two state stores in a stream-stream joiner should both return checkpoint ID or not")

    JoinerStateStoreCkptInfo(keyToNumValuesCkptInfo, keyWithIndexToValueCkptInfo)
  }

  override def metrics: StateStoreMetrics = {
    val keyToNumValuesMetrics = keyToNumValues.metrics
    val keyWithIndexToValueMetrics = keyWithIndexToValue.metrics
    def newDesc(desc: String): String = s"${joinSide.toString.toUpperCase(Locale.ROOT)}: $desc"

    StateStoreMetrics(
      keyWithIndexToValueMetrics.numKeys,       // represent each buffered row only once
      keyToNumValuesMetrics.memoryUsedBytes + keyWithIndexToValueMetrics.memoryUsedBytes,
      keyWithIndexToValueMetrics.customMetrics.map {
        case (metric, value) => (metric.withNewDesc(desc = newDesc(metric.desc)), value)
      },
      // We want to collect instance metrics from both state stores
      keyWithIndexToValueMetrics.instanceMetrics ++ keyToNumValuesMetrics.instanceMetrics
    )
  }
}

/**
 * Streaming join state manager that uses 1 state store with virtual column families enabled.
 * Instead of creating a new state store per join side and store type, this manager
 * uses column families to distinguish data between the original 4 state stores.
 */
class SymmetricHashJoinStateManagerV2(
    joinSide: JoinSide,
    inputValueAttributes: Seq[Attribute],
    joinKeys: Seq[Expression],
    stateInfo: Option[StatefulOperatorStateInfo],
    storeConf: StateStoreConf,
    hadoopConf: Configuration,
    partitionId: Int,
    keyToNumValuesStateStoreCkptId: Option[String],
    keyWithIndexToValueStateStoreCkptId: Option[String],
    stateFormatVersion: Int,
    skippedNullValueCount: Option[SQLMetric] = None,
    useStateStoreCoordinator: Boolean = true,
    snapshotOptions: Option[SnapshotOptions] = None,
    joinStoreGenerator: JoinStateManagerStoreGenerator)
  extends SymmetricHashJoinStateManagerBase(
    joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
    partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
    stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
    joinStoreGenerator) {

  /** Commit all the changes to the state store */
  override def commit(): Unit = {
    // Both keyToNumValues and keyWithIndexToValue are using the same state store, so only
    // one commit is needed.
    keyToNumValues.commit()
  }

  /** Abort any changes to the state store if needed */
  override def abortIfNeeded(): Unit = {
    keyToNumValues.abortIfNeeded()
  }

  /**
   * Get state store checkpoint information of the state store used for this joiner, after
   * they finished data processing.
   *
   * For [[SymmetricHashJoinStateManagerV2]], this returns the information of the single store
   * used for the entire joiner operator. Both fields of JoinerStateStoreCkptInfo will
   * be identical.
   */
  override def getLatestCheckpointInfo(): JoinerStateStoreCkptInfo = {
    // Note that both keyToNumValues and keyWithIndexToValue are using the same state store,
    // so the latest checkpoint info should be the same.
    // These are returned in a JoinerStateStoreCkptInfo object to remain consistent with
    // the V1 implementation.
    val keyToNumValuesCkptInfo = keyToNumValues.getLatestCheckpointInfo()
    val keyWithIndexToValueCkptInfo = keyWithIndexToValue.getLatestCheckpointInfo()

    assert(keyToNumValuesCkptInfo == keyWithIndexToValueCkptInfo)

    JoinerStateStoreCkptInfo(keyToNumValuesCkptInfo, keyWithIndexToValueCkptInfo)
  }

  /** Get the state store metrics from the state store manager */
  override def metrics: StateStoreMetrics = keyToNumValues.metrics
}

/** Class used to handle state store creation in SymmetricHashJoinStateManager V1 and V2 */
class JoinStateManagerStoreGenerator() extends Logging {

  // Store internally the store used for the manager if virtual column families are enabled
  private var _store: Option[StateStore] = None

  /**
   * Creates the state store used for join operations, or returns the existing instance
   * if it has been previously created and virtual column families are enabled.
   */
  // scalastyle:off argcount
  def getStore(
      storeProviderId: StateStoreProviderId,
      keySchema: StructType,
      valueSchema: StructType,
      keyStateEncoderSpec: KeyStateEncoderSpec,
      version: Long,
      stateStoreCkptId: Option[String],
      stateSchemaBroadcast: Option[StateSchemaBroadcast],
      useColumnFamilies: Boolean,
      useMultipleValuesPerKey: Boolean,
      storeConf: StateStoreConf,
      hadoopConf: Configuration): StateStore = {
    if (useColumnFamilies) {
      // Get the store if we haven't created it yet, otherwise use the one we just created
      if (_store.isEmpty) {
        _store = Some(
          StateStore.get(
            storeProviderId, keySchema, valueSchema, keyStateEncoderSpec, version,
            stateStoreCkptId, stateSchemaBroadcast, useColumnFamilies = useColumnFamilies,
            storeConf, hadoopConf, useMultipleValuesPerKey = useMultipleValuesPerKey
          )
        )
      }
      _store.get
    } else {
      // Do not use the store saved internally, as we need to create the four distinct stores
      StateStore.get(
        storeProviderId, keySchema, valueSchema, keyStateEncoderSpec, version,
        stateStoreCkptId, stateSchemaBroadcast, useColumnFamilies = useColumnFamilies,
        storeConf, hadoopConf, useMultipleValuesPerKey = useMultipleValuesPerKey
      )
    }
  }
  // scalastyle:on
}

object SymmetricHashJoinStateManager {
  val supportedVersions = Seq(1, 2, 3, 4)
  val legacyVersion = 1

  // scalastyle:off argcount
  /** Factory method to determines which version of the join state manager should be created */
  def apply(
      joinSide: JoinSide,
      inputValueAttributes: Seq[Attribute],
      joinKeys: Seq[Expression],
      stateInfo: Option[StatefulOperatorStateInfo],
      storeConf: StateStoreConf,
      hadoopConf: Configuration,
      partitionId: Int,
      keyToNumValuesStateStoreCkptId: Option[String],
      keyWithIndexToValueStateStoreCkptId: Option[String],
      stateFormatVersion: Int,
      skippedNullValueCount: Option[SQLMetric] = None,
      useStateStoreCoordinator: Boolean = true,
      snapshotOptions: Option[SnapshotOptions] = None,
      joinStoreGenerator: JoinStateManagerStoreGenerator,
      joinKeyOrdinalForWatermark: Option[Int] = None): SymmetricHashJoinStateManager = {
    if (stateFormatVersion == 4) {
      require(SQLConf.get.getConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_V4_ENABLED),
        "State format version 4 is under development.")
      new SymmetricHashJoinStateManagerV4(
        joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
        partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
        joinStoreGenerator, joinKeyOrdinalForWatermark
      )
    } else if (stateFormatVersion == 3) {
      new SymmetricHashJoinStateManagerV2(
        joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
        partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
        joinStoreGenerator
      )
    } else {
      new SymmetricHashJoinStateManagerV1(
        joinSide, inputValueAttributes, joinKeys, stateInfo, storeConf, hadoopConf,
        partitionId, keyToNumValuesStateStoreCkptId, keyWithIndexToValueStateStoreCkptId,
        stateFormatVersion, skippedNullValueCount, useStateStoreCoordinator, snapshotOptions,
        joinStoreGenerator
      )
    }
  }
  // scalastyle:on

  def allStateStoreNames(joinSides: JoinSide*): Seq[String] = {
    val allStateStoreTypes: Seq[StateStoreType] = Seq(KeyToNumValuesType, KeyWithIndexToValueType)
    for (joinSide <- joinSides; stateStoreType <- allStateStoreTypes) yield {
      getStateStoreName(joinSide, stateStoreType)
    }
  }

  def allStateStoreNamesV4(joinSides: JoinSide*): Seq[String] = {
    val allStateStoreTypes: Seq[StateStoreType] = Seq(KeyWithTsToValuesType, TsWithKeyType)
    for (joinSide <- joinSides; stateStoreType <- allStateStoreTypes) yield {
      getStateStoreName(joinSide, stateStoreType)
    }
  }

  def getSchemaForStateStores(
      joinSide: JoinSide,
      inputValueAttributes: Seq[Attribute],
      joinKeys: Seq[Expression],
      stateFormatVersion: Int): Map[String, (StructType, StructType)] = {
    val keySchema = StructType(
      joinKeys.zipWithIndex.map { case (k, i) => StructField(s"field$i", k.dataType, k.nullable) })

    if (stateFormatVersion == 4) {
      // V4 uses two column families: KeyWithTsToValues and TsWithKey
      val keySchemaWithTimestamp =
        TimestampKeyStateEncoder.keySchemaWithTimestamp(keySchema)
      val valueWithMatchedSchema =
        (inputValueAttributes :+ AttributeReference("matched", BooleanType)()).toStructType
      val dummyValueSchema = StructType(Array(StructField("__dummy__", NullType)))

      Map(
        getStateStoreName(joinSide, KeyWithTsToValuesType) ->
          (keySchemaWithTimestamp, valueWithMatchedSchema),
        getStateStoreName(joinSide, TsWithKeyType) ->
          (keySchemaWithTimestamp, dummyValueSchema))
    } else {
      var result: Map[String, (StructType, StructType)] = Map.empty

      // get the key and value schema for the KeyToNumValues state store
      val longValueSchema = new StructType().add("value", "long")
      result += (getStateStoreName(joinSide, KeyToNumValuesType) -> (keySchema, longValueSchema))

      // get the key and value schema for the KeyWithIndexToValue state store
      val keyWithIndexSchema = keySchema.add("index", LongType)
      val valueSchema = if (stateFormatVersion == 1) {
        inputValueAttributes
      } else if (stateFormatVersion == 2 || stateFormatVersion == 3) {
        inputValueAttributes :+ AttributeReference("matched", BooleanType)()
      } else {
        throw new IllegalArgumentException("Incorrect state format version! " +
          s"version=$stateFormatVersion")
      }
      result += (getStateStoreName(joinSide, KeyWithIndexToValueType) ->
        (keyWithIndexSchema, valueSchema.toStructType))

      result
    }
  }

  /** Retrieves the schemas used for join operator state stores that use column families */
  def getSchemasForStateStoreWithColFamily(
      joinSide: JoinSide,
      inputValueAttributes: Seq[Attribute],
      joinKeys: Seq[Expression],
      stateFormatVersion: Int): Map[String, StateStoreColFamilySchema] = {
    val schemas =
      getSchemaForStateStores(joinSide, inputValueAttributes, joinKeys, stateFormatVersion)

    schemas.map {
      case (colFamilyName, (keySchema, valueSchema)) =>
        val keyStateEncoderSpec = if (stateFormatVersion == 4) {
          if (colFamilyName == getStateStoreName(joinSide, KeyWithTsToValuesType)) {
            TimestampAsPostfixKeyStateEncoderSpec(keySchema)
          } else {
            TimestampAsPrefixKeyStateEncoderSpec(keySchema)
          }
        } else {
          NoPrefixKeyStateEncoderSpec(keySchema)
        }
        colFamilyName -> StateStoreColFamilySchema(
          colFamilyName, 0, keySchema, 0, valueSchema,
          Some(keyStateEncoderSpec)
        )
    }
  }

  /**
   * Stream-stream join has 4 state stores instead of one. So it will generate 4 different
   * checkpoint IDs. The approach we take here is to merge them into one array in the checkpointing
   * path. The driver will process this single checkpointID. When it is passed back to the
   * executors, they will split it back into 4 IDs and use them to load the state. This function is
   * used to merge two checkpoint IDs (each in the form of an array of 1) into one array.
   * The merged array is expected to read back by `getStateStoreCheckpointIds()`.
   */
  def mergeStateStoreCheckpointInfo(joinCkptInfo: JoinStateStoreCkptInfo):
      StatefulOpStateStoreCheckpointInfo = {
    assert(
      joinCkptInfo.left.keyToNumValues.partitionId == joinCkptInfo.right.keyToNumValues.partitionId,
      "state store info returned from two Stream-Stream Join sides have different partition IDs")
    assert(
      joinCkptInfo.left.keyToNumValues.batchVersion ==
      joinCkptInfo.right.keyToNumValues.batchVersion,
      "state store info returned from two Stream-Stream Join sides have different batch versions")
    assert(
      joinCkptInfo.left.keyToNumValues.stateStoreCkptId.isDefined ==
      joinCkptInfo.right.keyToNumValues.stateStoreCkptId.isDefined,
      "state store info returned from two Stream-Stream Join sides should both return " +
        "checkpoint ID or not")

    val ckptIds = joinCkptInfo.left.keyToNumValues.stateStoreCkptId.map(
      Array(
        _,
        joinCkptInfo.left.keyWithIndexToValue.stateStoreCkptId.get,
        joinCkptInfo.right.keyToNumValues.stateStoreCkptId.get,
        joinCkptInfo.right.keyWithIndexToValue.stateStoreCkptId.get
      )
    )
    val baseCkptIds = joinCkptInfo.left.keyToNumValues.baseStateStoreCkptId.map(
      Array(
        _,
        joinCkptInfo.left.keyWithIndexToValue.baseStateStoreCkptId.get,
        joinCkptInfo.right.keyToNumValues.baseStateStoreCkptId.get,
        joinCkptInfo.right.keyWithIndexToValue.baseStateStoreCkptId.get
      )
    )

    StatefulOpStateStoreCheckpointInfo(
      joinCkptInfo.left.keyToNumValues.partitionId,
      joinCkptInfo.left.keyToNumValues.batchVersion,
      ckptIds,
      baseCkptIds)
  }

  /**
   * Stream-stream join has 4 state stores instead of one. So it will generate 4 different
   * checkpoint IDs using stateStoreCkptIds. They are translated from each joiners' state
   * store into an array through mergeStateStoreCheckpointInfo(). This function is used to read
   * it back into individual state store checkpoint IDs for each store.
   * If useColumnFamiliesForJoins is true, then it will always return the first checkpoint ID.
   *
   * @param partitionId the partition ID of the state store
   * @param stateStoreCkptIds the array of checkpoint IDs for all the state stores
   * @param useColumnFamiliesForJoins whether virtual column families are used for the join
   *
   * @return the checkpoint IDs for all state stores used by this joiner
   */
  def getStateStoreCheckpointIds(
      partitionId: Int,
      stateStoreCkptIds: Option[Array[Array[String]]],
      useColumnFamiliesForJoins: Boolean): JoinStateStoreCheckpointId = {
    if (useColumnFamiliesForJoins) {
      val ckpt = stateStoreCkptIds.map(_(partitionId)).map(_.head)
      JoinStateStoreCheckpointId(
        left = JoinerStateStoreCheckpointId(keyToNumValues = ckpt, keyWithIndexToValue = ckpt),
        right = JoinerStateStoreCheckpointId(keyToNumValues = ckpt, keyWithIndexToValue = ckpt)
      )
    } else {
      val stateStoreCkptIdsOpt = stateStoreCkptIds
        .map(_(partitionId))
        .map(_.map(Option(_)))
        .getOrElse(Array.fill[Option[String]](4)(None))
      JoinStateStoreCheckpointId(
        left = JoinerStateStoreCheckpointId(
          keyToNumValues = stateStoreCkptIdsOpt(0),
          keyWithIndexToValue = stateStoreCkptIdsOpt(1)),
        right = JoinerStateStoreCheckpointId(
          keyToNumValues = stateStoreCkptIdsOpt(2),
          keyWithIndexToValue = stateStoreCkptIdsOpt(3)))
    }
  }

  /**
   * Stream-stream join has 4 state stores instead of one. So it will generate 4 different
   * checkpoint IDs when not using virtual column families.
   * This function is used to get the checkpoint ID for a specific state store
   * by the name of the store, partition ID and the stateStoreCkptIds array. The expected names
   * for the stores are generated by getStateStoreName().
   * If useColumnFamiliesForJoins is true, then it will always return the first checkpoint ID.
   *
   * @param storeName the name of the state store
   * @param partitionId the partition ID of the state store
   * @param stateStoreCkptIds the array of checkpoint IDs for all the state stores
   * @param useColumnFamiliesForJoins whether virtual column families are used for the join
   *
   * @return the checkpoint ID for the specific state store, or None if not found
   */
  def getStateStoreCheckpointId(
      storeName: String,
      partitionId: Int,
      stateStoreCkptIds: Option[Array[Array[String]]],
      useColumnFamiliesForJoins: Boolean = false) : Option[String] = {
    if (useColumnFamiliesForJoins || storeName == StateStoreId.DEFAULT_STORE_NAME) {
      stateStoreCkptIds.map(_(partitionId)).map(_.head)
    } else {
      val joinStateStoreCkptIds = getStateStoreCheckpointIds(
        partitionId, stateStoreCkptIds, useColumnFamiliesForJoins)

      if (storeName == getStateStoreName(LeftSide, KeyToNumValuesType)) {
        joinStateStoreCkptIds.left.keyToNumValues
      } else if (storeName == getStateStoreName(RightSide, KeyToNumValuesType)) {
        joinStateStoreCkptIds.right.keyToNumValues
      } else if (storeName == getStateStoreName(LeftSide, KeyWithIndexToValueType)) {
        joinStateStoreCkptIds.left.keyWithIndexToValue
      } else if (storeName == getStateStoreName(RightSide, KeyWithIndexToValueType)) {
        joinStateStoreCkptIds.right.keyWithIndexToValue
      } else {
        None
      }
    }
  }

  private[sql] sealed trait StateStoreType

  private[sql] case object KeyToNumValuesType extends StateStoreType {
    override def toString(): String = "keyToNumValues"
  }

  private[sql] case object KeyWithIndexToValueType extends StateStoreType {
    override def toString(): String = "keyWithIndexToValue"
  }

  private[sql] case object KeyWithTsToValuesType extends StateStoreType {
    override def toString(): String = "keyWithTsToValues"
  }

  private[sql] case object TsWithKeyType extends StateStoreType {
    override def toString(): String = "tsWithKey"
  }

  private[join] def getStateStoreName(
      joinSide: JoinSide, storeType: StateStoreType): String = {
    s"$joinSide-$storeType"
  }

  private[join] def getStoreType(storeName: String): StateStoreType = {
    if (storeName == getStateStoreName(LeftSide, KeyToNumValuesType) ||
      storeName == getStateStoreName(RightSide, KeyToNumValuesType)) {
      KeyToNumValuesType
    } else if (storeName == getStateStoreName(LeftSide, KeyWithIndexToValueType) ||
      storeName == getStateStoreName(RightSide, KeyWithIndexToValueType)) {
      KeyWithIndexToValueType
    } else if (storeName == getStateStoreName(LeftSide, KeyWithTsToValuesType) ||
      storeName == getStateStoreName(RightSide, KeyWithTsToValuesType)) {
      KeyWithTsToValuesType
    } else if (storeName == getStateStoreName(LeftSide, TsWithKeyType) ||
      storeName == getStateStoreName(RightSide, TsWithKeyType)) {
      TsWithKeyType
    } else {
      throw new IllegalArgumentException(s"Unsupported join store name: $storeName")
    }
  }

  /**
   * Returns the partition key extractor for the given join store and column family name.
   */
  def createPartitionKeyExtractor(
      storeName: String,
      colFamilyName: String,
      stateKeySchema: StructType,
      stateFormatVersion: Int): StatePartitionKeyExtractor = {
    assert(stateFormatVersion <= 4, "State format version must be less than or equal to 4")
    val name = if (stateFormatVersion >= 3) colFamilyName else storeName
    getStoreType(name) match {
      case KeyWithIndexToValueType =>
        // For KeyWithIndex, the index is added to the join (i.e. partition) key.
        // Drop the last field (index) to get the partition key
        new DropLastNFieldsStatePartitionKeyExtractor(stateKeySchema, numLastColsToDrop = 1)
      case KeyToNumValuesType =>
        new NoopStatePartitionKeyExtractor(stateKeySchema)
      case KeyWithTsToValuesType | TsWithKeyType =>
        // For v4 stores, the logical key schema in the schema file is just the join key
        // (timestamp is managed by the encoder), so the state key IS the partition key.
        new NoopStatePartitionKeyExtractor(stateKeySchema)
    }
  }

  /** Helper class for representing data (value, matched). */
  case class ValueAndMatchPair(value: UnsafeRow, matched: Boolean)

  /**
   * Helper class for representing data key to (value, matched).
   * Designed for object reuse.
   */
  case class KeyToValuePair(
      var key: UnsafeRow = null,
      var value: UnsafeRow = null,
      var matched: Boolean = false) {
    def withNew(newKey: UnsafeRow, newValue: UnsafeRow, newMatched: Boolean): this.type = {
      this.key = newKey
      this.value = newValue
      this.matched = newMatched
      this
    }

    def withNew(newKey: UnsafeRow, newValue: ValueAndMatchPair): this.type = {
      this.key = newKey
      if (newValue != null) {
        this.value = newValue.value
        this.matched = newValue.matched
      } else {
        this.value = null
        this.matched = false
      }
      this
    }
  }

  /**
   * Helper class for representing data (key, timestamp) to (value, matched).
   * Designed for object reuse.
   */
  case class KeyAndTsToValuePair(
      var key: UnsafeRow = null,
      var timestamp: Long = -1L,
      var value: UnsafeRow = null,
      var matched: Boolean = false) {
    def withNew(
        newKey: UnsafeRow,
        newTimestamp: Long,
        newValue: UnsafeRow,
        newMatched: Boolean): this.type = {
      this.key = newKey
      this.timestamp = newTimestamp
      this.value = newValue
      this.matched = newMatched
      this
    }

    def withNew(
        newKey: UnsafeRow,
        newTimestamp: Long,
        newValue: ValueAndMatchPair): this.type = {
      this.key = newKey
      this.timestamp = newTimestamp
      this.value = newValue.value
      this.matched = newValue.matched
      this
    }
  }
}

/**
 * Options controlling snapshot-based state replay for state data source reader.
 */
case class SnapshotOptions(
    snapshotVersion: Long,
    endVersion: Long,
    startKeyToNumValuesStateStoreCkptId: Option[String] = None,
    startKeyWithIndexToValueStateStoreCkptId: Option[String] = None,
    endKeyToNumValuesStateStoreCkptId: Option[String] = None,
    endKeyWithIndexToValueStateStoreCkptId: Option[String] = None) {

  def getKeyToNumValuesHandlerOpts(): HandlerSnapshotOptions =
    HandlerSnapshotOptions(
      snapshotVersion = snapshotVersion,
      endVersion = endVersion,
      startStateStoreCkptId = startKeyToNumValuesStateStoreCkptId,
      endStateStoreCkptId = endKeyToNumValuesStateStoreCkptId)

  def getKeyWithIndexToValueHandlerOpts(): HandlerSnapshotOptions =
    HandlerSnapshotOptions(
      snapshotVersion = snapshotVersion,
      endVersion = endVersion,
      startStateStoreCkptId = startKeyWithIndexToValueStateStoreCkptId,
      endStateStoreCkptId = endKeyWithIndexToValueStateStoreCkptId)
}

/** Snapshot options specialized for a single state store handler. */
private[join] case class HandlerSnapshotOptions(
    snapshotVersion: Long,
    endVersion: Long,
    startStateStoreCkptId: Option[String],
    endStateStoreCkptId: Option[String])
