Skip to content

[FEA] Support Project AST and AST JIT in cached broadcast hash join build-side projections #16509

Open

Description

@thirtiseven

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.

Activity

  1. thirtiseven commented on Oct 22, 2026

    @thirtiseven
    CollaboratorAuthor

    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.

  2. added
    performanceA performance related task/issue
    and removed
    ? - Needs TriageNeed team to review and classify
    on Oct 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    performanceA performance related task/issue

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions