Skip to content

feat: implement pmod natively for all numeric types#5029

Draft
andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:feat/pmod-native
Draft

feat: implement pmod natively for all numeric types#5029
andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:feat/pmod-native

Conversation

@andygrove

@andygrove andygrove commented Jul 24, 2026

Copy link
Copy Markdown
Member

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

pmod was 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 upstream datafusion-spark pmod UDF reads enable_ansi_mode from the session config, which Comet never sets, so a zero divisor would return NULL instead of raising an error under spark.sql.ansi.enabled=true.

This PR replaces the dispatch fallback with a native DataFusion implementation that mirrors the existing Remainder path, so ANSI behaviour is honoured natively and the per-batch JVM round trip is removed.

What changes are included in this PR?

  • New MathExpr pmod protobuf message and PmodBuilder, mirroring Remainder, so the eval mode reaches the native planner.
  • spark_pmod kernel and create_pmod_expr in modulo_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.
  • CometPmod serde now serializes to the new proto message instead of extending CometCodegenDispatch. All numeric input types are supported, including decimal.
  • Audit notes for pmod added to the math_funcs expression-audit page.

The implementation was scaffolded with the implement-comet-expression skill and reviewed with the audit-comet-expression skill, which confirmed the semantics are identical across Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1. The 4.1.1 NumericEvalContext constructor change and remainderByZeroError rename are both handled without a shim.

How are these changes tested?

  • New spark_pmod Rust unit tests cover integer and decimal inputs, all sign combinations, and zero-divisor handling in both legacy and ANSI modes.
  • New Comet SQL file tests pmod.sql (legacy) and pmod_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

CometPmodBenchmark measures SELECT count(*) FROM t WHERE pmod(c1, k) > n so 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.

Type dictionary Spark (ms) Comet (ms) Comet vs Spark
INT true 56 93 0.6x
INT false 51 92 0.6x
BIGINT true 49 119 0.4x
BIGINT false 89 152 0.6x
DOUBLE true 66 134 0.5x
DOUBLE false 250 304 0.8x
DECIMAL(18,2) true 379 236 1.6x
DECIMAL(18,2) false 3122 276 11.3x
DECIMAL(38,10) true 690 270 2.6x
DECIMAL(38,10) false 3867 321 12.0x

Decimal pmod is 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 to pmod.

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
andygrove marked this pull request as draft July 24, 2026 20:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant