anishshri-db commented on code in PR #45467:
URL: https://github.com/apache/spark/pull/45467#discussion_r1523670747
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala:
##########
@@ -64,31 +68,52 @@ case class TransformWithStateExec(
eventTimeWatermarkForLateEvents: Option[Long],
eventTimeWatermarkForEviction: Option[Long],
child: SparkPlan,
- isStreaming: Boolean = true)
- extends UnaryExecNode with StateStoreWriter with WatermarkSupport with
ObjectProducerExec {
+ isStreaming: Boolean = true,
+ hasInitialState: Boolean = false,
+ initialStateGroupingAttrs: Seq[Attribute],
+ initialStateDataAttrs: Seq[Attribute],
+ initialStateDeserializer: Expression,
+ initialState: SparkPlan)
+ extends BinaryExecNode with StateStoreWriter with WatermarkSupport with
ObjectProducerExec {
override def shortName: String = "transformWithStateExec"
// TODO: update this to run no-data batches when timer support is added
override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = false
- override protected def withNewChildInternal(
- newChild: SparkPlan): TransformWithStateExec = copy(child = newChild)
+ override def left: SparkPlan = child
+
+ override def right: SparkPlan = initialState
+
+ override protected def withNewChildrenInternal(
+ newLeft: SparkPlan, newRight: SparkPlan): TransformWithStateExec =
+ copy(child = newLeft, initialState = newRight)
override def keyExpressions: Seq[Attribute] = groupingAttributes
protected val schemaForKeyRow: StructType = new StructType().add("key",
BinaryType)
protected val schemaForValueRow: StructType = new StructType().add("value",
BinaryType)
+ /**
+ * Distribute by grouping attributes - We need the underlying data and the
initial state data
+ * to have the same grouping so that the data are co-lacated on the same
task.
+ */
override def requiredChildDistribution: Seq[Distribution] = {
- StatefulOperatorPartitioning.getCompatibleDistribution(groupingAttributes,
- getStateInfo, conf) ::
+ StatefulOperatorPartitioning.getCompatibleDistribution(
+ groupingAttributes, getStateInfo, conf) ::
+ StatefulOperatorPartitioning.getCompatibleDistribution(
Review Comment:
nit: indent ?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateExec.scala:
##########
@@ -111,6 +136,26 @@ case class TransformWithStateExec(
mappedIterator
}
+ private def processInitialStateRows(keyRow: UnsafeRow, initStateIter:
Iterator[InternalRow]):
+ Unit = {
Review Comment:
indent ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
