feat: implement pmod natively for all numeric types#5029
Draft
andygrove wants to merge 2 commits into
Draft
Conversation
Replace the codegen-dispatch fallback for `pmod` (Pmod) with a native DataFusion implementation. `CometPmod` now serializes to a `MathExpr pmod` protobuf message that mirrors `Remainder`, carrying the eval mode so ANSI behaviour is honoured natively rather than being invisible to the engine. The new `spark_pmod` kernel computes `((left % right) + right) % right`, returning NULL on a zero divisor in legacy mode and raising a remainder-by-zero error in ANSI mode. All numeric input types are supported, including decimal, with a Decimal256 intermediate for wide decimals to match the modulo path. Adds SQL file tests for legacy and ANSI modes plus Rust unit tests covering integer, decimal, sign combinations, and zero-divisor handling.
Add CometPmodBenchmark, which measures pmod evaluation cost using aggregate queries of the form 'SELECT count(*) FROM t WHERE pmod(c1, k) > n'. Aggregating to a single row keeps the result small so the benchmark reflects the scan and predicate evaluation rather than transferring a large result set back over Arrow FFI. Covers integer, long, double, and decimal inputs with and without Parquet dictionary encoding.
andygrove
marked this pull request as draft
July 24, 2026 20:16
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
There is no dedicated issue. This is part of the ongoing effort to provide native implementations for expressions that are currently handled through JVM codegen dispatch rather than native execution. Related to the expression coverage epic #240.
Rationale for this change
pmodwas previously handled by the JVM codegen dispatcher, which keeps the projection on the native pipeline but still evaluates the expression as Spark-generated Java per batch (Arrow to Spark rows to eval to Arrow). It also could not honour ANSI mode: the upstreamdatafusion-sparkpmodUDF readsenable_ansi_modefrom the session config, which Comet never sets, so a zero divisor would return NULL instead of raising an error underspark.sql.ansi.enabled=true.This PR replaces the dispatch fallback with a native DataFusion implementation that mirrors the existing
Remainderpath, so ANSI behaviour is honoured natively and the per-batch JVM round trip is removed.What changes are included in this PR?
MathExpr pmodprotobuf message andPmodBuilder, mirroringRemainder, so the eval mode reaches the native planner.spark_pmodkernel andcreate_pmod_exprinmodulo_expr.rs. The kernel computes((left % right) + right) % right, returning NULL on a zero divisor in legacy mode and raising a remainder-by-zero error in ANSI mode. Wide decimals use a Decimal256 intermediate, matching the modulo path.CometPmodserde now serializes to the new proto message instead of extendingCometCodegenDispatch. All numeric input types are supported, including decimal.pmodadded to themath_funcsexpression-audit page.The implementation was scaffolded with the
implement-comet-expressionskill and reviewed with theaudit-comet-expressionskill, which confirmed the semantics are identical across Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1. The 4.1.1NumericEvalContextconstructor change andremainderByZeroErrorrename are both handled without a shim.How are these changes tested?
spark_pmodRust unit tests cover integer and decimal inputs, all sign combinations, and zero-divisor handling in both legacy and ANSI modes.pmod.sql(legacy) andpmod_ansi.sql(ANSI) exercise column and literal arguments, every numeric type, negative operands, NaN and Infinity, float precision, INT_MIN and Long boundaries, mixed-scale decimals, the wide-decimal Decimal256 path, legacy NULL-on-zero, and the ANSI zero-divisor error.Performance
CometPmodBenchmarkmeasuresSELECT count(*) FROM t WHERE pmod(c1, k) > nso that the timing reflects the scan and predicate evaluation rather than transferring a large result set over Arrow FFI. All cases run fully natively in Comet. Numbers below are release builds on an Apple M3 Ultra (OpenJDK 17), 10M rows, best of the reported runs.Decimal
pmodis substantially faster natively (up to 12x), since Spark's decimal arithmetic path is comparatively slow. For the primitive types the predicate is cheap enough that the end-to-end query is dominated by scan and aggregation, where Spark's whole-stage codegen is very efficient; the gap there reflects Comet's general per-batch overhead on a trivial predicate rather than anything specific topmod.