feat: Use Datafusion's Sha2 and remove Comet's implementation. (#2063)

Author: rishvin

native/core/src/execution/jni_api.rs

@@ -40,6 +40,7 @@ use datafusion::{
     prelude::{SessionConfig, SessionContext},
 };
 use datafusion_comet_proto::spark_operator::Operator;
+use datafusion_spark::function::hash::sha2::SparkSha2;
 use datafusion_spark::function::math::expm1::SparkExpm1;
 use futures::poll;
 use futures::stream::StreamExt;
@@ -288,6 +289,7 @@ fn prepare_datafusion_session_context(
 
     // register UDFs from datafusion-spark crate
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkExpm1::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha2::default()));
 
     // Must be the last one to override existing functions with the same name
     datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?;

native/spark-expr/src/comet_scalar_funcs.rs

@@ -128,22 +128,6 @@ pub fn create_comet_physical_fun(
             let func = Arc::new(spark_isnan);
             make_comet_scalar_udf!("isnan", func, without data_type)
         }
-        "sha224" => {
-            let func = Arc::new(spark_sha224);
-            make_comet_scalar_udf!("sha224", func, without data_type)
-        }
-        "sha256" => {
-            let func = Arc::new(spark_sha256);
-            make_comet_scalar_udf!("sha256", func, without data_type)
-        }
-        "sha384" => {
-            let func = Arc::new(spark_sha384);
-            make_comet_scalar_udf!("sha384", func, without data_type)
-        }
-        "sha512" => {
-            let func = Arc::new(spark_sha512);
-            make_comet_scalar_udf!("sha512", func, without data_type)
-        }
         "date_add" => {
             let func = Arc::new(spark_date_add);
             make_comet_scalar_udf!("date_add", func, without data_type)

native/spark-expr/src/hash_funcs/mod.rs

@@ -16,10 +16,8 @@
 // under the License.
 
 pub mod murmur3;
-mod sha2;
 pub(super) mod utils;
 mod xxhash64;
 
 pub use murmur3::spark_murmur3_hash;
-pub use sha2::{spark_sha224, spark_sha256, spark_sha384, spark_sha512};
 pub use xxhash64::spark_xxhash64;

native/spark-expr/src/hash_funcs/sha2.rs

@@ -51,11 +51,6 @@ fn hex_encode<T: AsRef<[u8]>>(data: T, lower_case: bool) -> String {
     s
 }
 
-#[inline(always)]
-pub(crate) fn hex_strings<T: AsRef<[u8]>>(data: T) -> String {
-    hex_encode(data, true)
-}
-
 #[inline(always)]
 fn hex_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<String, std::fmt::Error> {
     let hex_string = hex_encode(bytes, false);

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

@@ -122,6 +122,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[Lower] -> CometLower,
     classOf[Murmur3Hash] -> CometMurmur3Hash,
     classOf[XxHash64] -> CometXxHash64,
+    classOf[Sha2] -> CometSha2,
     classOf[MapKeys] -> CometMapKeys,
     classOf[MapEntries] -> CometMapEntries,
     classOf[MapValues] -> CometMapValues,
@@ -1542,29 +1543,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
           None
         }
 
-      case Sha2(left, numBits) =>
-        if (!numBits.foldable) {
-          withInfo(expr, "non literal numBits is not supported")
-          return None
-        }
-        // it's possible for spark to dynamically compute the number of bits from input
-        // expression, however DataFusion does not support that yet.
-        val childExpr = exprToProtoInternal(left, inputs, binding)
-        val bits = numBits.eval().asInstanceOf[Int]
-        val algorithm = bits match {
-          case 224 => "sha224"
-          case 256 | 0 => "sha256"
-          case 384 => "sha384"
-          case 512 => "sha512"
-          case _ =>
-            null
-        }
-        if (algorithm == null) {
-          exprToProtoInternal(Literal(null, StringType), inputs, binding)
-        } else {
-          scalarFunctionExprToProtoWithReturnType(algorithm, StringType, childExpr)
-        }
-
       case struct @ CreateNamedStruct(_) =>
         if (struct.names.length != struct.names.distinct.length) {
           withInfo(expr, "CreateNamedStruct with duplicate field names are not supported")

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -19,8 +19,8 @@
 
 package org.apache.comet.serde
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, XxHash64}
-import org.apache.spark.sql.types.{DecimalType, IntegerType, LongType}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha2, XxHash64}
+import org.apache.spark.sql.types.{DecimalType, IntegerType, LongType, StringType}
 
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType}
@@ -65,6 +65,29 @@ object CometMurmur3Hash extends CometExpressionSerde {
   }
 }
 
+object CometSha2 extends CometExpressionSerde {
+  override def convert(
+      expr: Expression,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr] = {
+    if (!HashUtils.isSupportedType(expr)) {
+      return None
+    }
+
+    // It's possible for spark to dynamically compute the number of bits from input
+    // expression, however DataFusion does not support that yet.
+    val sha2Expr = expr.asInstanceOf[Sha2]
+    if (!sha2Expr.right.foldable) {
+      withInfo(expr, "For Sha2, non-foldable right argument is not supported")
+      return None
+    }
+
+    val leftExpr = exprToProtoInternal(sha2Expr.left, inputs, binding)
+    val numBitsExpr = exprToProtoInternal(sha2Expr.right, inputs, binding)
+    scalarFunctionExprToProtoWithReturnType("sha2", StringType, leftExpr, numBitsExpr)
+  }
+}
+
 private object HashUtils {
   def isSupportedType(expr: Expression): Boolean = {
     for (child <- expr.children) {

spark/src/main/scala/org/apache/comet/serde/hash.scala

@@ -1865,7 +1865,7 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               |md5(col), md5(cast(a as string)), md5(cast(b as string)),
               |hash(col), hash(col, 1), hash(col, 0), hash(col, a, b), hash(b, a, col),
               |xxhash64(col), xxhash64(col, 1), xxhash64(col, 0), xxhash64(col, a, b), xxhash64(b, a, col),
-              |sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128)
+              |sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128), sha2(col, -1)
               |from test
               |""".stripMargin)
         }
@@ -1977,7 +1977,7 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               |md5(col), md5(cast(a as string)), --md5(cast(b as string)),
               |hash(col), hash(col, 1), hash(col, 0), hash(col, a, b), hash(b, a, col),
               |xxhash64(col), xxhash64(col, 1), xxhash64(col, 0), xxhash64(col, a, b), xxhash64(b, a, col),
-              |sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128)
+              |sha2(col, 0), sha2(col, 256), sha2(col, 224), sha2(col, 384), sha2(col, 512), sha2(col, 128), sha2(col, -1)
               |from test
               |""".stripMargin)
         }