fix: support for spark-compatible null seed

apache · Dec 24, 2024 · 7e4ca2c · 7e4ca2c
1 parent c5c80e2
commit 7e4ca2c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/native/spark-expr/src/rand.rs b/native/spark-expr/src/rand.rs
@@ -87,8 +87,8 @@ impl RandExpr {
     }
 
     fn extract_init_state(seed: ScalarValue) -> Result<i64> {
-        if let ScalarValue::Int64(Some(init_seed)) = seed.cast_to(&DataType::Int64)? {
-            Ok(init_seed)
+        if let ScalarValue::Int64(seed_opt) = seed.cast_to(&DataType::Int64)? {
+            Ok(seed_opt.unwrap_or(0))
         } else {
             Err(DataFusionError::Internal(
                 "unexpected execution branch".to_string(),

diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -2531,6 +2531,9 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         val dfWithOverflowSeed =
           df.repartition(partitionsNumber).withColumn("rnd", rand(Long.MaxValue))
         checkSparkAnswer(dfWithOverflowSeed)
+        val dfWithNullSeed =
+          df.repartition(partitionsNumber).selectExpr("_1", "rand(null) as rnd")
+        checkSparkAnswer(dfWithNullSeed)
       }
     }
   }