Support ANSI JIT AST project expressions#15069
Conversation
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Greptile SummaryThis PR wires cuDF ANSI row IR JIT AST support into the GPU project operator, adding
Confidence Score: 3/5The change is guarded by an internal config defaulting to false and depends on a cuDF JNI API still under review upstream; the new Most of the new expression wiring looks correct and is well-tested. The key concern is the
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GpuProjectExecMeta.convertToGpu] --> B{isProjectAstEnabled?}
B -- No --> Z[GpuProjectExec]
B -- Yes --> C{canUseAnsiJitAst?}
C -- No --> W[warn: LIBCUDF_JIT_ENABLED missing]
C -- Yes --> D{allReturnTypesFixedWidth AND canThisBeAst?}
D -- No --> Y[log AST fallback reason]
Y --> Z
D -- Yes --> E[GpuProjectAstExec]
E --> F[RetryableCompiledAstExpressions.compile]
F --> G{JIT ops present?}
G -- Yes --> H[cudf JIT compile JitOperator.ADD/SUB/MUL etc]
G -- No --> I[cudf non-JIT compile]
H & I --> J[computeColumns per batch]
J --> K{OOM?}
K -- No --> L[Return projected batch]
K -- Yes --> M[runWithSplitRetry: split by rows, restore = recompile]
M --> F
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A[GpuProjectExecMeta.convertToGpu] --> B{isProjectAstEnabled?}
B -- No --> Z[GpuProjectExec]
B -- Yes --> C{canUseAnsiJitAst?}
C -- No --> W[warn: LIBCUDF_JIT_ENABLED missing]
C -- Yes --> D{allReturnTypesFixedWidth AND canThisBeAst?}
D -- No --> Y[log AST fallback reason]
Y --> Z
D -- Yes --> E[GpuProjectAstExec]
E --> F[RetryableCompiledAstExpressions.compile]
F --> G{JIT ops present?}
G -- Yes --> H[cudf JIT compile JitOperator.ADD/SUB/MUL etc]
G -- No --> I[cudf non-JIT compile]
H & I --> J[computeColumns per batch]
J --> K{OOM?}
K -- No --> L[Return projected batch]
K -- Yes --> M[runWithSplitRetry: split by rows, restore = recompile]
M --> F
Reviews (1): Last reviewed commit: "use new correctness branch" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
|
||
| private def isColumnSizeOverflow(ex: Throwable): Boolean = | ||
| ex.isInstanceOf[CudfColumnSizeOverflowException] | ||
| private def isColumnSizeOverflow(ex: Throwable): Boolean = ex match { | ||
| case e: CudfColumnSizeOverflowException => | ||
| // JNI maps std::overflow_error to CudfColumnSizeOverflowException, but transform UDF | ||
| // arithmetic errors should propagate instead of split-retrying as column size overflows. | ||
| !Option(e.getMessage).exists(_.contains("error in transform UDF")) | ||
| case _ => false | ||
| } | ||
|
|
There was a problem hiding this comment.
The filter !Option(e.getMessage).exists(_.contains("error in transform UDF")) relies on the cuDF JIT runtime producing a specific error message string to distinguish a JIT arithmetic overflow from a legitimate column-size overflow. If rapidsai/cudf#22224 changes this message (e.g., during ANSI JIT stabilization or localization), two silent failures become possible: (1) JIT arithmetic overflow errors are misclassified as column-size overflows and trigger unnecessary split-retries, or (2) a real CudfColumnSizeOverflowException that happens to contain that substring is incorrectly suppressed. Since this depends on an upstream API that is explicitly named as "still under review," the message contract is not stable.
| ("lhs", TypeSig.LONG + TypeSig.DECIMAL_128, TypeSig.LONG + TypeSig.DECIMAL_128), | ||
| ("rhs", TypeSig.LONG + TypeSig.DECIMAL_128, TypeSig.LONG + TypeSig.DECIMAL_128)), | ||
| (a, conf, p, r) => new BinaryExprMeta[IntegralDivide](a, conf, p, r) { | ||
| (a, conf, p, r) => new BinaryAstExprMeta[IntegralDivide](a, conf, p, r) { | ||
| override def tagSelfForAst(): Unit = { | ||
| super.tagSelfForAst() | ||
| if (!SQLConf.get.ansiEnabled || !conf.isProjectAstAnsiArithmeticEnabled || | ||
| !GpuAnsi.supportsAnsiArithmeticAst(a.dataType)) { |
There was a problem hiding this comment.
IntegralDivide AST gate uses a weaker check than the spark330db shim
spark330 checks !GpuAnsi.supportsAnsiArithmeticAst(a.dataType) (output type = LONG, always true for IntegralDivide), while spark330db uses !GpuAnsi.shouldUseAnsiDivModAst(...), which also validates that LHS and RHS types match and are themselves supported. Today the ExprChecks AST signature limits inputs to LONG, so the weaker check is safe, but the inconsistency across shims means a future type-signature widening in this shim could silently skip the input-type guard. Using shouldUseAnsiDivModAst here would align both shims.
| case LongType => ast.Literal.ofLong(value.asInstanceOf[java.lang.Long]) | ||
| case FloatType => ast.Literal.ofFloat(value.asInstanceOf[java.lang.Float]) | ||
| case DoubleType => ast.Literal.ofDouble(value.asInstanceOf[java.lang.Double]) | ||
| case dt: DecimalType => | ||
| val unscaledValue = value match { | ||
| case null => null | ||
| case d: Decimal => d.toBigDecimal.bigDecimal.unscaledValue() | ||
| case d: java.math.BigDecimal => d.unscaledValue() |
There was a problem hiding this comment.
BigInteger passed to ast.Literal.ofDecimal when value is null
When value is null, unscaledValue is null, and ast.Literal.ofDecimal(DecimalUtil.createCudfDecimal(dt), null) is called. The cuDF Java binding behavior for a null BigInteger argument is not specified in the change description, and the PR explicitly blocks decimal literals in AST via GpuOverrides (willNotWorkInAst), so this path is currently unreachable. However, a future enablement of decimal literals in AST would hit this NPE (or silently wrong-value) path. Consider either throwing explicitly on null or delegating null-decimal literals to a cuDF null-literal API.
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
# Conflicts: # integration_tests/src/main/python/ast_test.py
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
2 participants
Fixes #xxxx.
Current cuDF/JNI baseline
The current code is based on https://github.com/thirtiseven/cudf/tree/codex-ansi-jit-22602-22680-22836-jni-test.
This branch combines the current cuDF work from rapidsai/cudf#22602, rapidsai/cudf#22680, and rapidsai/cudf#22836 with the Java/JNI changes required by this prototype.
Description
ANSI JIT Project Benchmark
10000000016313Cache modes:
GPU_AST_JIT_COLD: runs in a fresh Spark application with an empty JIT cache.GPU_AST_JIT_DISK_WARM: runs in a fresh application with the target kernel already in the on-disk cache, but not loaded in the process.GPU_AST_JIT_PCH_WARM_KERNEL_COLD: creates a PCH with a distinct probe, then measures the first compilation of the target kernel.GPU_AST_JIT_HOT: warms up the exact target expression and measures reuse from the in-process kernel cache.PCH generation, disk-cache population, and hot warmup time are excluded from their corresponding measured results.
Previous benchmark results
based on draft pr
based on new pr
based on new pr, after RTCX cache optimization:
results based on merge of 22602 + 22680 + 22836
on cudf main branch
This PR adds initial spark-rapids support for cuDF ANSI row IR JIT project expressions.
The intent is to evaluate and review the Spark-side wiring against the cuDF/JNI baseline linked above. The feature remains guarded by existing AST/JIT configuration and by expression-level gating.
The changes include:
DIVand%project AST support where Spark and cuDF row IR semantics are expected to matchDIVIFexpressions withNULLIFY_IF+COALESCEOpen items for this draft:
%support depends on cuDFANSI_MODpreserving cuDFMODsemantics. If cuDF intentionally keepsANSI_MODas floor/Python-style modulo, Spark should gate%out of AST and fall back.Validation in this PR creation step:
git diff --checkChecklists
Documentation
Testing
(Please provide the names of the existing tests in the PR description.)
Performance