Github user tdas commented on a diff in the pull request: https://github.com/apache/spark/pull/11645#discussion_r56386237 --- Diff: sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala --- @@ -0,0 +1,465 @@ +/* + * 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 scala.collection.mutable +import scala.util.Random +import scala.util.control.NonFatal + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{FileStatus, Path} + +import org.apache.spark.{Logging, SparkConf} +import org.apache.spark.serializer.{DeserializationStream, KryoSerializer, SerializationStream} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.JoinedRow +import org.apache.spark.util.{CompletionIterator, Utils} + + +/** + * An implementation of [[StateStoreProvider]] and [[StateStore]] in which all the data is backed + * by files in a HDFS-compatible file system. All updates to the store has to be done in sets + * transactionally, and each set of updates increments the store's version. These versions can + * be used to re-execute the updates (by retries in RDD operations) on the correct version of + * the store, and regenerate the store version. + * + * Usage: + * To update the data in the state store, the following order of operations are needed. + * + * - val store = StateStore.get(operatorId, partitionId, version) // to get the right store + * - store.update(...) + * - store.remove(...) + * - store.commit() // commits all the updates to made with version number + * - store.iterator() // key-value data after last commit as an iterator + * - store.updates() // updates made in the last as an iterator + * + * Fault-tolerance model: + * - Every set of updates is written to a delta file before committing. + * - The state store is responsible for managing, collapsing and cleaning up of delta files. + * - Multiple attempts to commit the same version of updates must have the same updates. + * - Background management of files ensures that last versions of the store is always recoverable + * to ensure re-executed RDD operations re-apply updates on the correct past version of the + * store. + */ +private[state] class HDFSBackedStateStoreProvider( + val id: StateStoreId, + val directory: String, + numBatchesToRetain: Int = 2, + maxDeltaChainForSnapshots: Int = 10 + ) extends StateStoreProvider with Logging { + type MapType = mutable.HashMap[InternalRow, InternalRow] + + import StateStore._ + + /** Implementation of [[StateStore]] API which is backed by a HDFS-compatible file system */ + class HDFSBackedStateStore( val version: Long, mapToUpdate: MapType) + extends StateStore { + + /** Trait and classes representing the internal state of the store */ + trait STATE + case object UPDATING extends STATE + case object COMMITTED extends STATE + case object CANCELLED extends STATE + + private val newVersion = version + 1 + private val tempDeltaFile = new Path(baseDir, s"temp-${Random.nextLong}") + private val tempDeltaFileStream = + serializer.newInstance().serializeStream(fs.create(tempDeltaFile, true)) + private val allUpdates = new mutable.HashMap[InternalRow, StoreUpdate] + + @volatile private var state: STATE = UPDATING + @volatile private var finalDeltaFile: Path = null + + override def id: StateStoreId = HDFSBackedStateStoreProvider.this.id + + /** Update the value of a key using the value generated by the update function */ + override def update(key: InternalRow, updateFunc: Option[InternalRow] => InternalRow): Unit = { + verify(state == UPDATING, "Cannot update after already committed or cancelled") + val oldValueOption = mapToUpdate.get(key) + val value = updateFunc(oldValueOption) + mapToUpdate.put(key, value) + allUpdates.get(key) match { + case Some(ValueAdded(_, _)) => + // Value did not exist in previous version and was added already, keep it marked as added + allUpdates.put(key, ValueAdded(key, value)) + case Some(ValueUpdated(_, _)) | Some(KeyRemoved(_)) => + // Value existed in prev version and updated/removed, mark it as updated + allUpdates.put(key, ValueUpdated(key, value)) + case None => + // There was no prior update, so mark this as added or updated according to its presence + // in previous version. + val update = + if (oldValueOption.nonEmpty) ValueUpdated(key, value) else ValueAdded(key, value) + allUpdates.put(key, update) + } + tempDeltaFileStream.writeObject(ValueUpdated(key, value)) + } + + /** Remove keys that match the following condition */ + override def remove(condition: InternalRow => Boolean): Unit = { + verify(state == UPDATING, "Cannot remove after already committed or cancelled") + val keyIter = mapToUpdate.keysIterator + while (keyIter.hasNext) { + val key = keyIter.next + if (condition(key)) { + mapToUpdate.remove(key) + + allUpdates.get(key) match { + case Some(ValueUpdated(_, _)) | None => + // Value existed in previous version and maybe was updated, mark removed + allUpdates.put(key, KeyRemoved(key)) + case Some(ValueAdded(_, _)) => + // Value did not exist in previous version and was added, should not appear in updates + allUpdates.remove(key) + case Some(KeyRemoved(_)) => + // Remove already in update map, no need to change + } + tempDeltaFileStream.writeObject(KeyRemoved(key)) + } + } + } + + /** Commit all the updates that have been made to the store. */ + override def commit(): Long = { + verify(state == UPDATING, "Cannot commit again after already committed or cancelled") + + try { + tempDeltaFileStream.close() + finalDeltaFile = commitUpdates(newVersion, mapToUpdate, tempDeltaFile) + state = COMMITTED + newVersion + } catch { + case NonFatal(e) => + throw new IllegalStateException( + s"Error committing version $newVersion into ${HDFSBackedStateStoreProvider.this}", e) + } + } + + /** Cancel all the updates made on this store. This store will not be usable any more. */ + override def cancel(): Unit = { + state = CANCELLED + if (tempDeltaFileStream != null) { + tempDeltaFileStream.close() + } + if (tempDeltaFile != null && fs.exists(tempDeltaFile)) { + fs.delete(tempDeltaFile, true) + } + } + + /** + * Get an iterator of all the store data. This can be called only after committing the + * updates. + */ + override def iterator(): Iterator[InternalRow] = { + verify(state == COMMITTED, "Cannot get iterator of store data before comitting") + HDFSBackedStateStoreProvider.this.iterator(newVersion) + } + + /** + * Get an iterator of all the updates made to the store in the current version. + * This can be called only after committing the updates. + */ + override def updates(): Iterator[StoreUpdate] = { + verify(state == COMMITTED, "Cannot get iterator of updates before committing") + allUpdates.valuesIterator + } + + /** + * Whether all updates have been committed + */ + override def hasCommitted: Boolean = { + state == COMMITTED + } + } + + /** Get the state store for making updates to create a new `version` of the store. */ + override def getStore(version: Long): StateStore = synchronized { + require(version >= 0, "Version cannot be less than 0") + val newMap = new MapType() + if (version > 0) { + newMap ++= loadMap(version) + } + new HDFSBackedStateStore(version, newMap) + } + + /** Manage backing files, including creating snapshots and cleaning up old files */ + override def manage(): Unit = { --- End diff -- No strong feelings either way, but this is kind of long. How about `doMaintenance`?
--- If your project is set up for it, you can reply to this email and have your reply appear on GitHub as well. If your project does not have this feature enabled and wishes so, or if the feature is enabled but not working, please contact infrastructure at infrastruct...@apache.org or file a JIRA ticket with INFRA. --- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org