/*
 * 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

import java.util.UUID
import java.util.concurrent.TimeUnit._

import scala.collection.mutable
import scala.jdk.CollectionConverters._

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.Path

import org.apache.spark.SparkContext
import org.apache.spark.internal.Logging
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{AnalysisException, SparkSession}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection
import org.apache.spark.sql.catalyst.optimizer.NormalizeFloatingNumbers
import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark
import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, Distribution, Partitioning}
import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import org.apache.spark.sql.execution.python.PythonSQLMetrics
import org.apache.spark.sql.execution.streaming.state._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.{OutputMode, StateOperatorProgress}
import org.apache.spark.sql.types._
import org.apache.spark.util.{CollectionAccumulator, CompletionIterator, NextIterator, Utils}


/** Used to identify the state store for a given operator.
 *
 * stateStoreCkptIds is used to identify the checkpoint used for a specific stateful operator
 * The basic workflow works as following:
 * 1. When a stateful operator is created, it passes in the checkpoint IDs for each stateful
 *    operator through the StatefulOperatorStateInfo.
 * 2. When a stateful task starts to execute, it will find the checkpointID for its shuffle
 *    partition and use it to recover the state store. The ID is eventually passed into
 *    the StateStore layer and eventually  RocksDB State Store, where it is used to make sure
 *    the it loads the correct checkpoint
 * 3. When the stateful task is finishing, after the state store is committed, the checkpoint ID
 *    is fetched from the state store by calling StateStore.getStateStoreCheckpointInfo() and added
 *    to the stateStoreCkptIds accumulator by calling
 *    StateStoreWriter.setStateStoreCheckpointInfo().
 * 4. When ending the batch, MicroBatchExecution calls each stateful operator's
 *    getStateStoreCheckpointInfo() which aggregates checkpointIDs from different partitions. The
 *    driver will persistent it into commit logs (not implemented yet).
 * 5. When forming the next batch, the driver constructs the StatefulOperatorStateInfo with the
 *    checkpoint IDs for the previous batch.
 * */
case class StatefulOperatorStateInfo(
    checkpointLocation: String,
    queryRunId: UUID,
    operatorId: Long,
    storeVersion: Long,
    numPartitions: Int,
    stateSchemaMetadata: Option[StateSchemaBroadcast] = None,
    stateStoreCkptIds: Option[Array[Array[String]]] = None) {

  def getStateStoreCkptId(partitionId: Int): Option[Array[String]] = {
    stateStoreCkptIds.map(_(partitionId))
  }

  override def toString(): String = {
    s"state info [ checkpoint = $checkpointLocation, runId = $queryRunId, " +
      s"opId = $operatorId, ver = $storeVersion, numPartitions = $numPartitions] " +
      s"stateStoreCkptIds = $stateStoreCkptIds"
  }
}

object StatefulOperatorStateInfo {
  /**
   * Whether stateo store checkpoint version requires checkpointID to be used.
   * @return true if state store checkpointID should be used.
   */
  def enableStateStoreCheckpointIds(conf: SQLConf): Boolean = {
    conf.stateStoreCheckpointFormatVersion >= 2
  }
}

/**
 * An operator that reads or writes state from the [[StateStore]].
 * The [[StatefulOperatorStateInfo]] should be filled in by `prepareForExecution` in
 * [[IncrementalExecution]].
 */
trait StatefulOperator extends SparkPlan {
  def stateInfo: Option[StatefulOperatorStateInfo]

  def getStateInfo: StatefulOperatorStateInfo = {
    stateInfo.getOrElse {
      throw new IllegalStateException("State location not present for execution")
    }
  }

  def metadataFilePath(): Path = {
    val stateCheckpointPath =
      new Path(getStateInfo.checkpointLocation, getStateInfo.operatorId.toString)
    new Path(new Path(stateCheckpointPath, "_metadata"), "metadata")
  }

  // Function used to record state schema for the first time and validate it against proposed
  // schema changes in the future. Runs as part of a planning rule on the driver.
  // Returns the schema file path for operators that write this to the metadata file,
  // otherwise None
  def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
    List[StateSchemaValidationResult]
}

/**
 * Custom stateful operator metric definition to allow operators to expose their own custom metrics.
 * Also provides [[SQLMetric]] instance to show the metric in UI and accumulate it at the query
 * level.
 */
trait StatefulOperatorCustomMetric {
  def name: String
  def desc: String
  def createSQLMetric(sparkContext: SparkContext): SQLMetric
}

/** Custom stateful operator metric for simple "count" gauge */
case class StatefulOperatorCustomSumMetric(name: String, desc: String)
  extends StatefulOperatorCustomMetric {
  override def createSQLMetric(sparkContext: SparkContext): SQLMetric =
    SQLMetrics.createMetric(sparkContext, desc)
}

/** An operator that reads from a StateStore. */
trait StateStoreReader extends StatefulOperator with Logging {
  override lazy val metrics = Map(
    "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"))
}

/**
 * Used to pass state store checkpoint information to load the correct state store checkpoint for a
 * stateful operator. Will be passed from the driver to exeuctors to load state store.
 */
case class StatefulOpStateStoreCheckpointInfo(
    partitionId: Int,
    batchVersion: Long,
    // The checkpoint ID for a checkpoint at `batchVersion`. This is used to identify the checkpoint
    stateStoreCkptId: Option[Array[String]],
    // The checkpoint ID for `batchVersion` - 1, that is used to finish this batch. This is used
    // to validate the batch is processed based on the correct checkpoint.
    baseStateStoreCkptId: Option[Array[String]])

/** An operator that writes to a StateStore. */
trait StateStoreWriter
  extends StatefulOperator with PythonSQLMetrics with Logging { self: SparkPlan =>

  /**
   * Produce the output watermark for given input watermark (ms).
   *
   * In most cases, this is same as the criteria of state eviction, as most stateful operators
   * produce the output from two different kinds:
   *
   * 1. without buffering
   * 2. with buffering (state)
   *
   * The state eviction happens when event time exceeds a "certain threshold of timestamp", which
   * denotes a lower bound of event time values for output (output watermark).
   *
   * The default implementation provides the input watermark as it is. Most built-in operators
   * will evict based on min input watermark and ensure it will be minimum of the event time value
   * for the output so far (including output from eviction). Operators which behave differently
   * (e.g. different criteria on eviction) must override this method.
   *
   * Note that the default behavior wil advance the watermark aggressively to simplify the logic,
   * but it does not break the semantic of output watermark, which is following:
   *
   * An operator guarantees that it will not emit record with an event timestamp lower than its
   * output watermark.
   *
   * For example, for 5 minutes time window aggregation, the advancement of watermark can happen
   * "before" the window has been evicted and produced as output. Say, suppose there's an window
   * in state: [0, 5) and input watermark = 3. Although there is no output for this operator, this
   * operator will produce an output watermark as 3. It's still respecting the guarantee, as the
   * operator will produce the window [0, 5) only when the output watermark is equal or greater
   * than 5, and the downstream operator will process the input data, "and then" advance the
   * watermark. Hence this window is considered as "non-late" record.
   */
  def produceOutputWatermark(inputWatermarkMs: Long): Option[Long] = Some(inputWatermarkMs)

  def operatorStateMetadataVersion: Int = 1

  override lazy val metrics = {
    // Lazy initialize instance metrics, but do not include these with regular metrics
    instanceMetrics
    statefulOperatorCustomMetrics ++ Map(
      "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
      "numRowsDroppedByWatermark" -> SQLMetrics
        .createMetric(sparkContext, "number of rows which are dropped by watermark"),
      "numTotalStateRows" -> SQLMetrics.createMetric(sparkContext, "number of total state rows"),
      "numUpdatedStateRows" -> SQLMetrics
        .createMetric(sparkContext, "number of updated state rows"),
      "allUpdatesTimeMs" -> SQLMetrics.createTimingMetric(sparkContext, "time to update"),
      "numRemovedStateRows" -> SQLMetrics
        .createMetric(sparkContext, "number of removed state rows"),
      "allRemovalsTimeMs" -> SQLMetrics.createTimingMetric(sparkContext, "time to remove"),
      "commitTimeMs" -> SQLMetrics.createTimingMetric(sparkContext, "time to commit changes"),
      "stateMemory" -> SQLMetrics.createSizeMetric(sparkContext, "memory used by state"),
      "numStateStoreInstances" -> SQLMetrics
        .createMetric(sparkContext, "number of state store instances")
    ) ++ stateStoreCustomMetrics ++ pythonMetrics
  }

  /**
   * Map of all instance metrics (including partition ID and store names) to
   * their SQLMetric counterpart.
   *
   * The instance metric objects hold additional information on how to report these metrics,
   * while the SQLMetric objects store the metric values.
   *
   * This map is similar to the metrics map, but needs to be kept separate to prevent propagating
   * all initialized instance metrics to SparkUI.
   */
  lazy val instanceMetrics: Map[StateStoreInstanceMetric, SQLMetric] =
    stateStoreInstanceMetrics

  override def resetMetrics(): Unit = {
    super.resetMetrics()
    instanceMetrics.valuesIterator.foreach(_.reset())
  }

  val stateStoreNames: Seq[String] = Seq(StateStoreId.DEFAULT_STORE_NAME)

  // This method is only used to fetch the state schema directory path for
  // operators that use StateSchemaV3, as prior versions only use a single
  // set file path.
  def stateSchemaDirPath(
      storeName: Option[String] = None): Path = {
    val stateCheckpointPath =
      new Path(getStateInfo.checkpointLocation,
        s"${getStateInfo.operatorId.toString}")
    storeName match {
      case Some(storeName) =>
        new Path(new Path(stateCheckpointPath, "_stateSchema"), storeName)
      case None =>
        new Path(new Path(stateCheckpointPath, "_stateSchema"), "default")
    }
  }

  def stateSchemaList(
      stateSchemaValidationResults: List[StateSchemaValidationResult],
      oldMetadata: Option[OperatorStateMetadata]): List[List[String]] = {

    def getExistingSchemaFilePaths(metadata: OperatorStateMetadataV2): List[String] = {
      val ssInfo = metadata.stateStoreInfo.head
      ssInfo.stateSchemaFilePaths
    }

    // TODO: Make this work for Stream-Stream joins, where we use multiple
    // StateStores
    val validationResult = stateSchemaValidationResults.head

    oldMetadata match {
      case Some(v2: OperatorStateMetadataV2) =>
        val oldSchemaPaths = getExistingSchemaFilePaths(v2)
        if (validationResult.evolvedSchema) {
          List(oldSchemaPaths ++ List(validationResult.schemaPath))
        } else {
          List(oldSchemaPaths)
        }
      case _ =>
        // No previous metadata - start with schema ID 0
        List(List(validationResult.schemaPath))
    }
  }

  def supportsSchemaEvolution: Boolean = false

  /**
   * Aggregator used for the executors to pass new state store checkpoints' IDs to driver.
   * For the general checkpoint ID workflow, see comments of
   * class class [[StatefulOperatorStateInfo]]
   */
  val checkpointInfoAccumulator: CollectionAccumulator[StatefulOpStateStoreCheckpointInfo] = {
    SparkContext.getActive.map(_.collectionAccumulator[StatefulOpStateStoreCheckpointInfo]).get
  }

  /**
   * Get aggregated checkpoint ID info for all shuffle partitions
   * For the general checkpoint ID workflow, see comments of
   * class class [[StatefulOperatorStateInfo]]
   */
  def getStateStoreCheckpointInfo(): Array[StatefulOpStateStoreCheckpointInfo] = {
    assert(
      StatefulOperatorStateInfo.enableStateStoreCheckpointIds(conf),
      "Should not fetch checkpoint Info if the state store checkpoint IDs are not enabled")
    // Multiple entries can be returned for the same partitionID due to task speculative execution
    // or task failures. All of them should represent a valid state store checkpoint and we just
    // pick one of them.
    // In the end, we sort them by partitionID.
    val ret = checkpointInfoAccumulator
      .value
      .asScala
      .toSeq
      .groupBy(_.partitionId)
      .map {
        case (key, values) => key -> values.head
      }
      .toSeq
      .sortBy(_._1)
      .map(_._2)
      .toArray
    assert(
      // Normally, we should have checkpoint info for all partitions.
      // However, for globalLimit operator, there is only one partition (0) that has state.
      ret.length == getStateInfo.numPartitions
        || (outputPartitioning.numPartitions == 1 && ret.length == 1),
      s"CheckpointInfo length: ${ret.length}, numPartitions: ${getStateInfo.numPartitions}")
    ret
  }

  /**
   * The executor reports its state store checkpoint ID, which would be sent back to the driver.
   * For the general checkpoint ID workflow, see comments of
   * class class [[StatefulOperatorStateInfo]]
   */
  protected def setStateStoreCheckpointInfo(
      checkpointInfo: StatefulOpStateStoreCheckpointInfo): Unit = {
    if (StatefulOperatorStateInfo.enableStateStoreCheckpointIds(conf)) {
      checkpointInfoAccumulator.add(checkpointInfo)
    }
  }

  /**
   * Get the progress made by this stateful operator after execution. This should be called in
   * the driver after this SparkPlan has been executed and metrics have been updated.
   */
  def getProgress(): StateOperatorProgress = {
    val instanceMetricsToReport = instanceMetrics
      .filter {
        case (metricConf, sqlMetric) =>
          // Keep instance metrics that are updated or aren't marked to be ignored,
          // as their initial value could still be important.
          !metricConf.ignoreIfUnchanged || !sqlMetric.isZero
      }
      .groupBy {
        // Group all instance metrics underneath their common metric prefix
        // to ignore partition and store names.
        case (metricConf, sqlMetric) => metricConf.metricPrefix
      }
      .flatMap {
        case (_, metrics) =>
          // Select at most N metrics based on the metric's defined ordering
          // to report to the driver. For example, ascending order would be taking the N smallest.
          val metricConf = metrics.head._1
          metrics
            .map {
              case (metricConf, sqlMetric) =>
                // Use metric name as it will be combined with custom metrics in progress reports.
                // All metrics that are at their initial value at this stage should not be ignored
                // and should show their real initial value.
                metricConf.name -> (if (sqlMetric.isZero) metricConf.initValue
                                    else sqlMetric.value)
            }
            .toSeq
            .sortBy(_._2)(metricConf.ordering)
            .take(conf.numStateStoreInstanceMetricsToReport)
            .toMap
      }
    val customMetrics = (stateStoreCustomMetrics ++ statefulOperatorCustomMetrics)
      .map(entry => entry._1 -> longMetric(entry._1).value)
    val allCustomMetrics = customMetrics ++ instanceMetricsToReport

    val javaConvertedCustomMetrics: java.util.HashMap[String, java.lang.Long] =
      new java.util.HashMap(allCustomMetrics.transform((_, v) => long2Long(v)).asJava)

    // We now don't report number of shuffle partitions inside the state operator. Instead,
    // it will be filled when the stream query progress is reported
    new StateOperatorProgress(
      operatorName = shortName,
      numRowsTotal = longMetric("numTotalStateRows").value,
      numRowsUpdated = longMetric("numUpdatedStateRows").value,
      allUpdatesTimeMs = longMetric("allUpdatesTimeMs").value,
      numRowsRemoved = longMetric("numRemovedStateRows").value,
      allRemovalsTimeMs = longMetric("allRemovalsTimeMs").value,
      commitTimeMs = longMetric("commitTimeMs").value,
      memoryUsedBytes = longMetric("stateMemory").value,
      numRowsDroppedByWatermark = longMetric("numRowsDroppedByWatermark").value,
      numShufflePartitions = stateInfo.map(_.numPartitions.toLong).getOrElse(-1L),
      numStateStoreInstances = longMetric("numStateStoreInstances").value,
      javaConvertedCustomMetrics
    )
  }

  /** Records the duration of running `body` for the next query progress update. */
  protected def timeTakenMs(body: => Unit): Long = Utils.timeTakenMs(body)._2

  /** Metadata of this stateful operator and its states stores. */
  def operatorStateMetadata(
      stateSchemaPaths: List[List[String]] = List.empty): OperatorStateMetadata = {
    val info = getStateInfo
    val operatorInfo = OperatorInfoV1(info.operatorId, shortName)
    val stateStoreInfo =
      Array(StateStoreMetadataV1(StateStoreId.DEFAULT_STORE_NAME, 0, info.numPartitions))
    OperatorStateMetadataV1(operatorInfo, stateStoreInfo)
  }

  /** Set the operator level metrics */
  protected def setOperatorMetrics(numStateStoreInstances: Int = 1): Unit = {
    assert(numStateStoreInstances >= 1, s"invalid number of stores: $numStateStoreInstances")
    // Shuffle partitions capture the number of tasks that have this stateful operator instance.
    // For each task instance this number is incremented by one.
    longMetric("numStateStoreInstances") += numStateStoreInstances
  }

  /**
   * Set the SQL metrics related to the state store.
   * This should be called in that task after the store has been updated.
   */
  protected def setStoreMetrics(store: StateStore, setCheckpointInfo: Boolean = true): Unit = {
    val storeMetrics = store.metrics
    longMetric("numTotalStateRows") += storeMetrics.numKeys
    longMetric("stateMemory") += storeMetrics.memoryUsedBytes
    setStoreCustomMetrics(storeMetrics.customMetrics)
    setStoreInstanceMetrics(storeMetrics.instanceMetrics)

    if (StatefulOperatorStateInfo.enableStateStoreCheckpointIds(conf) && setCheckpointInfo) {
      // Set the state store checkpoint information for the driver to collect
      val ssInfo = store.getStateStoreCheckpointInfo()
      setStateStoreCheckpointInfo(
        StatefulOpStateStoreCheckpointInfo(
          ssInfo.partitionId,
          ssInfo.batchVersion,
          ssInfo.stateStoreCkptId.map(Array(_)),
          ssInfo.baseStateStoreCkptId.map(Array(_))
        )
      )
    }
  }

  protected def setStoreCustomMetrics(customMetrics: Map[StateStoreCustomMetric, Long]): Unit = {
    customMetrics.foreach {
      case (metric, value) =>
        longMetric(metric.name) += value
    }
  }

  protected def setStoreInstanceMetrics(
      otherStoreInstanceMetrics: Map[StateStoreInstanceMetric, Long]): Unit = {
    otherStoreInstanceMetrics.foreach {
      case (metric, value) =>
        // Update the metric's value based on the defined combine method
        instanceMetrics(metric).set(metric.combine(instanceMetrics(metric), value))
    }
  }

  private def stateStoreCustomMetrics: Map[String, SQLMetric] = {
    val provider = StateStoreProvider.create(conf.stateStoreProviderClass)
    provider.supportedCustomMetrics.map {
      metric => (metric.name, metric.createSQLMetric(sparkContext))
    }.toMap
  }

  private def stateStoreInstanceMetrics: Map[StateStoreInstanceMetric, SQLMetric] = {
    val provider = StateStoreProvider.create(conf.stateStoreProviderClass)
    val maxPartitions = stateInfo.map(_.numPartitions).getOrElse(conf.defaultNumShufflePartitions)

    (0 until maxPartitions).flatMap { partitionId =>
      provider.supportedInstanceMetrics.flatMap { metric =>
        stateStoreNames.map { storeName =>
          val metricWithPartition = metric.withNewId(partitionId, storeName)
          (metricWithPartition, metricWithPartition.createSQLMetric(sparkContext))
        }
      }
    }.toMap
  }

  /**
   * Set of stateful operator custom metrics. These are captured as part of the generic
   * key-value map [[StateOperatorProgress.customMetrics]].
   * Stateful operators can extend this method to provide their own unique custom metrics.
   */
  protected def customStatefulOperatorMetrics: Seq[StatefulOperatorCustomMetric] = Nil

  private def statefulOperatorCustomMetrics: Map[String, SQLMetric] = {
    customStatefulOperatorMetrics.map {
      metric => (metric.name, metric.createSQLMetric(sparkContext))
    }.toMap
  }

  protected def applyRemovingRowsOlderThanWatermark(
      iter: Iterator[InternalRow],
      predicateDropRowByWatermark: BasePredicate): Iterator[InternalRow] = {
    iter.filterNot { row =>
      val shouldDrop = predicateDropRowByWatermark.eval(row)
      if (shouldDrop) longMetric("numRowsDroppedByWatermark") += 1
      shouldDrop
    }
  }

  /** Name to output in [[StreamingOperatorProgress]] to identify operator type */
  def shortName: String = "defaultName"

  def validateNewMetadata(
      oldMetadata: OperatorStateMetadata,
      newMetadata: OperatorStateMetadata): Unit = {}

  /**
   * Should the MicroBatchExecution run another batch based on this stateful operator and the
   * new input watermark.
   */
  def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = false
}

/** An operator that supports watermark. */
trait WatermarkSupport extends SparkPlan {

  def child: SparkPlan

  /** The keys that may have a watermark attribute. */
  def keyExpressions: Seq[Attribute]

  /**
   * The watermark value for filtering late events/records. This should be the previous
   * batch state eviction watermark.
   */
  def eventTimeWatermarkForLateEvents: Option[Long]
  /**
   * The watermark value for closing aggregates and evicting state.
   * It is different from the late events filtering watermark (consider chained aggregators
   * agg1 -> agg2: agg1 evicts state which will be effectively late against the eviction watermark
   * but should not be late for agg2 input late record filtering watermark. Thus agg1 and agg2 use
   * the current batch watermark for state eviction but the previous batch watermark for late
   * record filtering.
   */
  def eventTimeWatermarkForEviction: Option[Long]

  /** Generate an expression that matches data older than late event filtering watermark */
  lazy val watermarkExpressionForLateEvents: Option[Expression] =
    watermarkExpression(eventTimeWatermarkForLateEvents)
  /** Generate an expression that matches data older than the state eviction watermark */
  lazy val watermarkExpressionForEviction: Option[Expression] =
    watermarkExpression(eventTimeWatermarkForEviction)

  lazy val allowMultipleStatefulOperators: Boolean =
    conf.getConf(SQLConf.STATEFUL_OPERATOR_ALLOW_MULTIPLE)

  /** Generate an expression that matches data older than the watermark */
  private def watermarkExpression(watermark: Option[Long]): Option[Expression] = {
    WatermarkSupport.watermarkExpression(
      WatermarkSupport.findEventTimeColumn(child.output,
        allowMultipleEventTimeColumns = !allowMultipleStatefulOperators), watermark)
  }

  /** Predicate based on keys that matches data older than the late event filtering watermark */
  lazy val watermarkPredicateForKeysForLateEvents: Option[BasePredicate] =
    watermarkPredicateForKeys(watermarkExpressionForLateEvents)

  /** Generate an expression that matches data older than the state eviction watermark */
  lazy val watermarkPredicateForKeysForEviction: Option[BasePredicate] =
    watermarkPredicateForKeys(watermarkExpressionForEviction)

  private def watermarkPredicateForKeys(
      watermarkExpression: Option[Expression]): Option[BasePredicate] = {
    watermarkExpression.flatMap { e =>
      if (keyExpressions.exists(_.metadata.contains(EventTimeWatermark.delayKey))) {
        Some(Predicate.create(e, keyExpressions))
      } else {
        None
      }
    }
  }

  /**
   * Predicate based on the child output that matches data older than the watermark for late events
   * filtering.
   */
  lazy val watermarkPredicateForDataForLateEvents: Option[BasePredicate] =
    watermarkPredicateForData(watermarkExpressionForLateEvents)

  lazy val watermarkPredicateForDataForEviction: Option[BasePredicate] =
    watermarkPredicateForData(watermarkExpressionForEviction)

  private def watermarkPredicateForData(
    watermarkExpression: Option[Expression]): Option[BasePredicate] = {
    watermarkExpression.map(Predicate.create(_, child.output))
  }

  protected def removeKeysOlderThanWatermark(store: StateStore): Unit = {
    if (watermarkPredicateForKeysForEviction.nonEmpty) {
      val numRemovedStateRows = longMetric("numRemovedStateRows")
      store.iterator().foreach { rowPair =>
        if (watermarkPredicateForKeysForEviction.get.eval(rowPair.key)) {
          store.remove(rowPair.key)
          numRemovedStateRows += 1
        }
      }
    }
  }

  protected def removeKeysOlderThanWatermark(
      storeManager: StreamingAggregationStateManager,
      store: StateStore): Unit = {
    if (watermarkPredicateForKeysForEviction.nonEmpty) {
      val numRemovedStateRows = longMetric("numRemovedStateRows")
      storeManager.keys(store).foreach { keyRow =>
        if (watermarkPredicateForKeysForEviction.get.eval(keyRow)) {
          storeManager.remove(store, keyRow)
          numRemovedStateRows += 1
        }
      }
    }
  }
}

object WatermarkSupport {

  /** Generate an expression on given attributes that matches data older than the watermark */
  def watermarkExpression(
      optionalWatermarkExpression: Option[Expression],
      optionalWatermarkMs: Option[Long]): Option[Expression] = {
    if (optionalWatermarkExpression.isEmpty || optionalWatermarkMs.isEmpty) return None

    val watermarkAttribute = optionalWatermarkExpression.get
    // If we are evicting based on a window, use the end of the window. Otherwise just
    // use the attribute itself.
    val evictionExpression =
      if (watermarkAttribute.dataType.isInstanceOf[StructType]) {
        LessThanOrEqual(
          GetStructField(watermarkAttribute, 1),
          Literal(optionalWatermarkMs.get * 1000))
      } else {
        LessThanOrEqual(
          watermarkAttribute,
          Literal(optionalWatermarkMs.get * 1000))
      }
    Some(evictionExpression)
  }

  /**
   * Find the column which is marked as "event time" column.
   *
   * If there are multiple event time columns in given column list, the behavior depends on the
   * parameter `allowMultipleEventTimeColumns`. If it's set to true, the first occurred column will
   * be returned. If not, this method will throw an AnalysisException as it is not allowed to have
   * multiple event time columns.
   */
  def findEventTimeColumn(
      attrs: Seq[Attribute],
      allowMultipleEventTimeColumns: Boolean): Option[Attribute] = {
    val eventTimeCols = attrs.filter(_.metadata.contains(EventTimeWatermark.delayKey))
    if (!allowMultipleEventTimeColumns) {
      // There is a case projection leads the same column (same exprId) to appear more than one
      // time. Allowing them does not hurt the correctness of state row eviction, hence let's start
      // with allowing them.
      val eventTimeColsSet = eventTimeCols.map(_.exprId).toSet
      if (eventTimeColsSet.size > 1) {
        throw new AnalysisException(
          // TODO: [SPARK-55731] Assign error class for _LEGACY_ERROR_TEMP_3077
          errorClass = "_LEGACY_ERROR_TEMP_3077",
          messageParameters = Map("eventTimeCols" -> eventTimeCols.mkString("(", ",", ")")))
      }

      // With above check, even there are multiple columns in eventTimeCols, all columns must be
      // the same.
    } else {
      // This is for compatibility with previous behavior - we allow multiple distinct event time
      // columns and pick up the first occurrence. This is incorrect if non-first occurrence is
      // not smaller than the first one, but allow this as "escape hatch" in case we break the
      // existing query.
    }
    // pick the first element if exists
    eventTimeCols.headOption
  }

  /**
   * Find the index of the column which is marked as "event time" column.
   *
   * If there are multiple event time columns in given column list, the behavior depends on the
   * parameter `allowMultipleEventTimeColumns`. If it's set to true, the first occurred column will
   * be returned. If not, this method will throw an AnalysisException as it is not allowed to have
   * multiple event time columns.
   */
  def findEventTimeColumnIndex(
      attrs: Seq[Attribute],
      allowMultipleEventTimeColumns: Boolean): Option[Int] = {
    val eventTimeCols = attrs.zipWithIndex
      .filter(_._1.metadata.contains(EventTimeWatermark.delayKey))
    if (!allowMultipleEventTimeColumns) {
      // There is a case projection leads the same column (same exprId) to appear more than one
      // time. Allowing them does not hurt the correctness of state row eviction, hence let's start
      // with allowing them.
      val eventTimeColsSet = eventTimeCols.map(_._1.exprId).toSet
      if (eventTimeColsSet.size > 1) {
        throw new AnalysisException(
          // TODO: [SPARK-55731] Assign error class for _LEGACY_ERROR_TEMP_3077
          errorClass = "_LEGACY_ERROR_TEMP_3077",
          messageParameters = Map("eventTimeCols" -> eventTimeCols.mkString("(", ",", ")")))
      }

      // With above check, even there are multiple columns in eventTimeCols, all columns must be
      // the same.
    } else {
      // This is for compatibility with previous behavior - we allow multiple distinct event time
      // columns and pick up the first occurrence. This is incorrect if non-first occurrence is
      // not smaller than the first one, but allow this as "escape hatch" in case we break the
      // existing query.
    }
    // pick the first element if exists
    eventTimeCols.headOption.map(_._2)
  }
}

/**
 * For each input tuple, the key is calculated and the value from the [[StateStore]] is added
 * to the stream (in addition to the input tuple) if present.
 */
case class StateStoreRestoreExec(
    keyExpressions: Seq[Attribute],
    stateInfo: Option[StatefulOperatorStateInfo],
    stateFormatVersion: Int,
    child: SparkPlan)
  extends UnaryExecNode with StateStoreReader {

  private[sql] val stateManager = StreamingAggregationStateManager.createStateManager(
    keyExpressions, child.output, stateFormatVersion)

  override def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
    List[StateSchemaValidationResult] = {
    val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME,
      0, keyExpressions.toStructType, 0, stateManager.getStateValueSchema))
    List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo,
      hadoopConf, newStateSchema, session.sessionState, stateSchemaVersion))
  }

  override protected def doExecute(): RDD[InternalRow] = {
    val numOutputRows = longMetric("numOutputRows")

    child.execute().mapPartitionsWithReadStateStore(
      getStateInfo,
      keyExpressions.toStructType,
      stateManager.getStateValueSchema,
      NoPrefixKeyStateEncoderSpec(keyExpressions.toStructType),
      session.sessionState,
      Some(session.streams.stateStoreCoordinator)) { case (store, iter) =>
      val hasInput = iter.hasNext
      if (!hasInput && keyExpressions.isEmpty) {
        // If our `keyExpressions` are empty, we're getting a global aggregation. In that case
        // the `HashAggregateExec` will output a 0 value for the partial merge. We need to
        // restore the value, so that we don't overwrite our state with a 0 value, but rather
        // merge the 0 with existing state.
        store.iterator().map(_.value)
      } else {
        iter.flatMap { row =>
          val key = stateManager.getKey(row.asInstanceOf[UnsafeRow])
          val restoredRow = stateManager.get(store, key)
          val outputRows = Option(restoredRow).toSeq :+ row
          numOutputRows += outputRows.size
          outputRows
        }
      }
    }
  }

  override def output: Seq[Attribute] = child.output

  override def outputPartitioning: Partitioning = child.outputPartitioning

  override def requiredChildDistribution: Seq[Distribution] = {
    if (keyExpressions.isEmpty) {
      AllTuples :: Nil
    } else {
      StatefulOperatorPartitioning.getCompatibleDistribution(
        keyExpressions, getStateInfo, conf) :: Nil
    }
  }

  override protected def withNewChildInternal(newChild: SparkPlan): StateStoreRestoreExec =
    copy(child = newChild)
}

/**
 * For each input tuple, the key is calculated and the tuple is `put` into the [[StateStore]].
 */
case class StateStoreSaveExec(
    keyExpressions: Seq[Attribute],
    stateInfo: Option[StatefulOperatorStateInfo] = None,
    outputMode: Option[OutputMode] = None,
    eventTimeWatermarkForLateEvents: Option[Long] = None,
    eventTimeWatermarkForEviction: Option[Long] = None,
    stateFormatVersion: Int,
    child: SparkPlan)
  extends UnaryExecNode with StateStoreWriter with WatermarkSupport {

  private[sql] val stateManager = StreamingAggregationStateManager.createStateManager(
    keyExpressions, child.output, stateFormatVersion)

  override def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration,
      batchId: Long,
      stateSchemaVersion: Int): List[StateSchemaValidationResult] = {
    val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME,
      keySchemaId = 0, keyExpressions.toStructType, valueSchemaId = 0,
      stateManager.getStateValueSchema))
    List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo,
      hadoopConf, newStateSchema, session.sessionState, stateSchemaVersion))
  }

  override protected def doExecute(): RDD[InternalRow] = {
    metrics // force lazy init at driver
    assert(outputMode.nonEmpty,
      "Incorrect planning in IncrementalExecution, outputMode has not been set")

    child.execute().mapPartitionsWithStateStore(
      getStateInfo,
      keyExpressions.toStructType,
      stateManager.getStateValueSchema,
      NoPrefixKeyStateEncoderSpec(keyExpressions.toStructType),
      session.sessionState,
      Some(session.streams.stateStoreCoordinator)) { (store, iter) =>
        val numOutputRows = longMetric("numOutputRows")
        val numUpdatedStateRows = longMetric("numUpdatedStateRows")
        val allUpdatesTimeMs = longMetric("allUpdatesTimeMs")
        val numRemovedStateRows = longMetric("numRemovedStateRows")
        val allRemovalsTimeMs = longMetric("allRemovalsTimeMs")
        val commitTimeMs = longMetric("commitTimeMs")

        outputMode match {
          // Update and output all rows in the StateStore.
          case Some(Complete) =>
            allUpdatesTimeMs += timeTakenMs {
              while (iter.hasNext) {
                val row = iter.next().asInstanceOf[UnsafeRow]
                stateManager.put(store, row)
                numUpdatedStateRows += 1
              }
            }

            // SPARK-45582 - Ensure that store instance is not used after commit is called
            // to invoke the iterator.
            val rangeIter = stateManager.values(store)

            new NextIterator[InternalRow] {
              override protected def getNext(): InternalRow = {
                if (rangeIter.hasNext) {
                  val valueRow = rangeIter.next()
                  numOutputRows += 1
                  valueRow
                } else {
                  finished = true
                  null
                }
              }

              override protected def close(): Unit = {
                allRemovalsTimeMs += 0
                commitTimeMs += timeTakenMs {
                  stateManager.commit(store)
                }
                setStoreMetrics(store)
                setOperatorMetrics()
              }
            }

          // Update and output only rows being evicted from the StateStore
          // Assumption: watermark predicates must be non-empty if append mode is allowed
          case Some(Append) =>
            if (watermarkPredicateForDataForLateEvents.isEmpty ||
              watermarkPredicateForKeysForEviction.isEmpty) {
              throw QueryExecutionErrors.unsupportedStreamingOperatorWithoutWatermark(
                "Append", "aggregations")
            }

            allUpdatesTimeMs += timeTakenMs {
              val filteredIter = applyRemovingRowsOlderThanWatermark(iter,
                watermarkPredicateForDataForLateEvents.get)
              while (filteredIter.hasNext) {
                val row = filteredIter.next().asInstanceOf[UnsafeRow]
                stateManager.put(store, row)
                numUpdatedStateRows += 1
              }
            }

            val removalStartTimeNs = System.nanoTime
            val rangeIter = stateManager.iterator(store)

            new NextIterator[InternalRow] {
              override protected def getNext(): InternalRow = {
                var removedValueRow: InternalRow = null
                while (rangeIter.hasNext && removedValueRow == null) {
                  val rowPair = rangeIter.next()
                  if (watermarkPredicateForKeysForEviction.get.eval(rowPair.key)) {
                    stateManager.remove(store, rowPair.key)
                    numRemovedStateRows += 1
                    removedValueRow = rowPair.value
                  }
                }
                if (removedValueRow == null) {
                  finished = true
                  null
                } else {
                  numOutputRows += 1
                  removedValueRow
                }
              }

              override protected def close(): Unit = {
                // Note: Due to the iterator lazy exec, this metric also captures the time taken
                // by the consumer operators in addition to the processing in this operator.
                allRemovalsTimeMs += NANOSECONDS.toMillis(System.nanoTime - removalStartTimeNs)
                commitTimeMs += timeTakenMs { stateManager.commit(store) }
                setStoreMetrics(store)
                setOperatorMetrics()
              }
            }

          // Update and output modified rows from the StateStore.
          case Some(Update) =>

            new NextIterator[InternalRow] {
              // Filter late date using watermark if specified
              private[this] val baseIterator = watermarkPredicateForDataForLateEvents match {
                case Some(predicate) => applyRemovingRowsOlderThanWatermark(iter, predicate)
                case None => iter
              }
              private val updatesStartTimeNs = System.nanoTime

              override protected def getNext(): InternalRow = {
                if (baseIterator.hasNext) {
                  val row = baseIterator.next().asInstanceOf[UnsafeRow]
                  stateManager.put(store, row)
                  numOutputRows += 1
                  numUpdatedStateRows += 1
                  row
                } else {
                  finished = true
                  null
                }
              }

              override protected def close(): Unit = {
                allUpdatesTimeMs += NANOSECONDS.toMillis(System.nanoTime - updatesStartTimeNs)

                // Remove old aggregates if watermark specified
                allRemovalsTimeMs += timeTakenMs {
                  removeKeysOlderThanWatermark(stateManager, store)
                }
                commitTimeMs += timeTakenMs { stateManager.commit(store) }
                setStoreMetrics(store)
                setOperatorMetrics()
              }
            }

          case _ => throw QueryExecutionErrors.unsupportedOutputModeForStreamingOperationError(
            outputMode.get, "streaming aggregations")
        }
    }
  }

  override def output: Seq[Attribute] = child.output

  override def outputPartitioning: Partitioning = child.outputPartitioning

  override def requiredChildDistribution: Seq[Distribution] = {
    if (keyExpressions.isEmpty) {
      AllTuples :: Nil
    } else {
      StatefulOperatorPartitioning.getCompatibleDistribution(
        keyExpressions, getStateInfo, conf) :: Nil
    }
  }

  override def shortName: String = StatefulOperatorsUtils.STATE_STORE_SAVE_EXEC_OP_NAME

  override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = {
    (outputMode.contains(Append) || outputMode.contains(Update)) &&
      eventTimeWatermarkForEviction.isDefined &&
      newInputWatermark > eventTimeWatermarkForEviction.get
  }

  override protected def withNewChildInternal(newChild: SparkPlan): StateStoreSaveExec =
    copy(child = newChild)
}

/**
 * This class sorts input rows and existing sessions in state and provides output rows as
 * sorted by "group keys + start time of session window".
 *
 * Refer [[MergingSortWithSessionWindowStateIterator]] for more details.
 */
case class SessionWindowStateStoreRestoreExec(
    keyWithoutSessionExpressions: Seq[Attribute],
    sessionExpression: Attribute,
    stateInfo: Option[StatefulOperatorStateInfo],
    eventTimeWatermarkForLateEvents: Option[Long],
    eventTimeWatermarkForEviction: Option[Long],
    stateFormatVersion: Int,
    child: SparkPlan)
  extends UnaryExecNode with StateStoreReader with WatermarkSupport {

  override lazy val metrics = Map(
    "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
    "numRowsDroppedByWatermark" -> SQLMetrics.createMetric(sparkContext,
      "number of rows which are dropped by watermark")
  )

  override def keyExpressions: Seq[Attribute] = keyWithoutSessionExpressions

  assert(keyExpressions.nonEmpty, "Grouping key must be specified when using sessionWindow")

  private val stateManager = StreamingSessionWindowStateManager.createStateManager(
    keyWithoutSessionExpressions, sessionExpression, child.output, stateFormatVersion)

  override def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
    List[StateSchemaValidationResult] = {
    val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME,
      keySchemaId = 0, stateManager.getStateKeySchema, valueSchemaId = 0,
      stateManager.getStateValueSchema))
    List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo, hadoopConf,
      newStateSchema, session.sessionState, stateSchemaVersion))
  }

  override protected def doExecute(): RDD[InternalRow] = {
    val numOutputRows = longMetric("numOutputRows")

    child.execute().mapPartitionsWithReadStateStore(
      getStateInfo,
      stateManager.getStateKeySchema,
      stateManager.getStateValueSchema,
      PrefixKeyScanStateEncoderSpec(stateManager.getStateKeySchema,
        stateManager.getNumColsForPrefixKey),
      session.sessionState,
      Some(session.streams.stateStoreCoordinator)) { case (store, iter) =>

      // We need to filter out outdated inputs
      val filteredIterator = watermarkPredicateForDataForLateEvents match {
        case Some(predicate) => iter.filter((row: InternalRow) => {
          val shouldKeep = !predicate.eval(row)
          if (!shouldKeep) longMetric("numRowsDroppedByWatermark") += 1
          shouldKeep
        })
        case None => iter
      }

      new MergingSortWithSessionWindowStateIterator(
        filteredIterator,
        stateManager,
        store,
        keyWithoutSessionExpressions,
        sessionExpression,
        child.output).map { row =>
          numOutputRows += 1
          row
      }
    }
  }

  override def output: Seq[Attribute] = child.output

  override def outputPartitioning: Partitioning = child.outputPartitioning

  override def outputOrdering: Seq[SortOrder] = {
    (keyWithoutSessionExpressions ++ Seq(sessionExpression)).map(SortOrder(_, Ascending))
  }

  override def requiredChildDistribution: Seq[Distribution] = {
    StatefulOperatorPartitioning.getCompatibleDistribution(
      keyWithoutSessionExpressions, getStateInfo, conf) :: Nil
  }

  override def requiredChildOrdering: Seq[Seq[SortOrder]] = {
    Seq((keyWithoutSessionExpressions ++ Seq(sessionExpression)).map(SortOrder(_, Ascending)))
  }

  override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
    copy(child = newChild)
}

/**
 * This class replaces existing sessions for the grouping key with new sessions in state store.
 * All inputs are valid on storing into state store; don't filter out via watermark while storing.
 * Refer the method doc of [[StreamingSessionWindowStateManager.updateSessions]] for more details.
 *
 * This class will provide the output according to the output mode.
 * Update mode is not supported as the semantic is not feasible for session window.
 */
case class SessionWindowStateStoreSaveExec(
    keyWithoutSessionExpressions: Seq[Attribute],
    sessionExpression: Attribute,
    stateInfo: Option[StatefulOperatorStateInfo] = None,
    outputMode: Option[OutputMode] = None,
    eventTimeWatermarkForLateEvents: Option[Long] = None,
    eventTimeWatermarkForEviction: Option[Long] = None,
    stateFormatVersion: Int,
    child: SparkPlan)
  extends UnaryExecNode with StateStoreWriter with WatermarkSupport {

  override def keyExpressions: Seq[Attribute] = keyWithoutSessionExpressions

  override def shortName: String =
    StatefulOperatorsUtils.SESSION_WINDOW_STATE_STORE_SAVE_EXEC_OP_NAME

  private val stateManager = StreamingSessionWindowStateManager.createStateManager(
    keyWithoutSessionExpressions, sessionExpression, child.output, stateFormatVersion)

  override def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
    List[StateSchemaValidationResult] = {
    val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME, 0,
      stateManager.getStateKeySchema, 0, stateManager.getStateValueSchema))
    List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo, hadoopConf,
      newStateSchema, session.sessionState, stateSchemaVersion))
  }

  override protected def doExecute(): RDD[InternalRow] = {
    metrics // force lazy init at driver
    assert(outputMode.nonEmpty,
      "Incorrect planning in IncrementalExecution, outputMode has not been set")
    assert(keyExpressions.nonEmpty,
      "Grouping key must be specified when using sessionWindow")

    child.execute().mapPartitionsWithStateStore(
      getStateInfo,
      stateManager.getStateKeySchema,
      stateManager.getStateValueSchema,
      PrefixKeyScanStateEncoderSpec(stateManager.getStateKeySchema,
        stateManager.getNumColsForPrefixKey),
      session.sessionState,
      Some(session.streams.stateStoreCoordinator)) { case (store, iter) =>

      val numOutputRows = longMetric("numOutputRows")
      val numRemovedStateRows = longMetric("numRemovedStateRows")
      val allUpdatesTimeMs = longMetric("allUpdatesTimeMs")
      val allRemovalsTimeMs = longMetric("allRemovalsTimeMs")
      val commitTimeMs = longMetric("commitTimeMs")

      outputMode match {
        // Update and output all rows in the StateStore.
        case Some(Complete) =>
          allUpdatesTimeMs += timeTakenMs {
            putToStore(iter, store)
          }

          // SPARK-45582 - Ensure that store instance is not used after commit is called
          // to invoke the iterator.
          val rangeIter = stateManager.iterator(store)

          new NextIterator[InternalRow] {
            override protected def getNext(): InternalRow = {
              if (rangeIter.hasNext) {
                val valueRow = rangeIter.next()
                numOutputRows += 1
                valueRow
              } else {
                finished = true
                null
              }
            }

            override protected def close(): Unit = {
              commitTimeMs += timeTakenMs {
                stateManager.commit(store)
              }
              setStoreMetrics(store)
            }
          }

        // Update and output only rows being evicted from the StateStore
        // Assumption: watermark predicates must be non-empty if append mode is allowed
        case Some(Append) =>
          if (watermarkPredicateForDataForEviction.isEmpty) {
            throw QueryExecutionErrors.unsupportedStreamingOperatorWithoutWatermark(
              "Append", "session window aggregations")
          }

          allUpdatesTimeMs += timeTakenMs {
            putToStore(iter, store)
          }

          val removalStartTimeNs = System.nanoTime
          new NextIterator[InternalRow] {
            private val removedIter = stateManager.removeByValueCondition(
              store, watermarkPredicateForDataForEviction.get.eval)

            override protected def getNext(): InternalRow = {
              if (!removedIter.hasNext) {
                finished = true
                null
              } else {
                numRemovedStateRows += 1
                numOutputRows += 1
                removedIter.next()
              }
            }

            override protected def close(): Unit = {
              allRemovalsTimeMs += NANOSECONDS.toMillis(System.nanoTime - removalStartTimeNs)
              commitTimeMs += timeTakenMs { store.commit() }
              setStoreMetrics(store)
              setOperatorMetrics()
            }
          }

        case _ => throw throw QueryExecutionErrors.unsupportedOutputModeForStreamingOperationError(
          outputMode.get, "session window streaming aggregations")
      }
    }
  }

  override def output: Seq[Attribute] = child.output

  override def outputPartitioning: Partitioning = child.outputPartitioning

  override def requiredChildDistribution: Seq[Distribution] = {
    StatefulOperatorPartitioning.getCompatibleDistribution(
      keyWithoutSessionExpressions, getStateInfo, conf) :: Nil
  }

  override def operatorStateMetadata(
      stateSchemaPaths: List[List[String]] = List.empty): OperatorStateMetadata = {
    val info = getStateInfo
    val operatorInfo = OperatorInfoV1(info.operatorId, shortName)
    val stateStoreInfo = Array(StateStoreMetadataV1(
      StateStoreId.DEFAULT_STORE_NAME, stateManager.getNumColsForPrefixKey, info.numPartitions))
    OperatorStateMetadataV1(operatorInfo, stateStoreInfo)
  }

  override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = {
    (outputMode.contains(Append) || outputMode.contains(Update)) &&
      eventTimeWatermarkForEviction.isDefined &&
      newInputWatermark > eventTimeWatermarkForEviction.get
  }

  private def putToStore(iter: Iterator[InternalRow], store: StateStore): Unit = {
    val numUpdatedStateRows = longMetric("numUpdatedStateRows")
    val numRemovedStateRows = longMetric("numRemovedStateRows")

    var curKey: UnsafeRow = null
    val curValuesOnKey = new mutable.ArrayBuffer[UnsafeRow]()

    def applyChangesOnKey(): Unit = {
      if (curValuesOnKey.nonEmpty) {
        val (upserted, deleted) = stateManager.updateSessions(store, curKey, curValuesOnKey.toSeq)
        numUpdatedStateRows += upserted
        numRemovedStateRows += deleted
        curValuesOnKey.clear()
      }
    }

    while (iter.hasNext) {
      val row = iter.next().asInstanceOf[UnsafeRow]
      val key = stateManager.extractKeyWithoutSession(row)

      if (curKey == null || curKey != key) {
        // new group appears
        applyChangesOnKey()
        curKey = key.copy()
      }

      // must copy the row, for this row is a reference in iterator and
      // will change when iter.next
      curValuesOnKey += row.copy
    }

    applyChangesOnKey()
  }

  override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
    copy(child = newChild)

  /**
   * The class overrides this method since dropping late events are happening in the upstream node
   * [[SessionWindowStateStoreRestoreExec]], and this class has responsibility to report the number
   * of dropped late events as a part of StateOperatorProgress.
   *
   * This method should be called in the driver after this SparkPlan has been executed and metrics
   * have been updated.
   */
  override def getProgress(): StateOperatorProgress = {
    val stateOpProgress = super.getProgress()

    // This should be safe, since the method is called in the driver after the plan has been
    // executed and metrics have been updated.
    val numRowsDroppedByWatermark = child.collectFirst {
      case s: SessionWindowStateStoreRestoreExec =>
        s.longMetric("numRowsDroppedByWatermark").value
    }.getOrElse(0L)

    stateOpProgress.copy(
      newNumRowsUpdated = stateOpProgress.numRowsUpdated,
      newNumRowsDroppedByWatermark = numRowsDroppedByWatermark,
      newNumRowsRemoved = stateOpProgress.numRowsRemoved)
  }
}

abstract class BaseStreamingDeduplicateExec
  extends UnaryExecNode with StateStoreWriter with WatermarkSupport {

  def keyExpressions: Seq[Attribute]
  def child: SparkPlan
  def stateInfo: Option[StatefulOperatorStateInfo]
  def eventTimeWatermarkForLateEvents: Option[Long]
  def eventTimeWatermarkForEviction: Option[Long]

  /** Distribute by grouping attributes */
  override def requiredChildDistribution: Seq[Distribution] = {
    StatefulOperatorPartitioning.getCompatibleDistribution(
      keyExpressions, getStateInfo, conf) :: Nil
  }

  protected val schemaForValueRow: StructType
  protected val extraOptionOnStateStore: Map[String, String]

  override protected def doExecute(): RDD[InternalRow] = {
    metrics // force lazy init at driver

    child.execute().mapPartitionsWithStateStore(
      getStateInfo,
      keyExpressions.toStructType,
      schemaForValueRow,
      NoPrefixKeyStateEncoderSpec(keyExpressions.toStructType),
      session.sessionState,
      Some(session.streams.stateStoreCoordinator),
      extraOptions = extraOptionOnStateStore) { (store, iter) =>
      // Normalize NaN and -0.0 in floating-point key values so that semantically equal
      // values produce identical UnsafeRow bytes.
      val getKey = GenerateUnsafeProjection.generate(
        keyExpressions.map(NormalizeFloatingNumbers.normalize), child.output)
      val numOutputRows = longMetric("numOutputRows")
      val numUpdatedStateRows = longMetric("numUpdatedStateRows")
      val allUpdatesTimeMs = longMetric("allUpdatesTimeMs")
      val allRemovalsTimeMs = longMetric("allRemovalsTimeMs")
      val commitTimeMs = longMetric("commitTimeMs")
      val numDroppedDuplicateRows = longMetric("numDroppedDuplicateRows")

      val baseIterator = watermarkPredicateForDataForLateEvents match {
        case Some(predicate) => applyRemovingRowsOlderThanWatermark(iter, predicate)
        case None => iter
      }

      val reusedDupInfoRow = initializeReusedDupInfoRow()

      val updatesStartTimeNs = System.nanoTime

      val result = baseIterator.filter { r =>
        val row = r.asInstanceOf[UnsafeRow]
        val key = getKey(row)
        val keyExists = store.keyExists(key)
        if (!keyExists) {
          putDupInfoIntoState(store, row, key, reusedDupInfoRow)
          numUpdatedStateRows += 1
          numOutputRows += 1
          true
        } else {
          // Drop duplicated rows
          numDroppedDuplicateRows += 1
          false
        }
      }

      CompletionIterator[InternalRow, Iterator[InternalRow]](result, {
        allUpdatesTimeMs += NANOSECONDS.toMillis(System.nanoTime - updatesStartTimeNs)
        allRemovalsTimeMs += timeTakenMs { evictDupInfoFromState(store) }
        commitTimeMs += timeTakenMs { store.commit() }
        setStoreMetrics(store)
        setOperatorMetrics()
      })
    }
  }

  protected def initializeReusedDupInfoRow(): Option[UnsafeRow]

  protected def putDupInfoIntoState(
      store: StateStore,
      data: UnsafeRow,
      key: UnsafeRow,
      reusedDupInfoRow: Option[UnsafeRow]): Unit

  protected def evictDupInfoFromState(store: StateStore): Unit

  override def output: Seq[Attribute] = child.output

  override def outputPartitioning: Partitioning = child.outputPartitioning

  override def customStatefulOperatorMetrics: Seq[StatefulOperatorCustomMetric] = {
    Seq(StatefulOperatorCustomSumMetric("numDroppedDuplicateRows", "number of duplicates dropped"))
  }

  override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = {
    eventTimeWatermarkForEviction.isDefined &&
      newInputWatermark > eventTimeWatermarkForEviction.get
  }
}

/** Physical operator for executing streaming Deduplicate. */
case class StreamingDeduplicateExec(
    keyExpressions: Seq[Attribute],
    child: SparkPlan,
    stateInfo: Option[StatefulOperatorStateInfo] = None,
    eventTimeWatermarkForLateEvents: Option[Long] = None,
    eventTimeWatermarkForEviction: Option[Long] = None)
  extends BaseStreamingDeduplicateExec {

  protected val schemaForValueRow: StructType =
    StructType(Array(StructField("__dummy__", NullType)))

  // We won't check value row in state store since the value StreamingDeduplicateExec.EMPTY_ROW
  // is unrelated to the output schema.
  protected val extraOptionOnStateStore: Map[String, String] =
    Map(StateStoreConf.FORMAT_VALIDATION_CHECK_VALUE_CONFIG -> "false")

  protected def initializeReusedDupInfoRow(): Option[UnsafeRow] = None

  protected def putDupInfoIntoState(
      store: StateStore,
      data: UnsafeRow,
      key: UnsafeRow,
      reusedDupInfoRow: Option[UnsafeRow]): Unit = {
    store.put(key, StreamingDeduplicateExec.EMPTY_ROW)
  }

  protected def evictDupInfoFromState(store: StateStore): Unit = {
    removeKeysOlderThanWatermark(store)
  }

  override def shortName: String = StatefulOperatorsUtils.DEDUPLICATE_EXEC_OP_NAME

  override protected def withNewChildInternal(newChild: SparkPlan): StreamingDeduplicateExec =
    copy(child = newChild)

  override def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
    List[StateSchemaValidationResult] = {
    val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME, 0,
      keyExpressions.toStructType, 0, schemaForValueRow))
    List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo, hadoopConf,
      newStateSchema, session.sessionState, stateSchemaVersion,
      extraOptions = extraOptionOnStateStore))
  }
}

object StreamingDeduplicateExec {
  private val EMPTY_ROW =
    UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null))
}

/**
 * For Deduplicate, the state key is the partition key i.e. the dedup key
 */
class StreamingDeduplicateStatePartitionKeyExtractor(stateKeySchema: StructType)
  extends NoopStatePartitionKeyExtractor(stateKeySchema)

case class StreamingDeduplicateWithinWatermarkExec(
    keyExpressions: Seq[Attribute],
    child: SparkPlan,
    stateInfo: Option[StatefulOperatorStateInfo] = None,
    eventTimeWatermarkForLateEvents: Option[Long] = None,
    eventTimeWatermarkForEviction: Option[Long] = None)
  extends BaseStreamingDeduplicateExec {

  protected val schemaForValueRow: StructType = StructType(
    Array(StructField("expiresAtMicros", LongType, nullable = false)))

  protected val extraOptionOnStateStore: Map[String, String] = Map.empty

  // Below three variables are defined as lazy, as evaluating these variables does not work with
  // canonicalized plan. Specifically, attributes in child won't have an event time column in
  // the canonicalized plan. These variables are NOT referenced in canonicalized plan, hence
  // defining these variables as lazy would avoid such error.
  private lazy val eventTimeCol: Attribute = WatermarkSupport.findEventTimeColumn(child.output,
    allowMultipleEventTimeColumns = false).get
  private lazy val delayThresholdMs = eventTimeCol.metadata.getLong(EventTimeWatermark.delayKey)
  private lazy val eventTimeColOrdinal: Int = child.output.indexOf(eventTimeCol)

  protected def initializeReusedDupInfoRow(): Option[UnsafeRow] = {
    val timeoutToUnsafeRow = UnsafeProjection.create(schemaForValueRow)
    val timeoutRow = timeoutToUnsafeRow(new SpecificInternalRow(schemaForValueRow))
    Some(timeoutRow)
  }

  protected def putDupInfoIntoState(
      store: StateStore,
      data: UnsafeRow,
      key: UnsafeRow,
      reusedDupInfoRow: Option[UnsafeRow]): Unit = {
    assert(reusedDupInfoRow.isDefined, "This should have reused row.")
    val timeoutRow = reusedDupInfoRow.get

    // We expect data type of event time column to be TimestampType or TimestampNTZType which both
    // are internally represented as Long.
    val timestamp = data.getLong(eventTimeColOrdinal)
    // The unit of timestamp in Spark is microseconds, convert the delay threshold to micros.
    val expiresAt = timestamp + DateTimeUtils.millisToMicros(delayThresholdMs)

    timeoutRow.setLong(0, expiresAt)
    store.put(key, timeoutRow)
  }

  protected def evictDupInfoFromState(store: StateStore): Unit = {
    val numRemovedStateRows = longMetric("numRemovedStateRows")

    // Convert watermark value to micros.
    val watermarkForEviction = DateTimeUtils.millisToMicros(eventTimeWatermarkForEviction.get)
    store.iterator().foreach { rowPair =>
      val valueRow = rowPair.value

      val expiresAt = valueRow.getLong(0)
      if (watermarkForEviction >= expiresAt) {
        store.remove(rowPair.key)
        numRemovedStateRows += 1
      }
    }
  }

  override def shortName: String = StatefulOperatorsUtils.DEDUPLICATE_WITHIN_WATERMARK_EXEC_OP_NAME

  override def validateAndMaybeEvolveStateSchema(
      hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int):
    List[StateSchemaValidationResult] = {
    val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME, 0,
      keyExpressions.toStructType, 0, schemaForValueRow))
    List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo, hadoopConf,
      newStateSchema, session.sessionState, stateSchemaVersion,
      extraOptions = extraOptionOnStateStore))
  }

  override protected def withNewChildInternal(
      newChild: SparkPlan): StreamingDeduplicateWithinWatermarkExec = copy(child = newChild)
}

/**
 * For DeduplicateWithinWatermark, the state key is the partition key i.e. the dedup key
 */
class StreamingDedupWithinWatermarkStatePartitionKeyExtractor(stateKeySchema: StructType)
  extends NoopStatePartitionKeyExtractor(stateKeySchema)

trait SchemaValidationUtils extends Logging {

  // Determines whether the operator should be able to evolve their schema
  val schemaEvolutionEnabledForOperator: Boolean = false

  // This method will return the column family schemas, and check whether the fields in the
  // schema are nullable. If Avro encoding is used, we want to enforce nullability
  def getColFamilySchemas(shouldBeNullable: Boolean): Map[String, StateStoreColFamilySchema]

  def validateAndWriteStateSchema(
      hadoopConf: Configuration,
      batchId: Long,
      stateSchemaVersion: Int,
      info: StatefulOperatorStateInfo,
      stateSchemaDir: Path,
      session: SparkSession,
      operatorStateMetadataVersion: Int = 2,
      stateStoreEncodingFormat: String = StateStoreEncoding.UnsafeRow.toString
  ): List[StateSchemaValidationResult] = {
    assert(stateSchemaVersion >= 3)
    val usingAvro = stateStoreEncodingFormat == StateStoreEncoding.Avro.toString
    val newSchemas = getColFamilySchemas(usingAvro)
    val newStateSchemaFilePath =
      new Path(stateSchemaDir, s"${batchId}_${UUID.randomUUID().toString}")
    val metadataPath = new Path(info.checkpointLocation, s"${info.operatorId}")
    val metadataReader = OperatorStateMetadataReader.createReader(
      metadataPath, hadoopConf, operatorStateMetadataVersion, batchId)
    val operatorStateMetadata = try {
      metadataReader.read()
    } catch {
      // If this is the first time we are running the query, there will be no metadata
      // and this error is expected. In this case, we return None.
      case _: Exception if batchId == 0 =>
        None
    }

    val oldStateSchemaFilePaths: List[Path] = operatorStateMetadata match {
      case Some(metadata) =>
        metadata match {
          case v2: OperatorStateMetadataV2 =>
            v2.stateStoreInfo.head.stateSchemaFilePaths.map(new Path(_))
          case _ => List.empty
        }
      case None => List.empty
    }
    // state schema file written here, writing the new schema list we passed here
    List(StateSchemaCompatibilityChecker.
      validateAndMaybeEvolveStateSchema(info, hadoopConf,
        newSchemas.values.toList, session.sessionState, stateSchemaVersion,
        storeName = StateStoreId.DEFAULT_STORE_NAME,
        oldSchemaFilePaths = oldStateSchemaFilePaths,
        newSchemaFilePath = Some(newStateSchemaFilePath),
        schemaEvolutionEnabled = usingAvro && schemaEvolutionEnabledForOperator))
  }
}

object StatefulOperatorsUtils {
  val TRANSFORM_WITH_STATE_EXEC_OP_NAME = "transformWithStateExec"
  val TRANSFORM_WITH_STATE_IN_PANDAS_EXEC_OP_NAME = "transformWithStateInPandasExec"
  val TRANSFORM_WITH_STATE_IN_PYSPARK_EXEC_OP_NAME = "transformWithStateInPySparkExec"
  // Seq of operator names who uses state schema v3 and TWS related options.
  val TRANSFORM_WITH_STATE_OP_NAMES: Seq[String] = Seq(
    TRANSFORM_WITH_STATE_EXEC_OP_NAME,
    TRANSFORM_WITH_STATE_IN_PANDAS_EXEC_OP_NAME,
    TRANSFORM_WITH_STATE_IN_PYSPARK_EXEC_OP_NAME
  )
  val SYMMETRIC_HASH_JOIN_EXEC_OP_NAME = "symmetricHashJoin"
  val STATE_STORE_SAVE_EXEC_OP_NAME = "stateStoreSave"
  val DEDUPLICATE_EXEC_OP_NAME = "dedupe"
  val DEDUPLICATE_WITHIN_WATERMARK_EXEC_OP_NAME = "dedupeWithinWatermark"
  val SESSION_WINDOW_STATE_STORE_SAVE_EXEC_OP_NAME = "sessionWindowStateStoreSaveExec"
  val FLAT_MAP_GROUPS_WITH_STATE_EXEC_OP_NAME = "flatMapGroupsWithState"
  val FLAT_MAP_GROUPS_IN_PANDAS_WITH_STATE_EXEC_OP_NAME = "applyInPandasWithState"
  val FLAT_MAP_GROUPS_OP_NAMES: Seq[String] = Seq(
    FLAT_MAP_GROUPS_WITH_STATE_EXEC_OP_NAME,
    FLAT_MAP_GROUPS_IN_PANDAS_WITH_STATE_EXEC_OP_NAME
  )
}
