[SPARK-58205][SQL] Mark JSON/CSV/XML expressions stateful to prevent shared-evaluator data races#57354
[SPARK-58205][SQL] Mark JSON/CSV/XML expressions stateful to prevent shared-evaluator data races#57354vinodkc wants to merge 2 commits into
Conversation
128a44f to
0427728
Compare
HyukjinKwon
left a comment
HyukjinKwon
left a comment
There was a problem hiding this comment.
3 blocking, 0 non-blocking, 0 nits.
The stateful override is correct and effective for 7 of the 10 expressions. The 3 RuntimeReplaceable expressions and a missed sibling need another look.
Correctness (2)
- jsonExpressions.scala:524 / 593, csvExpressions.scala:191:
statefulis dead at execution for theRuntimeReplaceablesubset (StructsToJson,SchemaOfJson,SchemaOfCsv) — see inline - JsonExpressionsSuite.scala:1060: the new test passes but doesn't exercise the executed path for the
RuntimeReplaceableclaim — see inline
Design / architecture (1)
MultiGetJsonObject(jsonExpressions.scala:176, not in this diff) holds per-row mutable evaluator state but is not marked stateful — same pattern as the fix, on a node where it would be effective.MultiGetJsonObjectEvaluatorresets a sharedoutputBufferper call (JsonExpressionEvalUtils.scala:614,761) and holdsfallbackEvaluators, aSeqof mutableGetJsonObjectEvaluators that aresetJson-mutated (:610-620). It's created byOptimizeCsvJsonExprs(:231) and can appear duplicated in aProject. Consider addingoverride def stateful: Boolean = truehere too.
Verification
Traced how stateful reaches execution for each expression kind. For the 7 that own eval/nullSafeEval and are not RuntimeReplaceable, freshCopyIfContainsStatefulExpression (Expression.scala:176) copies the node at execution and its @transient lazy val evaluator re-initializes fresh per copy — the fix is effective. For the 3 RuntimeReplaceable expressions, ReplaceExpressions (Finish Analysis batch, Optimizer.scala:338) rewrites the node to Invoke/StaticInvoke(Literal(evaluator), ...) before any execution-time freshCopy; RuntimeReplaceable.eval asserts input == null (Expression.scala:462), so the node is never row-evaluated in place and its stateful override no longer exists at execution. After replacement the evaluator is pinned inside a Literal (a non-stateful LeafExpression), and freshCopy on a Literal returns this — so both copies keep the same evaluator. SchemaOfJson/SchemaOfCsv additionally require a foldable child and hold only config lazy-vals (no per-row var), so they are constant-folded and can't race — the overrides are unnecessary and inconsistent with the unmarked SchemaOfXml.
PR description suggestions
- Document: the description lists the ten expressions but does not distinguish the
RuntimeReplaceablesubset (StructsToJson,SchemaOfJson,SchemaOfCsv), where the override behaves differently — worth noting the mechanism/scope and justifying or dropping them.
| override protected def withNewChildInternal(newChild: Expression): StructsToJson = | ||
| copy(child = newChild) | ||
|
|
||
| override def stateful: Boolean = true |
There was a problem hiding this comment.
StructsToJson is RuntimeReplaceable, so this override is dead at execution. ReplaceExpressions (Finish Analysis batch, Optimizer.scala:338) rewrites this node to Invoke(Literal(evaluator), ...) before any execution-time freshCopyIfContainsStatefulExpression, and RuntimeReplaceable.eval asserts input == null (Expression.scala:462) so it's never row-evaluated in place. After replacement the evaluator is pinned inside a Literal (non-stateful LeafExpression), and freshCopy on a Literal returns this — both copies keep the same evaluator. So if a shared-node race is real for StructsToJson (its evaluator does hold mutable writer/gen and it accepts a non-foldable child), this override does not fix it; the sharing would have to be broken where the evaluator actually lives (inside replacement).
| } | ||
| } | ||
|
|
||
| override def stateful: Boolean = true |
There was a problem hiding this comment.
SchemaOfJson is RuntimeReplaceable and requires a foldable child (NON_FOLDABLE_INPUT check above), and SchemaOfJsonEvaluator holds only config lazy-vals (no per-row var). A foldable input is constant-folded via a single eval, so no multi-eval race is possible — this override is unnecessary. It's also inconsistent with SchemaOfXml (same RuntimeReplaceable + config-only-evaluator shape), which is correctly left unmarked. Suggest dropping it (and the matching one on SchemaOfCsv).
| override protected def withNewChildInternal(newChild: Expression): SchemaOfCsv = | ||
| copy(child = newChild) | ||
|
|
||
| override def stateful: Boolean = true |
There was a problem hiding this comment.
Same as SchemaOfJson: SchemaOfCsv is RuntimeReplaceable, requires a foldable child, and SchemaOfCsvEvaluator holds only config lazy-vals — constant-folded, so no race is possible and the override is unnecessary/dead at execution. Inconsistent with the unmarked SchemaOfXml. Suggest dropping it.
| assert(jsonTuple.freshCopyIfContainsStatefulExpression() ne jsonTuple) | ||
| } | ||
|
|
||
| test("StructsToJson and SchemaOfJson are stateful and produce fresh copies with " + |
There was a problem hiding this comment.
This test asserts copy1.replacement...evaluator ne copy2.replacement...evaluator on the pre-replacement node — each freshCopy trivially has its own @transient lazy val evaluator, so this passes regardless. But it doesn't reflect the executed plan: ReplaceExpressions materializes the Invoke once and the evaluator is shared through the resulting Literal (see the StructsToJson comment). So for the RuntimeReplaceable expressions this assertion is green without exercising the runtime sharing it's meant to guard. If the intent is to prove StructsToJson no longer races, the test needs to go through replacement (e.g. run ReplaceExpressions / an end-to-end duplicated-projection query) rather than inspecting .replacement on the un-replaced node.
2 participants
What changes were proposed in this pull request?
Adds
override def stateful: Boolean = trueto ten JSON/CSV/XML expressions:GetJsonObject,JsonTuple,JsonToStructs,StructsToJson,SchemaOfJson,CsvToStructs,StructsToCsv,SchemaOfCsv,XmlToStructs, andStructsToXml.Why are the changes needed?
These expressions hold mutable state in
@transient lazy val evaluatorfields but did not overridestateful(which defaults tofalse). As a result,freshCopyIfContainsStatefulExpression()never made fresh copies of them. When the same expression node appears more than once in a query plan — e.g.df.select(from_json(col, schema), from_json(col, schema))— both references share the same underlying evaluator instance. In interpreted execution this causes data-race corruption and silently incorrect results.Does this PR introduce any user-facing change?
No. Queries that were previously correct continue to produce the same results. Queries that previously silently produced incorrect results due to evaluator sharing now produce correct results.
How was this patch tested?
Added one unit test per suite (
JsonExpressionsSuite,CsvExpressionsSuite,XmlExpressionsSuite)Was this patch authored or co-authored using generative AI tooling?
No