| /* |
| * 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.state |
| |
| import java.util.UUID |
| import java.util.concurrent.{ScheduledFuture, TimeUnit} |
| import java.util.concurrent.atomic.AtomicReference |
| import javax.annotation.concurrent.GuardedBy |
| |
| import scala.collection.mutable |
| import scala.util.Try |
| import scala.util.control.NonFatal |
| |
| import org.apache.hadoop.conf.Configuration |
| import org.apache.hadoop.fs.Path |
| |
| import org.apache.spark.{SparkContext, SparkEnv, SparkUnsupportedOperationException} |
| import org.apache.spark.internal.{Logging, MDC} |
| import org.apache.spark.internal.LogKey._ |
| import org.apache.spark.sql.catalyst.expressions.UnsafeRow |
| import org.apache.spark.sql.catalyst.util.UnsafeRowUtils |
| import org.apache.spark.sql.errors.QueryExecutionErrors |
| import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} |
| import org.apache.spark.sql.execution.streaming.StatefulOperatorStateInfo |
| import org.apache.spark.sql.types.StructType |
| import org.apache.spark.util.{ThreadUtils, Utils} |
| |
| /** |
| * Base trait for a versioned key-value store which provides read operations. Each instance of a |
| * `ReadStateStore` represents a specific version of state data, and such instances are created |
| * through a [[StateStoreProvider]]. |
| * |
| * `abort` method will be called when the task is completed - please clean up the resources in |
| * the method. |
| * |
| * IMPLEMENTATION NOTES: |
| * * The implementation can throw exception on calling prefixScan method if the functionality is |
| * not supported yet from the implementation. Note that some stateful operations would not work |
| * on disabling prefixScan functionality. |
| */ |
| trait ReadStateStore { |
| |
| /** Unique identifier of the store */ |
| def id: StateStoreId |
| |
| /** Version of the data in this store before committing updates. */ |
| def version: Long |
| |
| /** |
| * Get the current value of a non-null key. |
| * @return a non-null row if the key exists in the store, otherwise null. |
| */ |
| def get( |
| key: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): UnsafeRow |
| |
| /** |
| * Provides an iterator containing all values of a non-null key. If key does not exist, |
| * an empty iterator is returned. Implementations should make sure to return an empty |
| * iterator if the key does not exist. |
| * |
| * It is expected to throw exception if Spark calls this method without setting |
| * multipleValuesPerKey as true for the column family. |
| */ |
| def valuesIterator(key: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRow] |
| |
| /** |
| * Return an iterator containing all the key-value pairs which are matched with |
| * the given prefix key. |
| * |
| * The operator will provide numColsPrefixKey greater than 0 in StateStoreProvider.init method |
| * if the operator needs to leverage the "prefix scan" feature. The schema of the prefix key |
| * should be same with the leftmost `numColsPrefixKey` columns of the key schema. |
| * |
| * It is expected to throw exception if Spark calls this method without setting numColsPrefixKey |
| * to the greater than 0. |
| */ |
| def prefixScan( |
| prefixKey: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair] |
| |
| /** Return an iterator containing all the key-value pairs in the StateStore. */ |
| def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair] |
| |
| /** |
| * Clean up the resource. |
| * |
| * The method name is to respect backward compatibility on [[StateStore]]. |
| */ |
| def abort(): Unit |
| } |
| |
| /** |
| * Base trait for a versioned key-value store which provides both read and write operations. Each |
| * instance of a `StateStore` represents a specific version of state data, and such instances are |
| * created through a [[StateStoreProvider]]. |
| * |
| * Unlike [[ReadStateStore]], `abort` method may not be called if the `commit` method succeeds |
| * to commit the change. (`hasCommitted` returns `true`.) Otherwise, `abort` method will be called. |
| * Implementation should deal with resource cleanup in both methods, and also need to guard with |
| * double resource cleanup. |
| */ |
| trait StateStore extends ReadStateStore { |
| |
| /** |
| * Remove column family with given name, if present. |
| */ |
| def removeColFamilyIfExists(colFamilyName: String): Boolean |
| |
| /** |
| * Create column family with given name, if absent. |
| */ |
| def createColFamilyIfAbsent( |
| colFamilyName: String, |
| keySchema: StructType, |
| valueSchema: StructType, |
| keyStateEncoderSpec: KeyStateEncoderSpec, |
| useMultipleValuesPerKey: Boolean = false, |
| isInternal: Boolean = false): Unit |
| |
| /** |
| * Put a new non-null value for a non-null key. Implementations must be aware that the UnsafeRows |
| * in the params can be reused, and must make copies of the data as needed for persistence. |
| */ |
| def put( |
| key: UnsafeRow, |
| value: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Unit |
| |
| /** |
| * Remove a single non-null key. |
| */ |
| def remove( |
| key: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Unit |
| |
| /** |
| * Merges the provided value with existing values of a non-null key. If a existing |
| * value does not exist, this operation behaves as [[StateStore.put()]]. |
| * |
| * It is expected to throw exception if Spark calls this method without setting |
| * multipleValuesPerKey as true for the column family. |
| */ |
| def merge(key: UnsafeRow, value: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Unit |
| |
| /** |
| * Commit all the updates that have been made to the store, and return the new version. |
| * Implementations should ensure that no more updates (puts, removes) can be after a commit in |
| * order to avoid incorrect usage. |
| */ |
| def commit(): Long |
| |
| /** |
| * Abort all the updates that have been made to the store. Implementations should ensure that |
| * no more updates (puts, removes) can be after an abort in order to avoid incorrect usage. |
| */ |
| override def abort(): Unit |
| |
| /** |
| * Return an iterator containing all the key-value pairs in the StateStore. Implementations must |
| * ensure that updates (puts, removes) can be made while iterating over this iterator. |
| * |
| * It is not required for implementations to ensure the iterator reflects all updates being |
| * performed after initialization of the iterator. Callers should perform all updates before |
| * calling this method if all updates should be visible in the returned iterator. |
| */ |
| override def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): |
| Iterator[UnsafeRowPair] |
| |
| /** Current metrics of the state store */ |
| def metrics: StateStoreMetrics |
| |
| /** |
| * Whether all updates have been committed |
| */ |
| def hasCommitted: Boolean |
| } |
| |
| /** Wraps the instance of StateStore to make the instance read-only. */ |
| class WrappedReadStateStore(store: StateStore) extends ReadStateStore { |
| override def id: StateStoreId = store.id |
| |
| override def version: Long = store.version |
| |
| override def get(key: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): UnsafeRow = store.get(key, |
| colFamilyName) |
| |
| override def iterator(colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): |
| Iterator[UnsafeRowPair] = store.iterator(colFamilyName) |
| |
| override def abort(): Unit = store.abort() |
| |
| override def prefixScan(prefixKey: UnsafeRow, |
| colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): Iterator[UnsafeRowPair] = |
| store.prefixScan(prefixKey, colFamilyName) |
| |
| override def valuesIterator(key: UnsafeRow, colFamilyName: String): Iterator[UnsafeRow] = { |
| store.valuesIterator(key, colFamilyName) |
| } |
| } |
| |
| /** |
| * Metrics reported by a state store |
| * @param numKeys Number of keys in the state store |
| * @param memoryUsedBytes Memory used by the state store |
| * @param customMetrics Custom implementation-specific metrics |
| * The metrics reported through this must have the same `name` as those |
| * reported by `StateStoreProvider.customMetrics`. |
| */ |
| case class StateStoreMetrics( |
| numKeys: Long, |
| memoryUsedBytes: Long, |
| customMetrics: Map[StateStoreCustomMetric, Long]) |
| |
| object StateStoreMetrics { |
| def combine(allMetrics: Seq[StateStoreMetrics]): StateStoreMetrics = { |
| val distinctCustomMetrics = allMetrics.flatMap(_.customMetrics.keys).distinct |
| val customMetrics = allMetrics.flatMap(_.customMetrics) |
| val combinedCustomMetrics = distinctCustomMetrics.map { customMetric => |
| val sameMetrics = customMetrics.filter(_._1 == customMetric) |
| val sumOfMetrics = sameMetrics.map(_._2).sum |
| customMetric -> sumOfMetrics |
| }.toMap |
| |
| StateStoreMetrics( |
| allMetrics.map(_.numKeys).sum, |
| allMetrics.map(_.memoryUsedBytes).sum, |
| combinedCustomMetrics) |
| } |
| } |
| |
| /** |
| * Name and description of custom implementation-specific metrics that a |
| * state store may wish to expose. Also provides [[SQLMetric]] instance to |
| * show the metric in UI and accumulate it at the query level. |
| */ |
| trait StateStoreCustomMetric { |
| def name: String |
| def desc: String |
| def withNewDesc(desc: String): StateStoreCustomMetric |
| def createSQLMetric(sparkContext: SparkContext): SQLMetric |
| } |
| |
| case class StateStoreCustomSumMetric(name: String, desc: String) extends StateStoreCustomMetric { |
| override def withNewDesc(newDesc: String): StateStoreCustomSumMetric = copy(desc = desc) |
| |
| override def createSQLMetric(sparkContext: SparkContext): SQLMetric = |
| SQLMetrics.createMetric(sparkContext, desc) |
| } |
| |
| case class StateStoreCustomSizeMetric(name: String, desc: String) extends StateStoreCustomMetric { |
| override def withNewDesc(desc: String): StateStoreCustomSizeMetric = copy(desc = desc) |
| |
| override def createSQLMetric(sparkContext: SparkContext): SQLMetric = |
| SQLMetrics.createSizeMetric(sparkContext, desc) |
| } |
| |
| case class StateStoreCustomTimingMetric(name: String, desc: String) extends StateStoreCustomMetric { |
| override def withNewDesc(desc: String): StateStoreCustomTimingMetric = copy(desc = desc) |
| |
| override def createSQLMetric(sparkContext: SparkContext): SQLMetric = |
| SQLMetrics.createTimingMetric(sparkContext, desc) |
| } |
| |
| /** |
| * An exception thrown when an invalid UnsafeRow is detected in state store. |
| */ |
| class InvalidUnsafeRowException(error: String) |
| extends RuntimeException("The streaming query failed by state format invalidation. " + |
| "The following reasons may cause this: 1. An old Spark version wrote the checkpoint that is " + |
| "incompatible with the current one; 2. Broken checkpoint files; 3. The query is changed " + |
| "among restart. For the first case, you can try to restart the application without " + |
| s"checkpoint or use the legacy Spark version to process the streaming state.\n$error", null) |
| |
| sealed trait KeyStateEncoderSpec |
| |
| case class NoPrefixKeyStateEncoderSpec(keySchema: StructType) extends KeyStateEncoderSpec |
| |
| case class PrefixKeyScanStateEncoderSpec( |
| keySchema: StructType, |
| numColsPrefixKey: Int) extends KeyStateEncoderSpec { |
| if (numColsPrefixKey == 0 || numColsPrefixKey >= keySchema.length) { |
| throw StateStoreErrors.incorrectNumOrderingColsForPrefixScan(numColsPrefixKey.toString) |
| } |
| } |
| |
| /** Encodes rows so that they can be range-scanned based on orderingOrdinals */ |
| case class RangeKeyScanStateEncoderSpec( |
| keySchema: StructType, |
| orderingOrdinals: Seq[Int]) extends KeyStateEncoderSpec { |
| if (orderingOrdinals.isEmpty || orderingOrdinals.length > keySchema.length) { |
| throw StateStoreErrors.incorrectNumOrderingColsForRangeScan(orderingOrdinals.length.toString) |
| } |
| } |
| |
| /** |
| * Trait representing a provider that provide [[StateStore]] instances representing |
| * versions of state data. |
| * |
| * The life cycle of a provider and its provide stores are as follows. |
| * |
| * - A StateStoreProvider is created in a executor for each unique [[StateStoreId]] when |
| * the first batch of a streaming query is executed on the executor. All subsequent batches reuse |
| * this provider instance until the query is stopped. |
| * |
| * - Every batch of streaming data request a specific version of the state data by invoking |
| * `getStore(version)` which returns an instance of [[StateStore]] through which the required |
| * version of the data can be accessed. It is the responsible of the provider to populate |
| * this store with context information like the schema of keys and values, etc. |
| * |
| * - After the streaming query is stopped, the created provider instances are lazily disposed off. |
| */ |
| trait StateStoreProvider { |
| |
| /** |
| * Initialize the provide with more contextual information from the SQL operator. |
| * This method will be called first after creating an instance of the StateStoreProvider by |
| * reflection. |
| * |
| * @param stateStoreId Id of the versioned StateStores that this provider will generate |
| * @param keySchema Schema of keys to be stored |
| * @param valueSchema Schema of value to be stored |
| * @param numColsPrefixKey The number of leftmost columns to be used as prefix key. |
| * A value not greater than 0 means the operator doesn't activate prefix |
| * key, and the operator should not call prefixScan method in StateStore. |
| * @param useColumnFamilies Whether the underlying state store uses a single or multiple column |
| * families |
| * @param storeConfs Configurations used by the StateStores |
| * @param hadoopConf Hadoop configuration that could be used by StateStore to save state data |
| * @param useMultipleValuesPerKey Whether the underlying state store needs to support multiple |
| * values for a single key. |
| */ |
| def init( |
| stateStoreId: StateStoreId, |
| keySchema: StructType, |
| valueSchema: StructType, |
| keyStateEncoderSpec: KeyStateEncoderSpec, |
| useColumnFamilies: Boolean, |
| storeConfs: StateStoreConf, |
| hadoopConf: Configuration, |
| useMultipleValuesPerKey: Boolean = false): Unit |
| |
| /** |
| * Return the id of the StateStores this provider will generate. |
| * Should be the same as the one passed in init(). |
| */ |
| def stateStoreId: StateStoreId |
| |
| /** Called when the provider instance is unloaded from the executor */ |
| def close(): Unit |
| |
| /** Return an instance of [[StateStore]] representing state data of the given version */ |
| def getStore(version: Long): StateStore |
| |
| /** |
| * Return an instance of [[ReadStateStore]] representing state data of the given version. |
| * By default it will return the same instance as getStore(version) but wrapped to prevent |
| * modification. Providers can override and return optimized version of [[ReadStateStore]] |
| * based on the fact the instance will be only used for reading. |
| */ |
| def getReadStore(version: Long): ReadStateStore = |
| new WrappedReadStateStore(getStore(version)) |
| |
| /** Optional method for providers to allow for background maintenance (e.g. compactions) */ |
| def doMaintenance(): Unit = { } |
| |
| /** |
| * Optional custom metrics that the implementation may want to report. |
| * @note The StateStore objects created by this provider must report the same custom metrics |
| * (specifically, same names) through `StateStore.metrics`. |
| */ |
| def supportedCustomMetrics: Seq[StateStoreCustomMetric] = Nil |
| } |
| |
| object StateStoreProvider { |
| |
| /** |
| * Return a instance of the given provider class name. The instance will not be initialized. |
| */ |
| def create(providerClassName: String): StateStoreProvider = { |
| val providerClass = Utils.classForName(providerClassName) |
| providerClass.getConstructor().newInstance().asInstanceOf[StateStoreProvider] |
| } |
| |
| /** |
| * Return a instance of the required provider, initialized with the given configurations. |
| */ |
| def createAndInit( |
| providerId: StateStoreProviderId, |
| keySchema: StructType, |
| valueSchema: StructType, |
| keyStateEncoderSpec: KeyStateEncoderSpec, |
| useColumnFamilies: Boolean, |
| storeConf: StateStoreConf, |
| hadoopConf: Configuration, |
| useMultipleValuesPerKey: Boolean): StateStoreProvider = { |
| val provider = create(storeConf.providerClass) |
| provider.init(providerId.storeId, keySchema, valueSchema, keyStateEncoderSpec, |
| useColumnFamilies, storeConf, hadoopConf, useMultipleValuesPerKey) |
| provider |
| } |
| |
| /** |
| * Use the expected schema to check whether the UnsafeRow is valid. |
| */ |
| def validateStateRowFormat( |
| keyRow: UnsafeRow, |
| keySchema: StructType, |
| valueRow: UnsafeRow, |
| valueSchema: StructType, |
| conf: StateStoreConf): Unit = { |
| if (conf.formatValidationEnabled) { |
| val validationError = UnsafeRowUtils.validateStructuralIntegrityWithReason(keyRow, keySchema) |
| validationError.foreach { error => throw new InvalidUnsafeRowException(error) } |
| |
| if (conf.formatValidationCheckValue) { |
| val validationError = |
| UnsafeRowUtils.validateStructuralIntegrityWithReason(valueRow, valueSchema) |
| validationError.foreach { error => throw new InvalidUnsafeRowException(error) } |
| } |
| } |
| } |
| } |
| |
| /** |
| * Unique identifier for a provider, used to identify when providers can be reused. |
| * Note that `queryRunId` is used uniquely identify a provider, so that the same provider |
| * instance is not reused across query restarts. |
| */ |
| case class StateStoreProviderId(storeId: StateStoreId, queryRunId: UUID) |
| |
| object StateStoreProviderId { |
| private[sql] def apply( |
| stateInfo: StatefulOperatorStateInfo, |
| partitionIndex: Int, |
| storeName: String): StateStoreProviderId = { |
| val storeId = StateStoreId( |
| stateInfo.checkpointLocation, stateInfo.operatorId, partitionIndex, storeName) |
| StateStoreProviderId(storeId, stateInfo.queryRunId) |
| } |
| } |
| |
| /** |
| * Unique identifier for a bunch of keyed state data. |
| * @param checkpointRootLocation Root directory where all the state data of a query is stored |
| * @param operatorId Unique id of a stateful operator |
| * @param partitionId Index of the partition of an operators state data |
| * @param storeName Optional, name of the store. Each partition can optionally use multiple state |
| * stores, but they have to be identified by distinct names. |
| */ |
| case class StateStoreId( |
| checkpointRootLocation: String, |
| operatorId: Long, |
| partitionId: Int, |
| storeName: String = StateStoreId.DEFAULT_STORE_NAME) { |
| |
| /** |
| * Checkpoint directory to be used by a single state store, identified uniquely by the tuple |
| * (operatorId, partitionId, storeName). All implementations of [[StateStoreProvider]] should |
| * use this path for saving state data, as this ensures that distinct stores will write to |
| * different locations. |
| */ |
| def storeCheckpointLocation(): Path = { |
| if (storeName == StateStoreId.DEFAULT_STORE_NAME) { |
| // For reading state store data that was generated before store names were used (Spark <= 2.2) |
| new Path(checkpointRootLocation, s"$operatorId/$partitionId") |
| } else { |
| new Path(checkpointRootLocation, s"$operatorId/$partitionId/$storeName") |
| } |
| } |
| |
| override def toString: String = { |
| s"""StateStoreId[ checkpointRootLocation=$checkpointRootLocation, operatorId=$operatorId, |
| | partitionId=$partitionId, storeName=$storeName ] |
| |""".stripMargin.replaceAll("\n", "") |
| } |
| } |
| |
| object StateStoreId { |
| val DEFAULT_STORE_NAME = "default" |
| } |
| |
| /** Mutable, and reusable class for representing a pair of UnsafeRows. */ |
| class UnsafeRowPair(var key: UnsafeRow = null, var value: UnsafeRow = null) { |
| def withRows(key: UnsafeRow, value: UnsafeRow): UnsafeRowPair = { |
| this.key = key |
| this.value = value |
| this |
| } |
| } |
| |
| |
| /** |
| * Companion object to [[StateStore]] that provides helper methods to create and retrieve stores |
| * by their unique ids. In addition, when a SparkContext is active (i.e. SparkEnv.get is not null), |
| * it also runs a periodic background task to do maintenance on the loaded stores. For each |
| * store, it uses the [[StateStoreCoordinator]] to ensure whether the current loaded instance of |
| * the store is the active instance. Accordingly, it either keeps it loaded and performs |
| * maintenance, or unloads the store. |
| */ |
| object StateStore extends Logging { |
| |
| val PARTITION_ID_TO_CHECK_SCHEMA = 0 |
| |
| val DEFAULT_COL_FAMILY_NAME = "default" |
| |
| @GuardedBy("loadedProviders") |
| private val loadedProviders = new mutable.HashMap[StateStoreProviderId, StateStoreProvider]() |
| |
| @GuardedBy("loadedProviders") |
| private val schemaValidated = new mutable.HashMap[StateStoreProviderId, Option[Throwable]]() |
| |
| private val maintenanceThreadPoolLock = new Object |
| |
| // Shared exception between threads in thread pool that the scheduling thread |
| // checks to see if an exception has been thrown in the maintenance task |
| private val threadPoolException = new AtomicReference[Throwable](null) |
| |
| // This set is to keep track of the partitions that are queued |
| // for maintenance or currently have maintenance running on them |
| // to prevent the same partition from being processed concurrently. |
| @GuardedBy("maintenanceThreadPoolLock") |
| private val maintenancePartitions = new mutable.HashSet[StateStoreProviderId] |
| |
| /** |
| * Runs the `task` periodically and automatically cancels it if there is an exception. `onError` |
| * will be called when an exception happens. |
| */ |
| class MaintenanceTask(periodMs: Long, task: => Unit, onError: => Unit) { |
| private val executor = |
| ThreadUtils.newDaemonSingleThreadScheduledExecutor("state-store-maintenance-task") |
| |
| private val runnable = new Runnable { |
| override def run(): Unit = { |
| try { |
| task |
| } catch { |
| case NonFatal(e) => |
| logWarning("Error running maintenance thread", e) |
| onError |
| throw e |
| } |
| } |
| } |
| |
| private val future: ScheduledFuture[_] = executor.scheduleAtFixedRate( |
| runnable, periodMs, periodMs, TimeUnit.MILLISECONDS) |
| |
| def stop(): Unit = { |
| future.cancel(false) |
| executor.shutdown() |
| } |
| |
| def isRunning: Boolean = !future.isDone |
| } |
| |
| /** |
| * Thread Pool that runs maintenance on partitions that are scheduled by |
| * MaintenanceTask periodically |
| */ |
| class MaintenanceThreadPool(numThreads: Int) { |
| private val threadPool = ThreadUtils.newDaemonFixedThreadPool( |
| numThreads, "state-store-maintenance-thread") |
| |
| def execute(runnable: Runnable): Unit = { |
| threadPool.execute(runnable) |
| } |
| |
| def stop(): Unit = { |
| threadPool.shutdown() |
| } |
| } |
| |
| @GuardedBy("loadedProviders") |
| private var maintenanceTask: MaintenanceTask = null |
| |
| @GuardedBy("loadedProviders") |
| private var maintenanceThreadPool: MaintenanceThreadPool = null |
| |
| @GuardedBy("loadedProviders") |
| private var _coordRef: StateStoreCoordinatorRef = null |
| |
| /** Get or create a read-only store associated with the id. */ |
| def getReadOnly( |
| storeProviderId: StateStoreProviderId, |
| keySchema: StructType, |
| valueSchema: StructType, |
| keyStateEncoderSpec: KeyStateEncoderSpec, |
| version: Long, |
| useColumnFamilies: Boolean, |
| storeConf: StateStoreConf, |
| hadoopConf: Configuration, |
| useMultipleValuesPerKey: Boolean = false): ReadStateStore = { |
| if (version < 0) { |
| throw QueryExecutionErrors.unexpectedStateStoreVersion(version) |
| } |
| val storeProvider = getStateStoreProvider(storeProviderId, keySchema, valueSchema, |
| keyStateEncoderSpec, useColumnFamilies, storeConf, hadoopConf, useMultipleValuesPerKey) |
| storeProvider.getReadStore(version) |
| } |
| |
| /** Get or create a store associated with the id. */ |
| def get( |
| storeProviderId: StateStoreProviderId, |
| keySchema: StructType, |
| valueSchema: StructType, |
| keyStateEncoderSpec: KeyStateEncoderSpec, |
| version: Long, |
| useColumnFamilies: Boolean, |
| storeConf: StateStoreConf, |
| hadoopConf: Configuration, |
| useMultipleValuesPerKey: Boolean = false): StateStore = { |
| if (version < 0) { |
| throw QueryExecutionErrors.unexpectedStateStoreVersion(version) |
| } |
| val storeProvider = getStateStoreProvider(storeProviderId, keySchema, valueSchema, |
| keyStateEncoderSpec, useColumnFamilies, storeConf, hadoopConf, useMultipleValuesPerKey) |
| storeProvider.getStore(version) |
| } |
| |
| private def disallowBinaryInequalityColumn(schema: StructType): Unit = { |
| if (!UnsafeRowUtils.isBinaryStable(schema)) { |
| throw new SparkUnsupportedOperationException( |
| errorClass = "STATE_STORE_UNSUPPORTED_OPERATION_BINARY_INEQUALITY", |
| messageParameters = Map("schema" -> schema.json) |
| ) |
| } |
| } |
| |
| private def getStateStoreProvider( |
| storeProviderId: StateStoreProviderId, |
| keySchema: StructType, |
| valueSchema: StructType, |
| keyStateEncoderSpec: KeyStateEncoderSpec, |
| useColumnFamilies: Boolean, |
| storeConf: StateStoreConf, |
| hadoopConf: Configuration, |
| useMultipleValuesPerKey: Boolean): StateStoreProvider = { |
| loadedProviders.synchronized { |
| startMaintenanceIfNeeded(storeConf) |
| |
| if (storeProviderId.storeId.partitionId == PARTITION_ID_TO_CHECK_SCHEMA) { |
| val result = schemaValidated.getOrElseUpdate(storeProviderId, { |
| // SPARK-47776: collation introduces the concept of binary (in)equality, which means |
| // in some collation we no longer be able to just compare the binary format of two |
| // UnsafeRows to determine equality. For example, 'aaa' and 'AAA' can be "semantically" |
| // same in case insensitive collation. |
| // State store is basically key-value storage, and the most provider implementations |
| // rely on the fact that all the columns in the key schema support binary equality. |
| // We need to disallow using binary inequality column in the key schema, before we |
| // could support this in majority of state store providers (or high-level of state |
| // store.) |
| disallowBinaryInequalityColumn(keySchema) |
| |
| val checker = new StateSchemaCompatibilityChecker(storeProviderId, hadoopConf) |
| // regardless of configuration, we check compatibility to at least write schema file |
| // if necessary |
| // if the format validation for value schema is disabled, we also disable the schema |
| // compatibility checker for value schema as well. |
| val ret = Try( |
| checker.check(keySchema, valueSchema, |
| ignoreValueSchema = !storeConf.formatValidationCheckValue) |
| ).toEither.fold(Some(_), _ => None) |
| if (storeConf.stateSchemaCheckEnabled) { |
| ret |
| } else { |
| None |
| } |
| }) |
| |
| if (result.isDefined) { |
| throw result.get |
| } |
| } |
| |
| // SPARK-42567 - Track load time for state store provider and log warning if takes longer |
| // than 2s. |
| val (provider, loadTimeMs) = Utils.timeTakenMs { |
| loadedProviders.getOrElseUpdate( |
| storeProviderId, |
| StateStoreProvider.createAndInit( |
| storeProviderId, keySchema, valueSchema, keyStateEncoderSpec, |
| useColumnFamilies, storeConf, hadoopConf, useMultipleValuesPerKey) |
| ) |
| } |
| |
| if (loadTimeMs > 2000L) { |
| logWarning(log"Loaded state store provider in loadTimeMs=${MDC(LOAD_TIME, loadTimeMs)} " + |
| log"for storeId=${MDC(STORE_ID, storeProviderId.storeId.toString)} and " + |
| log"queryRunId=${MDC(QUERY_RUN_ID, storeProviderId.queryRunId)}") |
| } |
| |
| val otherProviderIds = loadedProviders.keys.filter(_ != storeProviderId).toSeq |
| val providerIdsToUnload = reportActiveStoreInstance(storeProviderId, otherProviderIds) |
| providerIdsToUnload.foreach(unload(_)) |
| provider |
| } |
| } |
| |
| /** Unload a state store provider */ |
| def unload(storeProviderId: StateStoreProviderId): Unit = loadedProviders.synchronized { |
| loadedProviders.remove(storeProviderId).foreach(_.close()) |
| } |
| |
| /** Unload all state store providers: unit test purpose */ |
| private[sql] def unloadAll(): Unit = loadedProviders.synchronized { |
| loadedProviders.keySet.foreach { key => unload(key) } |
| loadedProviders.clear() |
| } |
| |
| /** Whether a state store provider is loaded or not */ |
| def isLoaded(storeProviderId: StateStoreProviderId): Boolean = loadedProviders.synchronized { |
| loadedProviders.contains(storeProviderId) |
| } |
| |
| /** Check if maintenance thread is running and scheduled future is not done */ |
| def isMaintenanceRunning: Boolean = loadedProviders.synchronized { |
| maintenanceTask != null && maintenanceTask.isRunning |
| } |
| |
| /** Stop maintenance thread and reset the maintenance task */ |
| def stopMaintenanceTask(): Unit = loadedProviders.synchronized { |
| if (maintenanceThreadPool != null) { |
| threadPoolException.set(null) |
| maintenanceThreadPoolLock.synchronized { |
| maintenancePartitions.clear() |
| } |
| maintenanceThreadPool.stop() |
| maintenanceThreadPool = null |
| } |
| if (maintenanceTask != null) { |
| maintenanceTask.stop() |
| maintenanceTask = null |
| } |
| } |
| |
| /** Unload and stop all state store providers */ |
| def stop(): Unit = loadedProviders.synchronized { |
| loadedProviders.keySet.foreach { key => unload(key) } |
| loadedProviders.clear() |
| _coordRef = null |
| stopMaintenanceTask() |
| logInfo("StateStore stopped") |
| } |
| |
| /** Start the periodic maintenance task if not already started and if Spark active */ |
| private def startMaintenanceIfNeeded(storeConf: StateStoreConf): Unit = { |
| val numMaintenanceThreads = storeConf.numStateStoreMaintenanceThreads |
| loadedProviders.synchronized { |
| if (SparkEnv.get != null && !isMaintenanceRunning) { |
| maintenanceTask = new MaintenanceTask( |
| storeConf.maintenanceInterval, |
| task = { doMaintenance() }, |
| onError = { loadedProviders.synchronized { |
| logInfo("Stopping maintenance task since an error was encountered.") |
| stopMaintenanceTask() |
| // SPARK-44504 - Unload explicitly to force closing underlying DB instance |
| // and releasing allocated resources, especially for RocksDBStateStoreProvider. |
| loadedProviders.keySet.foreach { key => unload(key) } |
| loadedProviders.clear() |
| } |
| } |
| ) |
| maintenanceThreadPool = new MaintenanceThreadPool(numMaintenanceThreads) |
| logInfo("State Store maintenance task started") |
| } |
| } |
| } |
| |
| private def processThisPartition(id: StateStoreProviderId): Boolean = { |
| maintenanceThreadPoolLock.synchronized { |
| if (!maintenancePartitions.contains(id)) { |
| maintenancePartitions.add(id) |
| true |
| } else { |
| false |
| } |
| } |
| } |
| |
| /** |
| * Execute background maintenance task in all the loaded store providers if they are still |
| * the active instances according to the coordinator. |
| */ |
| private def doMaintenance(): Unit = { |
| logDebug("Doing maintenance") |
| if (SparkEnv.get == null) { |
| throw new IllegalStateException("SparkEnv not active, cannot do maintenance on StateStores") |
| } |
| loadedProviders.synchronized { |
| loadedProviders.toSeq |
| }.foreach { case (id, provider) => |
| // check exception |
| if (threadPoolException.get() != null) { |
| val exception = threadPoolException.get() |
| logWarning("Error in maintenanceThreadPool", exception) |
| throw exception |
| } |
| if (processThisPartition(id)) { |
| maintenanceThreadPool.execute(() => { |
| val startTime = System.currentTimeMillis() |
| try { |
| provider.doMaintenance() |
| if (!verifyIfStoreInstanceActive(id)) { |
| unload(id) |
| logInfo(log"Unloaded ${MDC(STATE_STORE_PROVIDER, provider)}") |
| } |
| } catch { |
| case NonFatal(e) => |
| logWarning(log"Error managing ${MDC(STATE_STORE_PROVIDER, provider)}, " + |
| log"stopping management thread", e) |
| threadPoolException.set(e) |
| } finally { |
| val duration = System.currentTimeMillis() - startTime |
| val logMsg = |
| log"Finished maintenance task for provider=${MDC(STATE_STORE_PROVIDER, id)}" + |
| log" in elapsed_time=${MDC(TIME_UNITS, duration)}\n" |
| if (duration > 5000) { |
| logInfo(logMsg) |
| } else { |
| logDebug(logMsg) |
| } |
| maintenanceThreadPoolLock.synchronized { |
| maintenancePartitions.remove(id) |
| } |
| } |
| }) |
| } else { |
| logInfo(log"Not processing partition ${MDC(PARTITION_ID, id)} " + |
| log"for maintenance because it is currently " + |
| log"being processed") |
| } |
| } |
| } |
| |
| private def reportActiveStoreInstance( |
| storeProviderId: StateStoreProviderId, |
| otherProviderIds: Seq[StateStoreProviderId]): Seq[StateStoreProviderId] = { |
| if (SparkEnv.get != null) { |
| val host = SparkEnv.get.blockManager.blockManagerId.host |
| val executorId = SparkEnv.get.blockManager.blockManagerId.executorId |
| val providerIdsToUnload = coordinatorRef |
| .map(_.reportActiveInstance(storeProviderId, host, executorId, otherProviderIds)) |
| .getOrElse(Seq.empty[StateStoreProviderId]) |
| logInfo(log"Reported that the loaded instance " + |
| log"${MDC(STATE_STORE_PROVIDER, storeProviderId)} is active") |
| logDebug(log"The loaded instances are going to unload: " + |
| log"${MDC(STATE_STORE_PROVIDER, providerIdsToUnload.mkString(", "))}") |
| providerIdsToUnload |
| } else { |
| Seq.empty[StateStoreProviderId] |
| } |
| } |
| |
| private def verifyIfStoreInstanceActive(storeProviderId: StateStoreProviderId): Boolean = { |
| if (SparkEnv.get != null) { |
| val executorId = SparkEnv.get.blockManager.blockManagerId.executorId |
| val verified = |
| coordinatorRef.map(_.verifyIfInstanceActive(storeProviderId, executorId)).getOrElse(false) |
| logDebug(s"Verified whether the loaded instance $storeProviderId is active: $verified") |
| verified |
| } else { |
| false |
| } |
| } |
| |
| private def coordinatorRef: Option[StateStoreCoordinatorRef] = loadedProviders.synchronized { |
| val env = SparkEnv.get |
| if (env != null) { |
| val isDriver = |
| env.executorId == SparkContext.DRIVER_IDENTIFIER |
| // If running locally, then the coordinator reference in _coordRef may be have become inactive |
| // as SparkContext + SparkEnv may have been restarted. Hence, when running in driver, |
| // always recreate the reference. |
| if (isDriver || _coordRef == null) { |
| logDebug("Getting StateStoreCoordinatorRef") |
| _coordRef = StateStoreCoordinatorRef.forExecutor(env) |
| } |
| logInfo(log"Retrieved reference to StateStoreCoordinator: " + |
| log"${MDC(STATE_STORE_PROVIDER, _coordRef)}") |
| Some(_coordRef) |
| } else { |
| _coordRef = null |
| None |
| } |
| } |
| } |
| |
