Repository navigation
[FEA] Support AST JIT projection #16385
Description
Is your feature request related to a problem? Please describe.
cuDF AST JIT can evaluate several expressions in one kernel and avoid materializing intermediate columns. #15312 adds the initial Project integration, limited to non-ANSI INT/LONG addition and multiplication and disabled by default.
We should extend expression coverage and address the cold run cost and backend selection needed to make this useful for more queries.
Describe the solution you'd like
Use this epic to track expression coverage and the performance work around AST JIT projection.
Work tracking
- Land the initial integration: #15312.
- Extend Spark expression coverage using the cuDF operators below, with scoped follow-ups for the first expression families selected for implementation.
- Reduce cold run overhead and evaluate when JIT pays off, accounting for estimated work and cache reuse per executor. Follow-up: #16383.
- Add cost-based selection between JIT, interpreted AST, and regular GPU execution. Follow-up: #16384.
- Re-evaluate multi-output group limits: #15853.
- Report the final Project backend mix in explain output: #15690.
- Support Project AST and AST JIT in cached broadcast hash join build-side projections: #16509.
Prioritize expression families using representative workloads and the effort needed to match Spark semantics. Each coverage follow-up should specify the Spark expressions, supported input/output types, ANSI and null behavior, and correctness/performance test matrix. The inventory below lists candidates; it does not require enabling all 76 operations before delivering useful coverage.
cuDF operator inventory
These 76 cuDF JIT operations are candidates for Spark expression coverage, based on cuDF at 3f1e3747a38a7cc18ea678769adc258127cde560. Enabling them in the plugin requires matching Spark's types, errors, and evaluation behavior.
Standard means Java UnaryOperation / BinaryOperation; JIT means Java JitOperation. Names below use cuDF JIT/Row IR spelling; Java standard BIT_INVERT and NOT map to BITWISE_INVERT and LOGICAL_NOT.
| Family | cuDF JIT operations | Java route | Spark enablement considerations |
|---|---|---|---|
| Existing baseline | ADD, MUL |
Standard or JIT | #15312 enables only non-ANSI INT/LONG. Other widths, floating point, and decimals require separate qualification. |
| Other arithmetic | SUB, DIV, MOD, PYMOD, TRUE_DIV, FLOOR_DIV, NEG, ABS |
Standard, except NEG uses JIT |
Resolve operand promotion, wraparound, division by zero, signed minima, and Spark remainder/pmod semantics. |
| Checked arithmetic | ADD_OVERFLOW, SUB_OVERFLOW, MUL_OVERFLOW, DIV_OVERFLOW, NEG_OVERFLOW, ABS_OVERFLOW, MOD_OVERFLOW |
JIT | Map ANSI and try/null-on-error behavior explicitly; preserve error classification and evaluation order. |
| Decimal precision | CHECK_PRECISION |
JIT | Decimal input plus INT32 precision; Spark precision/scale and overflow behavior require explicit lowering. |
| Bitwise | BITWISE_AND, BITWISE_OR, BITWISE_XOR, BITWISE_INVERT, BITWISE_SHIFT_LEFT, BITWISE_SHIFT_RIGHT |
Standard for AND/OR/XOR/invert; JIT for shifts | Integer operands; check Spark shift-distance masking and RHS typing. Unsigned right shift has no dedicated opcode here. |
| Numeric/boolean casts | CAST_TO_BOOL8, CAST_TO_INT8, CAST_TO_INT16, CAST_TO_INT32, CAST_TO_INT64, CAST_TO_UINT8, CAST_TO_UINT16, CAST_TO_UINT32, CAST_TO_UINT64, CAST_TO_FLOAT32, CAST_TO_FLOAT64 |
Standard for INT64/UINT64/FLOAT64; JIT for all listed targets | Validate source types, narrowing, NaN/infinity, ANSI failures, and Spark target types. Unsigned targets are cuDF capabilities, not additional Spark SQL numeric types. |
| Decimal casts/scaling | CAST_TO_DECIMAL32, CAST_TO_DECIMAL64, CAST_TO_DECIMAL128, RESCALE |
JIT | Current decimal casts take decimal inputs and change storage width; RESCALE specifies scale. Do not infer general integer/float-to-decimal support. |
| Comparisons | EQUAL, NOT_EQUAL, LESS, LESS_EQUAL, GREATER, GREATER_EQUAL, NULL_EQUAL |
Standard | Check matching operand types, NaN/signed-zero semantics, decimal scales, and null-safe equality. |
| Boolean logic | LOGICAL_AND, LOGICAL_OR, LOGICAL_NOT, NULL_LOGICAL_AND, NULL_LOGICAL_OR |
Standard | Use Spark-compatible three-valued logic and preserve short-circuit/error behavior. Ordinary AND/OR propagate nulls and are not interchangeable with the NULL_LOGICAL variants. |
| Identity/null handling | IDENTITY, IS_NULL, COALESCE, PREDICATE |
Standard for identity/is-null; JIT for coalesce/predicate | PREDICATE turns nullable bool into an always-valid predicate. Spark coalesce needs evaluation-semantics checks, not just matching output values. |
| Conditional selection | IF_ELSE |
JIT | cuDF argument order is true value, false value, predicate. Current Row IR emits argument computations before selection; do not enable arbitrary Spark branches containing errors/effects on this basis. |
| Math | CBRT, CEIL, FLOOR, RINT, SQRT, POW, EXP, LOG |
Standard | Most accept floating point; POW also accepts integers. Match Spark result types, domain handling, rounding, and floating-point compatibility. |
| Trigonometric/hyperbolic | ARCCOS, ARCCOSH, ARCSIN, ARCSINH, ARCTAN, ARCTANH, COS, COSH, SIN, SINH, TAN, TANH |
Standard | Floating-point inputs; qualify precision and large-value/domain behavior against each supported Spark version. |
References for the inventory:
Validation and scope
Compare new expression coverage with CPU and regular GPU execution, including nulls, overflow, and ANSI/try behavior where applicable. Measure cold run and warm run query times against regular GPU and interpreted AST where supported, including preparation costs, before considering default enablement. Keep interpreted AST and JIT independently configurable so each can be evaluated and enabled on its own evidence.
This epic focuses on projection. JIT join predicates, fused filters, aggregation, and nested-column support can be tracked separately.
Additional context
- #8954 covers the broader exploration of AST Project/filter performance and defaults.
- cuDF #24464 and PR #24407 cover work to reduce AST JIT compilation time.
- cuDF #24500 requests programmatic compiler/cache/dispatch statistics.
- #15887 explores
transform_ltofor custom kernels and RAPIDS UDFs.
Activity