You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[FEA] Support Project AST and AST JIT in cached broadcast hash join build-side projections #16509
Is your feature request related to a problem? Please describe.
When spark.rapids.sql.join.hashTable.reuse is enabled, a broadcast hash join keeps an executor-wide hash table for its build side. If Projects sit between the broadcast and the join, the hash-build cache keeps the bound build-side projection so it can rebuild the hash table after it is spilled. A rebuild can run in any later task.
Project AST and AST JIT evaluators compile native state lazily and release it when the task that compiled it completes. If that task completes while another task rebuilds with the same evaluators, the release can overlap CompiledExpression.computeColumn or AstJitProgram.computeTable, which cuDF does not support. #15312 avoids this by binding the cached projection for regular GPU evaluation only; per-task build-side projections still use AST and AST JIT.
Describe the solution you'd like
Give each cached build or rebuild invocation its own evaluators with explicit ownership:
Create fresh AST and AST JIT evaluators for each invocation, and reuse them across retries and tiers within that invocation.
Close them when the invocation finishes instead of registering task-completion cleanup.
Close the source CompiledExpressions once an AstJitProgram is built, since the program copies its inputs.
Measure program construction cost per rebuild with a warm native kernel cache.
Describe alternatives you've considered
Reference counting in the evaluators would avoid the use-after-free, but ownership would still depend on whichever task compiled the state first.
Additional context
Tests should cover a task that completes while another task rebuilds the cached hash table after a spill, and retry within one build or rebuild invocation. GpuBloomFilterMightContain uses the same lazily created, task-scoped native state, so the same ownership model may apply to other expressions retained by executor-wide caches.
GpuIn has the same lifetime pattern, independently of the Project AST switches. For an IN list with non-literal items that needs at least two comparisons, such as x IN (y, z) or x IN (1, y), GpuIn combines the comparisons with a GpuProjectAstExpression that it keeps in a @transient lazy val (orAst or fusedAst). That expression is not a child, so binding the cached build-side projection with both Project AST backends disabled (#15312) does not remove it. The first task that evaluates a cached projection containing such an IN compiles the AST and registers its release with its own completion, so the overlap described above can still happen.
The per-invocation ownership here should therefore also cover evaluators that expressions create internally, not only Project-level AST wrappers. A regression test could use a build-side Project containing x IN (y, z). This is from code inspection; I have not reproduced it, and it requires spark.rapids.sql.join.hashTable.reuse, which is disabled by default.
Activity
thirtiseven commentedon Oct 22, 2026
GpuInhas the same lifetime pattern, independently of the Project AST switches. For anINlist with non-literal items that needs at least two comparisons, such asx IN (y, z)orx IN (1, y),GpuIncombines the comparisons with aGpuProjectAstExpressionthat it keeps in a@transient lazy val(orAstorfusedAst). That expression is not a child, so binding the cached build-side projection with both Project AST backends disabled (#15312) does not remove it. The first task that evaluates a cached projection containing such anINcompiles the AST and registers its release with its own completion, so the overlap described above can still happen.The per-invocation ownership here should therefore also cover evaluators that expressions create internally, not only Project-level AST wrappers. A regression test could use a build-side Project containing
x IN (y, z). This is from code inspection; I have not reproduced it, and it requiresspark.rapids.sql.join.hashTable.reuse, which is disabled by default.Add experimental Project AST JIT for integral add and multiply [datab…