[SPARK-52195][PYTHON][SS] Fix initial state column dropping issue for Python TWS #50926
Conversation
| @@ -811,7 +811,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { | |||
| isStreaming = true, | |||
| hasInitialState, | |||
| planLater(initialState), | |||
| t.rightAttributes, | |||
| t.rightAttributes(), | |||
There was a problem hiding this comment.
why do we need this ?
There was a problem hiding this comment.
Because we have changed the rightAttributes to take parameters, even though the parameter has a default value, the compiler still requires parentheses
There was a problem hiding this comment.
Why do we pass includesInitialStateColumns as false here while we pass includesInitialStateColumns as true inside references?
There was a problem hiding this comment.
Here we need to pass initialStateGroupingAttrs as the input of TransformWithStateInPySparkExec which should not include other initial state columns. We only need to add these columns in references.
| assert(resolved, "This method is expected to be called after resolution.") | ||
| if (hasInitialState) { | ||
| right.output.take(initGroupingAttrsLen) | ||
| if (includesInitialStateColumns) { | ||
| // Include the initial state columns in the references to avoid being column pruned. |
There was a problem hiding this comment.
If I understand correctly for your PR descrption, the column pruning happens inside optimizer? Do you have a code pointer to where in the optimizer that the column get pruned?
There was a problem hiding this comment.
yeah it happened when Spark applies ColumnPruning rule. Since we didn't add these columns to references, ColumnPruning rule thinks these columns can be dropped.
|
cc @HeartSaVioR |
| @@ -215,10 +216,15 @@ case class TransformWithStateInPySpark( | |||
| left.output.take(groupingAttributesLen) | |||
| } | |||
|
|
|||
| def rightAttributes: Seq[Attribute] = { | |||
| def rightAttributes(includesInitialStateColumns: Boolean = false): Seq[Attribute] = { | |||
There was a problem hiding this comment.
Let's not make a single method to do two different purposes. Shall we have rightReferences to cover the new case?
|
@bogao007 Sorry, but could you please re-trigger the CI via empty commit, or just this module https://github.com/bogao007/spark/actions/runs/15123895213/job/42527399451 in the Github UI? I'd like to make sure any relevant modules aren't failing. Thanks! |
What changes were proposed in this pull request?
Fix initial state column dropping issue for Python TWS. This may occur when user adds extra transformations after
TransformWithStateInPandasoperator and those initial state columns will get pruned during optimization.Why are the changes needed?
This prevents users to use initial state with
TransformWithStateInPandasif they require extra transformations.Does this PR introduce any user-facing change?
No.
How was this patch tested?
Added unit test case.
Was this patch authored or co-authored using generative AI tooling?
No.