From cbaac8dbe7226afd0894fd43349d5683e70e1427 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 13:16:18 -0600 Subject: [PATCH 1/2] feat: implement pmod natively for all numeric types 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. --- .../expression-audits/math_funcs.md | 5 + .../src/execution/expressions/arithmetic.rs | 34 ++- .../execution/planner/expression_registry.rs | 4 + native/proto/src/proto/expr.proto | 1 + native/spark-expr/src/comet_scalar_funcs.rs | 5 +- native/spark-expr/src/lib.rs | 4 +- native/spark-expr/src/math_funcs/mod.rs | 2 +- .../spark-expr/src/math_funcs/modulo_expr.rs | 262 +++++++++++++++++- .../scala/org/apache/comet/serde/math.scala | 18 +- .../sql-tests/expressions/math/pmod.sql | 92 +++++- .../sql-tests/expressions/math/pmod_ansi.sql | 43 +++ 11 files changed, 459 insertions(+), 11 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/math/pmod_ansi.sql diff --git a/docs/source/contributor-guide/expression-audits/math_funcs.md b/docs/source/contributor-guide/expression-audits/math_funcs.md index 6c9607efc5..093778c681 100644 --- a/docs/source/contributor-guide/expression-audits/math_funcs.md +++ b/docs/source/contributor-guide/expression-audits/math_funcs.md @@ -188,6 +188,11 @@ Internal fused expression that rescales a Decimal128 value (changing scale) and - Spark 3.4.3, 3.5.8, 4.0.1, 4.1.1 (audited 2026-05-27): `LeafMathExpression(math.Pi, "PI")`; foldable, so Spark `ConstantFolding` rewrites it to a `Literal` before Comet sees the plan. The `CometScalarFunction("pi")` registration is exercised only when `ConstantFolding` is excluded. +## pmod + +- Spark 3.4.3, 3.5.8, 4.0.1 (audited 2026-07-24): `Pmod(left, right, evalMode)` signature identical across these versions. `CometPmod` serializes to the `MathExpr pmod` proto (mirroring `Remainder`), carrying the eval mode. The native `spark_pmod` UDF computes `((left % right) + right) % right`; non-ANSI returns NULL on a zero divisor, ANSI raises `DIVIDE_BY_ZERO`. All numeric input types are supported, including decimal (wide decimals use a Decimal256 intermediate, matching modulo). +- Spark 4.1.1 (audited 2026-07-24): constructor changed to `Pmod(left, right, evalContext: NumericEvalContext)`, but `BinaryArithmetic.evalMode` is still available so no shim is needed. The ANSI zero-divisor error changed to `REMAINDER_BY_ZERO`; Comet's native error is `RemainderByZero`, which matches. + ## positive - Spark 3.4.3, 3.5.8 (audited 2026-05-27): `UnaryPositive(child)` is a regular expression. There is no Comet serde for `UnaryPositive`, so projections containing `+col` silently disable Comet for the projection on 3.4/3.5. diff --git a/native/core/src/execution/expressions/arithmetic.rs b/native/core/src/execution/expressions/arithmetic.rs index 8d4c59a010..bb2fd296f2 100644 --- a/native/core/src/execution/expressions/arithmetic.rs +++ b/native/core/src/execution/expressions/arithmetic.rs @@ -167,7 +167,9 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use datafusion::logical_expr::Operator as DataFusionOperator; use datafusion_comet_proto::spark_expression::Expr; -use datafusion_comet_spark_expr::{create_modulo_expr, create_negate_expr, EvalMode}; +use datafusion_comet_spark_expr::{ + create_modulo_expr, create_negate_expr, create_pmod_expr, EvalMode, +}; use crate::execution::{ expressions::extract_expr, @@ -251,6 +253,36 @@ impl ExpressionBuilder for RemainderBuilder { } } +/// Builder for Pmod expressions (uses special positive-modulo function) +pub struct PmodBuilder; + +impl ExpressionBuilder for PmodBuilder { + fn build( + &self, + spark_expr: &Expr, + input_schema: SchemaRef, + planner: &PhysicalPlanner, + ) -> Result, ExecutionError> { + let expr = extract_expr!(spark_expr, Pmod); + let eval_mode = from_protobuf_eval_mode(expr.eval_mode)?; + let left = planner.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?; + let right = planner.create_expr(expr.right.as_ref().unwrap(), Arc::clone(&input_schema))?; + + let result = create_pmod_expr( + left, + right, + expr.return_type + .as_ref() + .map(crate::execution::serde::to_arrow_datatype) + .unwrap(), + input_schema, + eval_mode == EvalMode::Ansi, + &planner.session_ctx().state(), + ); + result.map_err(|e| ExecutionError::GeneralError(e.to_string())) + } +} + /// Builder for UnaryMinus expressions (uses special negate function) pub struct UnaryMinusBuilder; diff --git a/native/core/src/execution/planner/expression_registry.rs b/native/core/src/execution/planner/expression_registry.rs index 7fe7a477dd..0834544be2 100644 --- a/native/core/src/execution/planner/expression_registry.rs +++ b/native/core/src/execution/planner/expression_registry.rs @@ -47,6 +47,7 @@ pub enum ExpressionType { Divide, IntegralDivide, Remainder, + Pmod, UnaryMinus, // Comparison expressions @@ -212,6 +213,8 @@ impl ExpressionRegistry { ); self.builders .insert(ExpressionType::Remainder, Box::new(RemainderBuilder)); + self.builders + .insert(ExpressionType::Pmod, Box::new(PmodBuilder)); self.builders .insert(ExpressionType::UnaryMinus, Box::new(UnaryMinusBuilder)); } @@ -325,6 +328,7 @@ impl ExpressionRegistry { Some(ExprStruct::Divide(_)) => Ok(ExpressionType::Divide), Some(ExprStruct::IntegralDivide(_)) => Ok(ExpressionType::IntegralDivide), Some(ExprStruct::Remainder(_)) => Ok(ExpressionType::Remainder), + Some(ExprStruct::Pmod(_)) => Ok(ExpressionType::Pmod), Some(ExprStruct::UnaryMinus(_)) => Ok(ExpressionType::UnaryMinus), Some(ExprStruct::Eq(_)) => Ok(ExpressionType::Eq), diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index cbb1043f21..7db013004c 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -92,6 +92,7 @@ message Expr { JvmScalarUdf jvm_scalar_udf = 70; PreciseTimestampConversion precise_timestamp_conversion = 71; Shuffle shuffle = 72; + MathExpr pmod = 73; } reserved 20; diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 8e913dabab..b6065eda10 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -21,7 +21,7 @@ use crate::map_funcs::spark_map_sort; use crate::math_funcs::abs::abs; use crate::math_funcs::checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub}; use crate::math_funcs::log::spark_log; -use crate::math_funcs::modulo_expr::spark_modulo; +use crate::math_funcs::modulo_expr::{spark_modulo, spark_pmod}; use crate::{ spark_ceil, spark_day_name, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding, @@ -185,6 +185,9 @@ pub fn create_comet_physical_fun_with_eval_mode( let func = Arc::new(spark_modulo); make_comet_scalar_udf!("spark_modulo", func, without data_type, fail_on_error) } + "spark_pmod" => { + make_comet_scalar_udf!("spark_pmod", spark_pmod, data_type, fail_on_error) + } "abs" => { let func = Arc::new(abs); make_comet_scalar_udf!("abs", func, without data_type) diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 40422cde07..7abf54551d 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -85,8 +85,8 @@ pub use hash_funcs::*; pub use json_funcs::{FromJson, ToJson}; pub use math_funcs::{ checked_add, checked_div, checked_mul, checked_sub, create_modulo_expr, create_negate_expr, - spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_log, - spark_make_decimal, spark_round, spark_unhex, spark_unscaled_value, CheckOverflow, + create_pmod_expr, spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, + spark_log, spark_make_decimal, spark_round, spark_unhex, spark_unscaled_value, CheckOverflow, DecimalRescaleCheckOverflow, NegativeExpr, NormalizeNaNAndZero, WideDecimalBinaryExpr, WideDecimalOp, }; diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs index 51bd016a2a..1a4e5aac88 100644 --- a/native/spark-expr/src/math_funcs/mod.rs +++ b/native/spark-expr/src/math_funcs/mod.rs @@ -36,7 +36,7 @@ pub use div::spark_decimal_integral_div; pub use floor::spark_floor; pub use internal::*; pub use log::spark_log; -pub use modulo_expr::create_modulo_expr; +pub use modulo_expr::{create_modulo_expr, create_pmod_expr}; pub use negative::{create_negate_expr, NegativeExpr}; pub use round::spark_round; pub use unhex::spark_unhex; diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs index bc31b3a125..f83a3306b4 100644 --- a/native/spark-expr/src/math_funcs/modulo_expr.rs +++ b/native/spark-expr/src/math_funcs/modulo_expr.rs @@ -17,7 +17,11 @@ use crate::{create_comet_physical_fun, IfExpr}; use crate::{remainder_by_zero_error, Cast, EvalMode, SparkCastOptions}; -use arrow::compute::kernels::numeric::rem; +use arrow::array::{new_null_array, ArrayRef, Scalar}; +use arrow::compute::kernels::cast::cast; +use arrow::compute::kernels::cmp::{eq, lt}; +use arrow::compute::kernels::numeric::{add, rem}; +use arrow::compute::kernels::zip::zip; use arrow::datatypes::*; use datafusion::common::{exec_err, internal_err, DataFusionError, Result, ScalarValue}; use datafusion::config::ConfigOptions; @@ -136,6 +140,70 @@ pub fn create_modulo_expr( } } +/// Computes `left % right` with Spark-compliant divide-by-zero handling. +/// In ANSI mode (`fail_on_error`), a zero divisor raises Spark's remainder-by-zero error. +/// In legacy mode, zero divisors are replaced with `NULL` before the remainder is computed, so +/// those rows return `NULL` while the rest compute normally. +fn try_rem(left: &ArrayRef, right: &ArrayRef, fail_on_error: bool) -> Result { + if fail_on_error { + match rem(left, right) { + Ok(result) => Ok(result), + Err(e) if e.to_string().contains("Divide by zero") => { + Err(remainder_by_zero_error().into()) + } + Err(e) => Err(e.into()), + } + } else { + let zero = Scalar::new(ScalarValue::new_zero(right.data_type())?.to_array()?); + let null = Scalar::new(new_null_array(right.data_type(), 1)); + let is_zero = eq(right, &zero)?; + let safe_right = zip(&is_zero, &null, right)?; + Ok(rem(left, &safe_right)?) + } +} + +/// Spark-compliant `pmod` (positive modulo) function. Returns a non-negative remainder when the +/// divisor is positive, matching Spark's `MathUtils.pmod`. If `fail_on_error` is true (ANSI mode) +/// then a zero divisor raises an error, otherwise those rows return `NULL`. +/// +/// The result is computed as `((left % right) + right) % right`, adding the divisor back only where +/// the plain remainder is negative and taking the remainder a second time to normalise the +/// negative-divisor case. `data_type` is Spark's declared result type: Arrow's decimal arithmetic +/// can widen the precision or scale of the intermediates, so the final result is cast back to it. +/// The final value always fits in `data_type`, so the cast is lossless. +pub fn spark_pmod( + args: &[ColumnarValue], + data_type: &DataType, + fail_on_error: bool, +) -> Result { + if args.len() != 2 { + return exec_err!("pmod expects exactly two arguments"); + } + + let args = ColumnarValue::values_to_arrays(args)?; + let left = &args[0]; + let right = &args[1]; + + let remainder = try_rem(left, right, fail_on_error)?; + // Zero arrays are built from the operand types they are compared/combined with so that Arrow's + // decimal kernels see matching scales. + let remainder_zero = + ScalarValue::new_zero(remainder.data_type())?.to_array_of_size(remainder.len())?; + let right_zero = ScalarValue::new_zero(right.data_type())?.to_array_of_size(right.len())?; + // Where the remainder is negative, add the divisor back to make it non-negative; elsewhere add + // zero (a no-op). + let negative = lt(&remainder, &remainder_zero)?; + let addend = zip(&negative, right, &right_zero)?; + let shifted = add(&addend, &remainder)?; + let result = try_rem(&shifted, right, fail_on_error)?; + + if result.data_type() == data_type { + Ok(ColumnarValue::Array(result)) + } else { + Ok(ColumnarValue::Array(cast(&result, data_type)?)) + } +} + fn null_if_zero_primitive( expression: Arc, input_schema: &Schema, @@ -212,6 +280,85 @@ fn create_modulo_scalar_function( ))) } +/// Builds the physical expression for Spark's `pmod`. Zero-divisor handling lives inside the +/// `spark_pmod` kernel, so the operands are passed through unchanged (no null-if-zero wrapping is +/// needed, unlike modulo). As with modulo, when the decimal operands are wide enough that the +/// intermediate `(remainder + divisor)` could exceed `Decimal128`'s maximum precision, both operands +/// are promoted to `Decimal256` and the result is cast back. +pub fn create_pmod_expr( + left: Arc, + right: Arc, + data_type: DataType, + input_schema: SchemaRef, + fail_on_error: bool, + registry: &dyn FunctionRegistry, +) -> Result, DataFusionError> { + match ( + left.data_type(&input_schema), + right.data_type(&input_schema), + ) { + (Ok(DataType::Decimal128(p1, s1)), Ok(DataType::Decimal128(p2, s2))) + if max(s1, s2) as u8 + max(p1 - s1 as u8, p2 - s2 as u8) > DECIMAL128_MAX_PRECISION => + { + let left_256 = Arc::new(Cast::new( + left, + DataType::Decimal256(p1, s1), + SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + None, + None, + )); + let right_256 = Arc::new(Cast::new( + right, + DataType::Decimal256(p2, s2), + SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + None, + None, + )); + + // Operating on Decimal256 inputs, the kernel returns Decimal256. + let decimal256_return_type = match &data_type { + DataType::Decimal128(p, s) => DataType::Decimal256(*p, *s), + other => other.clone(), + }; + let pmod_scalar_func = create_pmod_scalar_function( + left_256, + right_256, + &decimal256_return_type, + registry, + fail_on_error, + )?; + + Ok(Arc::new(Cast::new( + pmod_scalar_func, + data_type, + SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + None, + None, + ))) + } + _ => create_pmod_scalar_function(left, right, &data_type, registry, fail_on_error), + } +} + +fn create_pmod_scalar_function( + left: Arc, + right: Arc, + data_type: &DataType, + registry: &dyn FunctionRegistry, + fail_on_error: bool, +) -> Result, DataFusionError> { + let func_name = "spark_pmod"; + let pmod_expr = + create_comet_physical_fun(func_name, data_type.clone(), registry, Some(fail_on_error))?; + Ok(Arc::new(ScalarFunctionExpr::new( + func_name, + pmod_expr, + vec![left, right], + Arc::new(Field::new(func_name, data_type.clone(), true)), + Arc::new(ConfigOptions::default()), + ))) +} + #[cfg(test)] mod tests { use super::*; @@ -441,4 +588,117 @@ mod tests { verify_result(modulo_expr, batch, fail_on_error, Some(expected_result)); }) } + + #[test] + fn test_pmod_basic_int() { + with_fail_on_error(|fail_on_error| { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + // Cover the sign combinations: the result takes the sign of the divisor, matching + // Spark's MathUtils.pmod. + let a_array = Arc::new(Int32Array::from(vec![7, -7, 7, -7, i32::MIN])); + let b_array = Arc::new(Int32Array::from(vec![3, 3, -3, -3, 3])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array]).unwrap(); + + let session_ctx = SessionContext::new(); + let pmod_expr = create_pmod_expr( + Arc::new(Column::new("a", 0)), + Arc::new(Column::new("b", 1)), + DataType::Int32, + schema, + fail_on_error, + &session_ctx.state(), + ) + .unwrap(); + + let should_fail = false; + // pmod adds the divisor back only when the plain remainder is negative: + // 7 pmod 3 = 1 (7 % 3 = 1, already >= 0) + // -7 pmod 3 = 2 (-7 % 3 = -1, then -1 + 3 = 2) + // 7 pmod -3 = 1 (7 % -3 = 1, already >= 0) + // -7 pmod -3 = -1 (-7 % -3 = -1, then (-1 + -3) % -3 = -1) + // i32::MIN pmod 3 = 1 (no overflow) + let expected_result = Arc::new(Int32Array::from(vec![1, 2, 1, -1, 1])); + verify_result(pmod_expr, batch, should_fail, Some(expected_result)); + }) + } + + #[test] + fn test_pmod_divide_by_zero_int() { + with_fail_on_error(|fail_on_error| { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + let a_array = Arc::new(Int32Array::from(vec![7])); + let b_array = Arc::new(Int32Array::from(vec![0])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array]).unwrap(); + + let session_ctx = SessionContext::new(); + let pmod_expr = create_pmod_expr( + Arc::new(Column::new("a", 0)), + Arc::new(Column::new("b", 1)), + DataType::Int32, + schema, + fail_on_error, + &session_ctx.state(), + ) + .unwrap(); + + // ANSI mode errors on a zero divisor; legacy mode returns NULL. + let expected_result = Arc::new(Int32Array::from(vec![None])); + verify_result(pmod_expr, batch, fail_on_error, Some(expected_result)); + }) + } + + #[test] + fn test_pmod_basic_decimal() { + with_fail_on_error(|fail_on_error| { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Decimal128(18, 4), false), + Field::new("b", DataType::Decimal128(18, 4), false), + ])); + + // 7.5, -7.5 at scale 4 + let mut a_builder = + Decimal128Builder::with_capacity(2).with_data_type(DataType::Decimal128(18, 4)); + a_builder.append_value(75000); + a_builder.append_value(-75000); + let a_array: ArrayRef = Arc::new(a_builder.finish()); + + // 3.0, 3.0 at scale 4 + let mut b_builder = + Decimal128Builder::with_capacity(2).with_data_type(DataType::Decimal128(18, 4)); + b_builder.append_value(30000); + b_builder.append_value(30000); + let b_array: ArrayRef = Arc::new(b_builder.finish()); + + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a_array, b_array]).unwrap(); + + let session_ctx = SessionContext::new(); + let pmod_expr = create_pmod_expr( + Arc::new(Column::new("a", 0)), + Arc::new(Column::new("b", 1)), + DataType::Decimal128(18, 4), + schema, + fail_on_error, + &session_ctx.state(), + ) + .unwrap(); + + let should_fail = false; + // 7.5 pmod 3.0 = 1.5 -> 15000 + // -7.5 pmod 3.0 = 1.5 -> 15000 (non-negative because divisor is positive) + let expected_result = Arc::new( + Decimal128Array::from(vec![Some(15000), Some(15000)]) + .with_precision_and_scale(18, 4) + .unwrap(), + ); + verify_result(pmod_expr, batch, should_fail, Some(expected_result)); + }) + } } diff --git a/spark/src/main/scala/org/apache/comet/serde/math.scala b/spark/src/main/scala/org/apache/comet/serde/math.scala index 0ec2b5cd9e..f24cf836b4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/math.scala +++ b/spark/src/main/scala/org/apache/comet/serde/math.scala @@ -282,7 +282,23 @@ object CometConv extends CometCodegenDispatch[Conv] object CometLog1p extends CometCodegenDispatch[Log1p] -object CometPmod extends CometCodegenDispatch[Pmod] +object CometPmod extends CometExpressionSerde[Pmod] with MathBase { + + override def convert( + expr: Pmod, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + createMathExpression( + expr, + expr.left, + expr.right, + inputs, + binding, + expr.dataType, + expr.evalMode, + (builder, mathExpr) => builder.setPmod(mathExpr)) + } +} object CometWidthBucket extends CometCodegenDispatch[WidthBucket] diff --git a/spark/src/test/resources/sql-tests/expressions/math/pmod.sql b/spark/src/test/resources/sql-tests/expressions/math/pmod.sql index 0ca2a0b3ce..7261b6d33b 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/pmod.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/pmod.sql @@ -15,17 +15,101 @@ -- specific language governing permissions and limitations -- under the License. --- Routes pmod through the codegen dispatcher so behavior matches Spark exactly. +-- pmod in legacy (non-ANSI) mode: a zero divisor returns NULL. +-- Config: spark.sql.ansi.enabled=false statement CREATE TABLE test_pmod(a int, b int) USING parquet statement -INSERT INTO test_pmod VALUES (7, 3), (-7, 3), (7, -3), (0, 5), (5, NULL), (NULL, 5) +INSERT INTO test_pmod VALUES + (7, 3), + (-7, 3), + (7, -3), + (-7, -3), + (0, 5), + (5, 0), + (5, NULL), + (NULL, 5), + (-2147483648, 3) +-- column arguments, including negative operands and a zero divisor (NULL in legacy mode) query SELECT a, b, pmod(a, b) FROM test_pmod --- literal arguments including doubles +-- tinyint / smallint / bigint +statement +CREATE TABLE test_pmod_int_types(t tinyint, s smallint, l bigint) USING parquet + +statement +INSERT INTO test_pmod_int_types VALUES + (-7, -7, -7), + (7, 7, 7), + (-1, -1, -1), + (0, 0, 0) + +query +SELECT + pmod(t, cast(3 as tinyint)), + pmod(s, cast(3 as smallint)), + pmod(l, cast(3 as bigint)) +FROM test_pmod_int_types + +-- float and double +statement +CREATE TABLE test_pmod_fp(f float, d double) USING parquet + +statement +INSERT INTO test_pmod_fp VALUES + (-7.5, -7.5), + (7.5, 7.5), + (10.5, 10.5), + (0.0, 0.0) + +query +SELECT pmod(f, cast(3.0 as float)), pmod(d, 3.0D) FROM test_pmod_fp + +-- floating-point special values: NaN pmod x = NaN, Inf pmod x = NaN, x pmod Inf = x +query +SELECT + pmod(cast('NaN' as double), 3.0D), + pmod(cast('Infinity' as double), 3.0D), + pmod(5.0D, cast('Infinity' as double)) + +-- literal arguments (constant folding is disabled by the test suite) +query +SELECT pmod(-7, 3), pmod(7, -3), pmod(-7, -3), pmod(5, 0), pmod(0, 5) + +-- precision and boundary cases exercised by Spark's own ArithmeticExpressionSuite +query +SELECT pmod(7.2D, 4.1D), pmod(2L, 9223372036854775807L), pmod(-9223372036854775808L, 3L) + +-- decimal, including negative operands and a zero divisor (NULL in legacy mode) +statement +CREATE TABLE test_pmod_dec(a decimal(10,2), b decimal(10,2)) USING parquet + +statement +INSERT INTO test_pmod_dec VALUES + (7.5, 3.0), + (-7.5, 3.0), + (7.5, -3.0), + (-7.5, -3.0), + (0.7, 0.2), + (5.0, 0.0), + (5.0, NULL), + (NULL, 3.0) + +query +SELECT a, b, pmod(a, b) FROM test_pmod_dec + +-- decimal literals with differing precision and scale (exercises scale coercion) +query +SELECT + pmod(cast(0.7 as decimal(2,1)), cast(0.2 as decimal(2,1))), + pmod(cast(7.25 as decimal(5,2)), cast(2.5 as decimal(3,1))), + pmod(cast(-7.25 as decimal(5,2)), cast(2.5 as decimal(3,1))) + +-- wide decimals that exceed Decimal128 precision for the intermediate, exercising the Decimal256 +-- promotion path query -SELECT pmod(-7, 3), pmod(10.5, 3.0), pmod(7, -3) +SELECT pmod(cast(1.5 as decimal(38,30)), cast(0.4 as decimal(38,20))) diff --git a/spark/src/test/resources/sql-tests/expressions/math/pmod_ansi.sql b/spark/src/test/resources/sql-tests/expressions/math/pmod_ansi.sql new file mode 100644 index 0000000000..be2f8922f2 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/pmod_ansi.sql @@ -0,0 +1,43 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- pmod in ANSI mode: a zero divisor raises an error instead of returning NULL. +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE test_pmod_ansi(a int, b int) USING parquet + +statement +INSERT INTO test_pmod_ansi VALUES (7, 3), (-7, 3), (7, -3), (-7, -3) + +-- Sentinel query: non-zero divisors must compute natively and match Spark. This guards against a +-- vacuous pass where Comet falls back to Spark for the error cases below. +query +SELECT a, b, pmod(a, b) FROM test_pmod_ansi + +-- column zero divisor throws. Spark 4.0 raises DIVIDE_BY_ZERO, Spark 4.1 raises REMAINDER_BY_ZERO; +-- match the common substring. +query expect_error(BY_ZERO) +SELECT pmod(a, 0) FROM test_pmod_ansi + +-- literal zero divisor throws (constant folding is disabled by the test suite) +query expect_error(BY_ZERO) +SELECT pmod(7, 0) + +-- decimal zero divisor also throws in ANSI mode +query expect_error(BY_ZERO) +SELECT pmod(cast(5.0 as decimal(10,2)), cast(0.0 as decimal(10,2))) From 76b73829f2196f01cac0e9dbe77b551fa045d013 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 13:51:59 -0600 Subject: [PATCH 2/2] test: add end-to-end pmod predicate microbenchmark 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. --- .../sql/benchmark/CometPmodBenchmark.scala | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometPmodBenchmark.scala diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPmodBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPmodBenchmark.scala new file mode 100644 index 0000000000..ac7ae80eb4 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPmodBenchmark.scala @@ -0,0 +1,101 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.sql.types._ + +/** + * Benchmark to measure the cost of evaluating `pmod` as a filter predicate. + * + * The queries take the form `SELECT count(*) FROM t WHERE pmod(c1, k) > n`. Aggregating to a + * single row keeps the result tiny, so the benchmark measures the scan and predicate evaluation + * rather than the cost of transferring a large result set back to the JVM over Arrow FFI. + * + * To run this benchmark: `SPARK_GENERATE_BENCHMARK_FILES=1 make + * benchmark-org.apache.spark.sql.benchmark.CometPmodBenchmark` Results will be written to + * "spark/benchmarks/CometPmodBenchmark-**results.txt". + */ +object CometPmodBenchmark extends CometBenchmarkBase { + private val table = "parquetV1Table" + + // A non-zero literal divisor and a threshold that selects part of the [0, divisor) output range. + private val divisor = 7 + private val threshold = 3 + + private def integerPmodBenchmark( + values: Int, + dataType: DataType, + useDictionary: Boolean): Unit = { + import spark.implicits._ + withTempPath { dir => + withTempTable(table) { + // A small distinct set encourages Parquet dictionary encoding; otherwise use the raw id. + val col = if (useDictionary) ($"id" % 8) else $"id" + prepareTable(dir, spark.range(values).select(col.cast(dataType).as("c1"))) + + val name = s"pmod(${dataType.sql}), dictionary = $useDictionary" + val query = s"SELECT count(*) FROM $table WHERE pmod(c1, $divisor) > $threshold" + runExpressionBenchmark(name, values, query) + } + } + } + + private def decimalPmodBenchmark( + values: Int, + dataType: DecimalType, + useDictionary: Boolean): Unit = { + import spark.implicits._ + withTempPath { dir => + withTempTable(table) { + // Bounded so the values fit the target precision and scale (no overflow-to-NULL). + val raw = if (useDictionary) ($"id" % 8) else ($"id" % 100000) + val col = (raw / 100.0).cast(dataType) + prepareTable(dir, spark.range(values).select(col.as("c1"))) + + val name = s"pmod(${dataType.sql}), dictionary = $useDictionary" + val query = + s"SELECT count(*) FROM $table " + + s"WHERE pmod(c1, CAST($divisor AS ${dataType.sql})) > CAST($threshold AS ${dataType.sql})" + runExpressionBenchmark(name, values, query) + } + } + } + + private val TOTAL: Int = 1024 * 1024 * 10 + + override def runCometBenchmark(args: Array[String]): Unit = { + Seq(true, false).foreach { useDictionary => + runBenchmark("pmod integer filter") { + integerPmodBenchmark(TOTAL, IntegerType, useDictionary) + } + runBenchmark("pmod long filter") { + integerPmodBenchmark(TOTAL, LongType, useDictionary) + } + runBenchmark("pmod double filter") { + integerPmodBenchmark(TOTAL, DoubleType, useDictionary) + } + for ((precision, scale) <- Seq((18, 2), (38, 10))) { + runBenchmark("pmod decimal filter") { + decimalPmodBenchmark(TOTAL, DecimalType(precision, scale), useDictionary) + } + } + } + } +}