fix: Modulus on decimal data type mismatch (#2922)

Author: andygrove

native/spark-expr/src/math_funcs/modulo_expr.rs

@@ -100,10 +100,16 @@ pub fn create_modulo_expr(
                 SparkCastOptions::new_without_timezone(EvalMode::Legacy, false),
             ));
 
+            // The UDF's return type must match what Arrow's rem function will actually return.
+            // Since we're operating on Decimal256 inputs, rem will return Decimal256.
+            let decimal256_return_type = match &data_type {
+                DataType::Decimal128(p, s) => DataType::Decimal256(*p, *s),
+                other => other.clone(),
+            };
             let modulo_scalar_func = create_modulo_scalar_function(
                 left_256,
                 right_256,
-                &data_type,
+                &decimal256_return_type,
                 registry,
                 fail_on_error,
             )?;

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -1772,6 +1772,22 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("Decimal modulus with Decimal256 intermediate type") {
+    // regression test for https://github.com/apache/datafusion-comet/issues/2911
+    withTable("test") {
+      sql("create table test(a decimal(33, 29), b decimal(28, 17)) using parquet")
+      sql(
+        "insert into test values (-6788.53035340376888409034576923353, " +
+          "70948216565.90127985418365471)")
+      withSQLConf(
+        "spark.comet.enabled" -> "true",
+        "spark.sql.decimalOperations.allowPrecisionLoss" -> "true") {
+        val df = sql("select a, b, a % b from test")
+        df.collect()
+      }
+    }
+  }
+
   test("Decimal random number tests") {
     val rand = new scala.util.Random(42)
     def makeNum(p: Int, s: Int): String = {