You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by as...@apache.org on 2026/01/09 20:09:22 UTC
(spark) branch master updated: [SPARK-54443][SS] Integrate PartitionKeyExtractor in Re-partition reader
This is an automated email from the ASF dual-hosted git repository.
ashrigondekar pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new d2b00ae84341 [SPARK-54443][SS] Integrate PartitionKeyExtractor in Re-partition reader
d2b00ae84341 is described below
commit d2b00ae843416ba2fd9c70fe0da9e041e9e04658
Author: zifeif2 <zi...@gmail.com>
AuthorDate: Fri Jan 9 12:09:08 2026 -0800
[SPARK-54443][SS] Integrate PartitionKeyExtractor in Re-partition reader
### What changes were proposed in this pull request?
Integrate the PartitionKeyExtractor introduced in [this PR](https://github.com/apache/spark/pull/53355/files) to StatePartitionAllColumnFamiliesReader. Before this change, StatePartitionAllColumnFamiliesReader returns the entire key value in partition_key field, and SchemaUtil will return `keySchema` as the partitionKey schema. After this change, StatePartitionAllColumnFamiliesReader will return the actual partition key, and SchemaUtil returns the actual partitionKey schema
### Why are the changes needed?
When creating a StatePartitionAllColumnFamiliesReader, we need to pass along the stateFormatVersion and operator name.
In SchemaUtil, we will create a `getExtractor` helper function. It's used when getSourceSchema is called (for default column family), as well as when StatePartitionAllColumnFamiliesReader is initialized, as the reader will use the extractor to get partitionKey for different column families in `iterator`
We also added checks of partitionKey in reader suite
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
See updated StatePartitionAllColumnFamiliesSuite. We have hard-coded function that extract partition key for different column families from normalDF, then we'll compare the extracted partition key against the partition key read from bytesDF
### Was this patch authored or co-authored using generative AI tooling?
Yes. claude-4.5-opus
Closes #53459 from zifeif2/integrate-key-extraction.
Authored-by: zifeif2 <zi...@gmail.com>
Signed-off-by: Anish Shrigondekar <an...@databricks.com>
---
.../src/main/resources/error/error-conditions.json | 7 ++
.../datasources/v2/state/StateDataSource.scala | 97 ++++++++++++------
.../v2/state/StatePartitionReader.scala | 71 ++++++++++---
.../datasources/v2/state/utils/SchemaUtil.scala | 48 +++++----
.../StateStoreColumnFamilySchemaUtils.scala | 59 ++++++++++-
.../TransformWithStateVariableUtils.scala | 5 +
.../transformwithstate/timers/TimerStateImpl.scala | 6 ++
.../streaming/state/StateStoreErrors.scala | 10 ++
...tatePartitionAllColumnFamiliesReaderSuite.scala | 112 ++++++++++++++-------
9 files changed, 314 insertions(+), 101 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
index 1979fbffe361..4a0bfa704b2d 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -5771,6 +5771,13 @@
],
"sqlState" : "XXKST"
},
+ "STATE_STORE_UNKNOWN_INTERNAL_COLUMN_FAMILY" : {
+ "message" : [
+ "Unknown internal column family: <colFamilyName>.",
+ "This internal column family is not recognized by the StateStoreColumnFamilySchemaUtils."
+ ],
+ "sqlState" : "XXKST"
+ },
"STATE_STORE_UNSUPPORTED_OPERATION" : {
"message" : [
"<operationType> operation not supported with <entity>"
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala
index 897bb8339197..40fba5a90cbd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala
@@ -44,6 +44,7 @@ import org.apache.spark.sql.execution.streaming.runtime.StreamingQueryCheckpoint
import org.apache.spark.sql.execution.streaming.state.{InMemoryStateSchemaProvider, KeyStateEncoderSpec, NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RocksDBStateStoreProvider, StateSchemaCompatibilityChecker, StateSchemaMetadata, StateSchemaProvider, StateStore, StateStoreColFamilySchema, StateStoreConf, StateStoreId, StateStoreProviderId}
import org.apache.spark.sql.execution.streaming.state.OfflineStateRepartitionErrors
import org.apache.spark.sql.execution.streaming.utils.StreamingUtils
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.sources.DataSourceRegister
import org.apache.spark.sql.streaming.TimeMode
import org.apache.spark.sql.types.StructType
@@ -75,9 +76,8 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
sourceOptions.resolvedCpLocation,
stateConf.providerClass)
}
- val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
- sourceOptions)
+ val stateStoreReaderInfo = getStoreMetadataAndRunChecks(sourceOptions)
// The key state encoder spec should be available for all operators except stream-stream joins
val keyStateEncoderSpec = if (stateStoreReaderInfo.keyStateEncoderSpecOpt.isDefined) {
stateStoreReaderInfo.keyStateEncoderSpecOpt.get
@@ -91,15 +91,14 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
stateStoreReaderInfo.stateStoreColFamilySchemaOpt,
stateStoreReaderInfo.stateSchemaProviderOpt,
stateStoreReaderInfo.joinColFamilyOpt,
- Option(stateStoreReaderInfo.allColumnFamiliesReaderInfo))
+ stateStoreReaderInfo.allColumnFamiliesReaderInfo)
}
override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
val sourceOptions = StateSourceOptions.modifySourceOptions(hadoopConf,
StateSourceOptions.apply(session, hadoopConf, options))
- val stateStoreReaderInfo: StateStoreReaderInfo = getStoreMetadataAndRunChecks(
- sourceOptions)
+ val stateStoreReaderInfo = getStoreMetadataAndRunChecks(sourceOptions)
val oldSchemaFilePaths = StateDataSource.getOldSchemaFilePaths(sourceOptions, hadoopConf)
val stateCheckpointLocation = sourceOptions.stateCheckpointLocation
@@ -123,7 +122,9 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
SchemaUtil.getSourceSchema(sourceOptions, keySchema,
valueSchema,
stateStoreReaderInfo.transformWithStateVariableInfoOpt,
- stateStoreReaderInfo.stateStoreColFamilySchemaOpt)
+ stateStoreReaderInfo.stateStoreColFamilySchemaOpt,
+ stateStoreReaderInfo.allColumnFamiliesReaderInfo.map(_.operatorName),
+ stateStoreReaderInfo.allColumnFamiliesReaderInfo.flatMap(_.stateFormatVersion))
} catch {
case NonFatal(e) =>
throw StateDataSourceErrors.failedToReadStateSchema(sourceOptions, e)
@@ -132,6 +133,29 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
override def supportsExternalMetadata(): Boolean = false
+ /**
+ * Return the state format version for SYMMETRIC_HASH_JOIN operators,
+ * otherwise None for non-join operators.
+ * This currently only supports join operators because this function is only used to
+ * create a PartitionKeyExtractor through PartitionKeyExtractorFactory where only join operators
+ * require state format version
+ */
+ private def getStateFormatVersion(
+ storeMetadata: Array[StateMetadataTableEntry],
+ checkpointLocation: String,
+ batchId: Long): Option[Int] = {
+ if (storeMetadata.nonEmpty &&
+ storeMetadata.head.operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME) {
+ new StreamingQueryCheckpointMetadata(session, checkpointLocation).offsetLog
+ .get(batchId)
+ .flatMap(_.metadataOpt)
+ .flatMap(_.conf.get(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION.key))
+ .map(_.toInt)
+ } else {
+ None
+ }
+ }
+
/**
* Returns true if this is a read-all-column-families request for a stream-stream join
* that uses virtual column families (state format version 3).
@@ -260,8 +284,8 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
}
}
- private def getStoreMetadataAndRunChecks(sourceOptions: StateSourceOptions):
- StateStoreReaderInfo = {
+ private def getStoreMetadataAndRunChecks(
+ sourceOptions: StateSourceOptions): StateStoreReaderInfo = {
val storeMetadata = StateDataSource.getStateStoreMetadata(sourceOptions, hadoopConf)
if (!sourceOptions.internalOnlyReadAllColumnFamilies) {
// Skip runStateVarChecks when reading all column families (for repartitioning) because:
@@ -297,28 +321,31 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
if (sourceOptions.readRegisteredTimers) {
stateVarName = TimerStateUtils.getTimerStateVarNames(timeMode)._1
}
- // When reading all column families (for repartitioning), we collect all state variable
- // infos instead of validating a specific stateVarName. This skips the normal validation
- // logic because we're not reading a specific state variable - we're reading all of them.
if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ // When reading all column families (for repartitioning) for TWS operator,
+ // we will just choose a random state as placeholder for default column family,
+ // because we need to use matching stateVariableInfo and stateStoreColFamilySchemaOpt
+ // to inferSchema (partitionKey in particular) later
+ stateVarName = operatorProperties.stateVariables.head.stateName
stateVariableInfos = operatorProperties.stateVariables
- } else {
- var stateVarInfoList = operatorProperties.stateVariables
- .filter(stateVar => stateVar.stateName == stateVarName)
- if (stateVarInfoList.isEmpty &&
+ }
+
+ var stateVarInfoList = operatorProperties.stateVariables
+ .filter(stateVar => stateVar.stateName == stateVarName)
+ if (!TimerStateUtils.isTimerCFName(stateVarName) &&
StateStoreColumnFamilySchemaUtils.isTestingInternalColFamily(stateVarName)) {
- // pass this dummy TWSStateVariableInfo for TWS internal column family during testing,
- // because internalColumns are not register in operatorProperties.stateVariables,
- // thus stateVarInfoList will be empty.
- stateVarInfoList = List(TransformWithStateVariableInfo(
- stateVarName, StateVariableType.ValueState, false
- ))
- }
- require(stateVarInfoList.size == 1, s"Failed to find unique state variable info " +
- s"for state variable $stateVarName in operator ${sourceOptions.operatorId}")
- val stateVarInfo = stateVarInfoList.head
- transformWithStateVariableInfoOpt = Some(stateVarInfo)
+ // pass this dummy TWSStateVariableInfo for TWS internal column family during testing,
+ // because internal column families are not registered in
+ // operatorProperties.stateVariables, thus stateVarInfoList will be empty.
+ stateVarInfoList = List(TransformWithStateVariableInfo(
+ stateVarName, StateVariableType.ValueState, false
+ ))
}
+ require(stateVarInfoList.size == 1, s"Failed to find unique state variable info " +
+ s"for state variable $stateVarName in operator ${sourceOptions.operatorId}")
+ val stateVarInfo = stateVarInfoList.head
+ transformWithStateVariableInfoOpt = Some(stateVarInfo)
+
val schemaFilePaths = storeMetadataEntry.stateSchemaFilePaths
val stateSchemaMetadata = StateSchemaMetadata.createStateSchemaMetadata(
sourceOptions.stateCheckpointLocation.toString,
@@ -374,13 +401,27 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging
}
}
+ val allColFamilyReaderInfoOpt: Option[AllColumnFamiliesReaderInfo] =
+ if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ assert(storeMetadata.nonEmpty, "storeMetadata shouldn't be empty")
+ val operatorName = storeMetadata.head.operatorName
+ val stateFormatVersion = getStateFormatVersion(
+ storeMetadata,
+ sourceOptions.resolvedCpLocation,
+ sourceOptions.batchId
+ )
+ Some(AllColumnFamiliesReaderInfo(
+ stateStoreColFamilySchemas, stateVariableInfos, operatorName, stateFormatVersion))
+ } else {
+ None
+ }
StateStoreReaderInfo(
keyStateEncoderSpecOpt,
stateStoreColFamilySchemaOpt,
transformWithStateVariableInfoOpt,
stateSchemaProvider,
joinColFamilyOpt,
- AllColumnFamiliesReaderInfo(stateStoreColFamilySchemas, stateVariableInfos)
+ allColFamilyReaderInfoOpt
)
}
@@ -769,7 +810,7 @@ case class StateStoreReaderInfo(
stateSchemaProviderOpt: Option[StateSchemaProvider],
joinColFamilyOpt: Option[String], // Only used for join op with state format v3
// List of all column family schemas - used when internalOnlyReadAllColumnFamilies=true
- allColumnFamiliesReaderInfo: AllColumnFamiliesReaderInfo
+ allColumnFamiliesReaderInfo: Option[AllColumnFamiliesReaderInfo]
)
object StateDataSource {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala
index 40c6999f962f..f55822af4ada 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala
@@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow}
import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory}
import org.apache.spark.sql.execution.datasources.v2.state.utils.SchemaUtil
+import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorsUtils, StatePartitionKeyExtractorFactory}
import org.apache.spark.sql.execution.streaming.operators.stateful.join.SymmetricHashJoinStateManager
import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.{StateStoreColumnFamilySchemaUtils, StateVariableType, TransformWithStateVariableInfo}
import org.apache.spark.sql.execution.streaming.state._
@@ -38,7 +39,9 @@ import org.apache.spark.util.{NextIterator, SerializableConfiguration}
*/
case class AllColumnFamiliesReaderInfo(
colFamilySchemas: Set[StateStoreColFamilySchema] = Set.empty,
- stateVariableInfos: List[TransformWithStateVariableInfo] = List.empty)
+ stateVariableInfos: List[TransformWithStateVariableInfo] = List.empty,
+ operatorName: String,
+ stateFormatVersion: Option[Int] = None)
/**
* An implementation of [[PartitionReaderFactory]] for State data source. This is used to support
@@ -285,6 +288,38 @@ class StatePartitionAllColumnFamiliesReader(
private val stateStoreColFamilySchemas = allColumnFamiliesReaderInfo.colFamilySchemas
private val stateVariableInfos = allColumnFamiliesReaderInfo.stateVariableInfos
+ private val operatorName = allColumnFamiliesReaderInfo.operatorName
+ private val stateFormatVersion = allColumnFamiliesReaderInfo.stateFormatVersion
+
+ private def isTWSOperator(operatorName: String): Boolean = {
+ StatefulOperatorsUtils.TRANSFORM_WITH_STATE_OP_NAMES.contains(operatorName)
+ }
+
+ private def isDefaultColFamilyInTWS(operatorName: String, colFamilyName: String): Boolean = {
+ isTWSOperator(operatorName) && colFamilyName == StateStore.DEFAULT_COL_FAMILY_NAME
+ }
+
+ // Create extractors for each column family - each column family may have different key schema
+ private lazy val cfPartitionKeyExtractors: Map[String, StatePartitionKeyExtractor] = {
+
+ stateStoreColFamilySchemas
+ // Filter out default column family for TWS operators because they are not in use
+ // and will not have a key extractor
+ .filter(schema => !isDefaultColFamilyInTWS(operatorName, schema.colFamilyName))
+ .map { cfSchema =>
+ val stateVariableInfoOpt = stateVariableInfos.find(
+ _.stateName == StateStoreColumnFamilySchemaUtils.getBaseStateName(
+ cfSchema.colFamilyName))
+ val extractor = StatePartitionKeyExtractorFactory.create(
+ operatorName,
+ cfSchema.keySchema,
+ partition.sourceOptions.storeName,
+ cfSchema.colFamilyName,
+ stateFormatVersion,
+ stateVariableInfoOpt)
+ cfSchema.colFamilyName -> extractor
+ }.toMap
+ }
private def isListType(colFamilyName: String): Boolean = {
SchemaUtil.checkVariableType(
@@ -364,21 +399,27 @@ class StatePartitionAllColumnFamiliesReader(
override lazy val iter: Iterator[InternalRow] = {
// Iterate all column families and concatenate results
- stateStoreColFamilySchemas.iterator.flatMap { cfSchema =>
- if (isListType(cfSchema.colFamilyName)) {
- store.iterator(cfSchema.colFamilyName).flatMap(
- pair =>
- store.valuesIterator(pair.key, cfSchema.colFamilyName).map {
- value =>
- SchemaUtil.unifyStateRowPairAsRawBytes((pair.key, value), cfSchema.colFamilyName)
- }
- )
- } else {
- store.iterator(cfSchema.colFamilyName).map { pair =>
- SchemaUtil.unifyStateRowPairAsRawBytes(
- (pair.key, pair.value), cfSchema.colFamilyName)
+ stateStoreColFamilySchemas.iterator
+ // Filter out default column family for TWS operators because they are not in use
+ // and will not have data
+ .filter(schema => !isDefaultColFamilyInTWS(operatorName, schema.colFamilyName))
+ .flatMap { cfSchema =>
+ val extractor = cfPartitionKeyExtractors(cfSchema.colFamilyName)
+ if (isListType(cfSchema.colFamilyName)) {
+ store.iterator(cfSchema.colFamilyName).flatMap(
+ pair =>
+ store.valuesIterator(pair.key, cfSchema.colFamilyName).map {
+ value =>
+ SchemaUtil.unifyStateRowPairAsRawBytes(
+ (pair.key, value), cfSchema.colFamilyName, extractor)
+ }
+ )
+ } else {
+ store.iterator(cfSchema.colFamilyName).map { pair =>
+ SchemaUtil.unifyStateRowPairAsRawBytes(
+ (pair.key, pair.value), cfSchema.colFamilyName, extractor)
+ }
}
- }
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala
index 44e032f5163a..1210f0605179 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/utils/SchemaUtil.scala
@@ -25,9 +25,10 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeRow}
import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData}
import org.apache.spark.sql.execution.datasources.v2.state.{StateDataSourceErrors, StateSourceOptions}
+import org.apache.spark.sql.execution.streaming.operators.stateful.StatePartitionKeyExtractorFactory
import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.{StateVariableType, TransformWithStateVariableInfo}
import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.StateVariableType._
-import org.apache.spark.sql.execution.streaming.state.{ReadStateStore, StateStoreColFamilySchema, UnsafeRowPair}
+import org.apache.spark.sql.execution.streaming.state.{ReadStateStore, StatePartitionKeyExtractor, StateStoreColFamilySchema, UnsafeRowPair}
import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, IntegerType, LongType, MapType, StringType, StructType}
import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.util.ArrayImplicits._
@@ -49,8 +50,26 @@ object SchemaUtil {
keySchema: StructType,
valueSchema: StructType,
transformWithStateVariableInfoOpt: Option[TransformWithStateVariableInfo],
- stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema]): StructType = {
- if (transformWithStateVariableInfoOpt.isDefined) {
+ stateStoreColFamilySchemaOpt: Option[StateStoreColFamilySchema],
+ operatorName: Option[String],
+ stateFormatVersion: Option[Int] = None): StructType = {
+ if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ require(stateStoreColFamilySchemaOpt.isDefined)
+ require(operatorName.isDefined)
+ val colFamilyName: String = stateStoreColFamilySchemaOpt.get.colFamilyName
+ val extractor = StatePartitionKeyExtractorFactory.create(
+ operatorName.get,
+ keySchema,
+ sourceOptions.storeName,
+ colFamilyName,
+ stateFormatVersion,
+ transformWithStateVariableInfoOpt)
+ new StructType()
+ .add("partition_key", extractor.partitionKeySchema)
+ .add("key_bytes", BinaryType)
+ .add("value_bytes", BinaryType)
+ .add("column_family_name", StringType)
+ } else if (transformWithStateVariableInfoOpt.isDefined) {
require(stateStoreColFamilySchemaOpt.isDefined)
generateSchemaForStateVar(transformWithStateVariableInfoOpt.get,
stateStoreColFamilySchemaOpt.get, sourceOptions)
@@ -61,14 +80,6 @@ object SchemaUtil {
.add("key", keySchema)
.add("value", valueSchema)
.add("partition_id", IntegerType)
- } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
- new StructType()
- // TODO [SPARK-54443]: change keySchema to a more specific type after we
- // can extract partition key from keySchema
- .add("partition_key", keySchema)
- .add("key_bytes", BinaryType)
- .add("value_bytes", BinaryType)
- .add("column_family_name", StringType)
} else {
new StructType()
.add("key", keySchema)
@@ -87,18 +98,17 @@ object SchemaUtil {
/**
* Returns an InternalRow representing
- * 1. partitionKey
+ * 1. partitionKey (extracted using the StatePartitionKeyExtractor)
* 2. key in bytes
* 3. value in bytes
* 4. column family name
*/
def unifyStateRowPairAsRawBytes(
pair: (UnsafeRow, UnsafeRow),
- colFamilyName: String): InternalRow = {
+ colFamilyName: String,
+ extractor: StatePartitionKeyExtractor): InternalRow = {
val row = new GenericInternalRow(4)
- // todo [SPARK-54443]: change keySchema to more specific type after we
- // can extract partition key from keySchema
- row.update(0, pair._1)
+ row.update(0, extractor.partitionKey(pair._1))
row.update(1, pair._1.getBytes)
row.update(2, pair._2.getBytes)
row.update(3, UTF8String.fromString(colFamilyName))
@@ -266,7 +276,9 @@ object SchemaUtil {
"value_bytes" -> classOf[BinaryType],
"column_family_name" -> classOf[StringType])
- val expectedFieldNames = if (transformWithStateVariableInfoOpt.isDefined) {
+ val expectedFieldNames = if (sourceOptions.internalOnlyReadAllColumnFamilies) {
+ Seq("partition_key", "key_bytes", "value_bytes", "column_family_name")
+ } else if (transformWithStateVariableInfoOpt.isDefined) {
val stateVarInfo = transformWithStateVariableInfoOpt.get
val stateVarType = stateVarInfo.stateVariableType
@@ -305,8 +317,6 @@ object SchemaUtil {
}
} else if (sourceOptions.readChangeFeed) {
Seq("batch_id", "change_type", "key", "value", "partition_id")
- } else if (sourceOptions.internalOnlyReadAllColumnFamilies) {
- Seq("partition_key", "key_bytes", "value_bytes", "column_family_name")
} else {
Seq("key", "value", "partition_id")
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateStoreColumnFamilySchemaUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateStoreColumnFamilySchemaUtils.scala
index bbcdb1ff5326..4271cafb2eb7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateStoreColumnFamilySchemaUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateStoreColumnFamilySchemaUtils.scala
@@ -21,8 +21,9 @@ import scala.collection.mutable
import org.apache.spark.sql.Encoder
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateKeyValueRowSchemaUtils._
-import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateVariableUtils.getRowCounterCFName
-import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema}
+import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateVariableUtils.{getRowCounterCFName, getStateNameFromRowCounterCFName, isRowCounterCFName}
+import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.timers.TimerStateUtils
+import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, PrefixKeyScanStateEncoderSpec, RangeKeyScanStateEncoderSpec, StateStoreColFamilySchema, StateStoreErrors}
import org.apache.spark.sql.types._
object StateStoreColumnFamilySchemaUtils {
@@ -114,6 +115,60 @@ object StateStoreColumnFamilySchemaUtils {
org.apache.spark.util.Utils.isTesting && isInternalColFamily(colFamilyName)
}
+ /**
+ * Extracts the base state variable name from internal column family names.
+ * Internal column families are auxiliary data structures (TTL index, min expiry index,
+ * count index, row counter) that are associated with a user-defined state variable.
+ *
+ * @param colFamilyName The internal column family name (must start with "$")
+ * @return The base state variable name
+ * @throws IllegalArgumentException if the column family name is not a recognized internal type
+ */
+ private def getStateNameForInternalCF(colFamilyName: String): String = {
+ if (isTtlColFamilyName(colFamilyName)) {
+ getStateNameFromTtlColFamily(colFamilyName)
+ } else if (isMinExpiryIndexCFName(colFamilyName)) {
+ getStateNameFromMinExpiryIndexCFName(colFamilyName)
+ } else if (isCountIndexCFName(colFamilyName)) {
+ getStateNameFromCountIndexCFName(colFamilyName)
+ } else if (isRowCounterCFName(colFamilyName)) {
+ getStateNameFromRowCounterCFName(colFamilyName)
+ } else if (TimerStateUtils.isTimerCFName(colFamilyName)) {
+ // Return the primary index for timer secondary index column family
+ // because we only store the primary index column family in the
+ // StateMetadataTableEntry.operatorProperies.stateVariables
+ if (TimerStateUtils.isTimerSecondaryIndexCF(colFamilyName)) {
+ TimerStateUtils.getPrimaryIndexFromSecondaryIndexCF(colFamilyName)
+ } else {
+ colFamilyName
+ }
+ } else {
+ throw StateStoreErrors.unknownInternalColumnFamily(colFamilyName)
+ }
+ }
+
+ /**
+ * Extracts the base state variable name from a column family name.
+ *
+ * This is useful for looking up the stateVariableInfo associated with a column family,
+ * since stateVariableInfo is only stored for primary/user-facing column families.
+ *
+ * Returns:
+ * - For internal CFs (e.g., $ttl_*, $min_*, $count_*): the associated state variable name
+ * - For timer secondary index CFs: the corresponding primary index timer CF name
+ * - For all other CFs (regular state variables, timer primary index): the name as-is
+ *
+ * @param colFamilyName The column family name
+ * @return The base state variable name for stateVariableInfo lookup
+ */
+ def getBaseStateName(colFamilyName: String): String = {
+ if (isInternalColFamily(colFamilyName)) {
+ getStateNameForInternalCF(colFamilyName)
+ } else {
+ colFamilyName
+ }
+ }
+
def getValueStateSchema[T](
stateName: String,
keyEncoder: ExpressionEncoder[Any],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateVariableUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateVariableUtils.scala
index 5b4100e9c256..64d73ce5f545 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateVariableUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateVariableUtils.scala
@@ -63,6 +63,11 @@ object TransformWithStateVariableUtils {
def isRowCounterCFName(colFamilyName: String): Boolean = {
colFamilyName.startsWith(ROW_COUNTER_CF_PREFIX)
}
+
+ def getStateNameFromRowCounterCFName(colFamilyName: String): String = {
+ require(isRowCounterCFName(colFamilyName))
+ colFamilyName.substring(ROW_COUNTER_CF_PREFIX.length)
+ }
}
// Enum of possible State Variable types
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala
index 75624c9af9ef..101265fd8d83 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala
@@ -64,6 +64,12 @@ object TimerStateUtils {
assert(isTimerCFName(colFamilyName), s"Column family name must be for a timer: $colFamilyName")
colFamilyName.endsWith(TIMESTAMP_TO_KEY_CF)
}
+
+ def getPrimaryIndexFromSecondaryIndexCF(colFamilyName: String): String = {
+ assert(isTimerSecondaryIndexCF(colFamilyName),
+ s"Column family name must be for a timer secondary index: $colFamilyName")
+ colFamilyName.replace(TIMESTAMP_TO_KEY_CF, KEY_TO_TIMESTAMP_CF)
+ }
}
/**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreErrors.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreErrors.scala
index cd29f8f30f6e..fc9c256735b2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreErrors.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreErrors.scala
@@ -291,6 +291,11 @@ object StateStoreErrors {
def stateStoreCheckpointIdsNotSupported(msg: String): StateStoreCheckpointIdsNotSupported = {
new StateStoreCheckpointIdsNotSupported(msg)
}
+
+ def unknownInternalColumnFamily(colFamilyName: String):
+ StateStoreUnknownInternalColumnFamily = {
+ new StateStoreUnknownInternalColumnFamily(colFamilyName)
+ }
}
trait ConvertableToCannotLoadStoreError {
@@ -637,3 +642,8 @@ class StateStoreAutoSnapshotRepairFailed(
"selectedSnapshots" -> selectedSnapshots.mkString(","),
"eligibleSnapshots" -> eligibleSnapshots.mkString(",")),
cause)
+
+class StateStoreUnknownInternalColumnFamily(colFamilyName: String)
+ extends SparkRuntimeException(
+ errorClass = "STATE_STORE_UNKNOWN_INTERNAL_COLUMN_FAMILY",
+ messageParameters = Map("colFamilyName" -> colFamilyName))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala
index cb864da48731..4eb571804935 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionAllColumnFamiliesReaderSuite.scala
@@ -88,13 +88,22 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
/**
* Compares normal read data with bytes read data for a specific column family.
* Converts normal rows to bytes then compares with bytes read.
+ *
+ * @param normalDf Normal read data with columns (partition_id, key, value)
+ * @param bytesDf Bytes read data with columns (partition_key, key_bytes, value_bytes, cf_name)
+ * @param columnFamily The column family name to filter on
+ * @param keySchema Schema of the full key
+ * @param valueSchema Schema of the value
+ * @param partitionKeyExtractor Function to extract partition key from full key Row.
+ * If None, assumes partition key equals the full key.
*/
private def compareNormalAndBytesData(
normalDf: Array[Row],
bytesDf: Array[Row],
columnFamily: String,
keySchema: StructType,
- valueSchema: StructType): Unit = {
+ valueSchema: StructType,
+ partitionKeyExtractor: Option[Row => Row] = None): Unit = {
// Filter bytes data for the specified column family and extract raw bytes directly
val filteredBytesData = bytesDf.filter { row =>
@@ -114,12 +123,17 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
val keyConverter = CatalystTypeConverters.createToCatalystConverter(keySchema)
val valueConverter = CatalystTypeConverters.createToCatalystConverter(valueSchema)
- // Convert normal data to bytes
- val normalAsBytes = normalDf.toSeq.map { row =>
+ // Convert normal data to (partitionKeyStruct, keyBytes, valueBytes)
+ val normalData = normalDf.toSeq.map { row =>
val key = row.getStruct(1)
val value = if (row.isNullAt(2)) null else row.getStruct(2)
- // Convert key to InternalRow, then to UnsafeRow, then get bytes
+ // Extract partition key - use extractor if provided, otherwise use full key
+ val partitionKey: Row = partitionKeyExtractor match {
+ case Some(extractor) => extractor(key)
+ case None => key
+ }
+ // Convert key to bytes
val keyInternalRow = keyConverter(key).asInstanceOf[InternalRow]
val keyUnsafeRow = keyProjection(keyInternalRow)
// IMPORTANT: Must clone the bytes array since getBytes() returns a reference
@@ -137,26 +151,29 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
valueUnsafeRow.getBytes.clone()
}
- (keyBytes, valueBytes)
+ (keyBytes, valueBytes, partitionKey)
}
- // Extract raw bytes from bytes read data (no deserialization/reserialization)
- val bytesAsBytes = filteredBytesData.map { row =>
+ // Extract (partitionKeyStruct, keyBytes, valueBytes) from bytes read data
+ val bytesData = filteredBytesData.map { row =>
+ val partitionKey = row.getStruct(0)
val keyBytes = row.getAs[Array[Byte]](1)
val valueBytes = row.getAs[Array[Byte]](2)
- (keyBytes, valueBytes)
+ (keyBytes, valueBytes, partitionKey)
}
- // Sort both for comparison (since Set equality doesn't work well with byte arrays)
- val normalSorted = normalAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
- val bytesSorted = bytesAsBytes.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+ // Sort both for comparison by key and value bytes
+ val normalSorted = normalData.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
+ val bytesSorted = bytesData.sortBy(x => (x._1.mkString(","), x._2.mkString(",")))
assert(normalSorted.length == bytesSorted.length,
s"Size mismatch: normal has ${normalSorted.length}, bytes has ${bytesSorted.length}")
- // Compare each pair
+ // Compare each tuple (partitionKeyStruct, keyBytes, valueBytes)
normalSorted.zip(bytesSorted).zipWithIndex.foreach {
- case (((normalKey, normalValue), (bytesKey, bytesValue)), idx) =>
+ case (((normalKey, normalValue, normalPartitionKey),
+ (bytesKey, bytesValue, bytesPartitionKey)), idx) =>
+ assert(normalPartitionKey == bytesPartitionKey)
assert(Arrays.equals(normalKey, bytesKey),
s"Key mismatch at index $idx:\n" +
s" Normal: ${normalKey.mkString("[", ",", "]")}\n" +
@@ -179,7 +196,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
keySchema: StructType,
valueSchema: StructType,
extraOptions: Map[String, String] = Map.empty,
- selectExprs: Seq[String] = Seq("partition_id", "key", "value")): Unit = {
+ selectExprs: Seq[String] = Seq("partition_id", "key", "value"),
+ partitionKeyExtractor: Option[Row => Row] = None): Unit = {
var reader = spark.read
.format("statestore")
.option(StateSourceOptions.PATH, checkpointDir)
@@ -194,7 +212,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
allBytesData,
stateVarName,
keySchema,
- valueSchema)
+ valueSchema,
+ partitionKeyExtractor)
}
/**
@@ -252,7 +271,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
allBytesData,
s"$$${timerPrefix}Timers_keyToTimestamp",
keyToTimestampSchema,
- dummySchema)
+ dummySchema,
+ partitionKeyExtractor = Some(compositeKey => compositeKey.getStruct(0)))
val timestampToKeyNormalDf = timerBaseDf.selectExpr(
TimerTestUtils.getTimerSelectExpressions(s"$$${timerPrefix}Timers_timestampToKey"): _*)
@@ -261,7 +281,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
allBytesData,
s"$$${timerPrefix}Timers_timestampToKey",
timestampToKeySchema,
- dummySchema)
+ dummySchema,
+ partitionKeyExtractor = Some(compositeKey => compositeKey.getStruct(1)))
}
/**
@@ -278,7 +299,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
tempDir: String,
keySchema: StructType,
valueSchema: StructType,
- storeName: Option[String] = None): Unit = {
+ storeName: Option[String] = None,
+ partitionKeyExtractor: Option[Row => Row] = None): Unit = {
val normalDf = getNormalReadDf(tempDir, storeName)
val bytesDf = getBytesReadDf(tempDir, storeName)
@@ -288,7 +310,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
bytesDf.collect(),
StateStore.DEFAULT_COL_FAMILY_NAME,
keySchema,
- valueSchema)
+ valueSchema,
+ partitionKeyExtractor)
}
/**
@@ -307,7 +330,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
colFamilyName: String,
sharedBytesDf: DataFrame,
keySchema: StructType,
- valueSchema: StructType): Unit = {
+ valueSchema: StructType,
+ partitionKeyExtractor: Option[Row => Row] = None): Unit = {
val normalDf = getNormalReadDf(tempDir, Option(colFamilyName))
compareNormalAndBytesData(
@@ -315,7 +339,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
sharedBytesDf.collect(),
colFamilyName,
keySchema,
- valueSchema)
+ valueSchema,
+ partitionKeyExtractor)
}
// Run all tests with both changelog checkpointing enabled and disabled
@@ -399,7 +424,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
val (keySchema, valueSchema) = SessionWindowTestUtils.getSchemas()
- validateStateStore(tempDir.getAbsolutePath, keySchema, valueSchema)
+ validateStateStore(tempDir.getAbsolutePath, keySchema, valueSchema,
+ partitionKeyExtractor = Some(row => Row(row.getString(0))))
}
}
@@ -414,13 +440,12 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
val normalData = getNormalReadDf(tempDir.getAbsolutePath).collect()
val bytesDf = getBytesReadDf(tempDir.getAbsolutePath)
- validateBytesReadDfSchema(bytesDf)
- compareNormalAndBytesData(
- normalData, bytesDf.collect(), "default", keySchema, valueSchema)
- }
+ validateBytesReadDfSchema(bytesDf)
+ compareNormalAndBytesData(
+ normalData, bytesDf.collect(), "default", keySchema, valueSchema)
}
}
- )
+ })
Seq(1, 2).foreach(version =>
testWithChangelogConfig(s"stream-stream join, state ver $version") {
@@ -448,7 +473,9 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
tempDir.getAbsolutePath,
keyWithIndexKeySchema,
keyWithIndexValueSchema,
- Some(storeName))
+ Some(storeName),
+ partitionKeyExtractor = Some(compositeKey =>
+ Row(compositeKey.getInt(0))))
}
}
}
@@ -514,7 +541,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
tempDir.getAbsolutePath, allBytesData,
MultiStateVarProcessorTestUtils.ITEMS_MAP, mapKeySchema, mapValueSchema,
selectExprs = MultiStateVarProcessorTestUtils.getSelectExpressions(
- MultiStateVarProcessorTestUtils.ITEMS_MAP))
+ MultiStateVarProcessorTestUtils.ITEMS_MAP),
+ partitionKeyExtractor = Some(compositeKey => compositeKey.getStruct(0)))
}
}
@@ -616,13 +644,18 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
val (ttlIndexKeySchema, ttlValueSchema) = schemas(TTLProcessorUtils.LIST_STATE_TTL_INDEX)
val (_, minExpiryValueSchema) = schemas(TTLProcessorUtils.LIST_STATE_MIN)
val (_, countValueSchema) = schemas(TTLProcessorUtils.LIST_STATE_COUNT)
+ val ttlColFamilyPartitionKeyExtractor: Option[Row => Row] =
+ Some(compositeKey => compositeKey.getStruct(1))
val columnFamilyAndKeyValueSchema = Seq(
- (TTLProcessorUtils.LIST_STATE_TTL_INDEX, ttlIndexKeySchema, ttlValueSchema),
- (TTLProcessorUtils.LIST_STATE_MIN, groupByKeySchema, minExpiryValueSchema),
- (TTLProcessorUtils.LIST_STATE_COUNT, groupByKeySchema, countValueSchema)
+ (TTLProcessorUtils.LIST_STATE_TTL_INDEX,
+ ttlIndexKeySchema, ttlValueSchema,
+ ttlColFamilyPartitionKeyExtractor),
+ (TTLProcessorUtils.LIST_STATE_MIN, groupByKeySchema, minExpiryValueSchema, None),
+ (TTLProcessorUtils.LIST_STATE_COUNT, groupByKeySchema, countValueSchema, None)
)
columnFamilyAndKeyValueSchema.foreach(pair => {
- validateColumnFamily(tempDir.getAbsolutePath, pair._1, bytesDf, pair._2, pair._3)
+ validateColumnFamily(
+ tempDir.getAbsolutePath, pair._1, bytesDf, pair._2, pair._3, pair._4)
})
}
}
@@ -665,12 +698,14 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
readAndValidateStateVar(
tempDir.getAbsolutePath, allBytesData,
stateVarName = TTLProcessorUtils.MAP_STATE, compositeKeySchema, mapStateValueSchema,
- selectExprs = TTLProcessorUtils.getTTLSelectExpressions(TTLProcessorUtils.MAP_STATE))
+ selectExprs = TTLProcessorUtils.getTTLSelectExpressions(TTLProcessorUtils.MAP_STATE),
+ partitionKeyExtractor = Some(compositeKey => compositeKey.getStruct(0)))
// Validate $ttl_mapState column family
readAndValidateStateVar(
tempDir.getAbsolutePath, allBytesData,
- stateVarName = TTLProcessorUtils.MAP_STATE_TTL_INDEX, ttlIndexKeySchema, dummyValueSchema)
+ stateVarName = TTLProcessorUtils.MAP_STATE_TTL_INDEX, ttlIndexKeySchema, dummyValueSchema,
+ partitionKeyExtractor = Some(ttlKey => ttlKey.getStruct(1).getStruct(0)))
}
}
@@ -728,7 +763,8 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
allBytesData,
TTLProcessorUtils.VALUE_STATE_TTL_INDEX,
ttlIndexKeySchema,
- dummyValueSchema)
+ dummyValueSchema,
+ partitionKeyExtractor = Some(ttlKey => ttlKey.getStruct(1)))
}
}
@@ -762,7 +798,9 @@ class StatePartitionAllColumnFamiliesReaderSuite extends StateDataSourceTestBase
colFamilyName,
stateBytesDf,
keyWithIndexKeySchema,
- keyWithIndexValueSchema)
+ keyWithIndexValueSchema,
+ partitionKeyExtractor = Some(compositeKey =>
+ Row(compositeKey.getInt(0))))
}
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org