int_to_binary_boolean_to_decimal

Author: coderfender

native/spark-expr/src/conversion_funcs/cast.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::utils::array_with_timezone;
+use crate::EvalMode::Legacy;
 use crate::{timezone, BinaryOutputStyle};
 use crate::{EvalMode, SparkError, SparkResult};
 use arrow::array::builder::StringBuilder;
@@ -25,8 +26,8 @@ use arrow::array::{
 };
 use arrow::compute::can_cast_types;
 use arrow::datatypes::{
-    i256, ArrowDictionaryKeyType, ArrowNativeType, DataType, Decimal256Type,
-    GenericBinaryType, Schema,
+    i256, ArrowDictionaryKeyType, ArrowNativeType, DataType, Decimal256Type, GenericBinaryType,
+    Schema,
 };
 use arrow::{
     array::{
@@ -66,7 +67,6 @@ use std::{
     num::Wrapping,
     sync::Arc,
 };
-use crate::EvalMode::Legacy;
 
 static TIMESTAMP_FORMAT: Option<&str> = Some("%Y-%m-%d %H:%M:%S%.f");
 
@@ -305,7 +305,10 @@ fn can_cast_from_timestamp(to_type: &DataType, _options: &SparkCastOptions) -> b
 
 fn can_cast_from_boolean(to_type: &DataType, _: &SparkCastOptions) -> bool {
     use DataType::*;
-    matches!(to_type, Int8 | Int16 | Int32 | Int64 | Float32 | Float64)
+    matches!(
+        to_type,
+        Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Decimal128(_, _)
+    )
 }
 
 fn can_cast_from_byte(to_type: &DataType, _: &SparkCastOptions) -> bool {
@@ -1125,9 +1128,18 @@ fn cast_array(
         }
         (Binary, Utf8) => Ok(cast_binary_to_string::<i32>(&array, cast_options)?),
         (Int8, Binary) if (eval_mode == Legacy) => cast_whole_num_to_binary!(&array, Int8Array, 1),
-        (Int16, Binary) if (eval_mode == Legacy)  => cast_whole_num_to_binary!(&array, Int16Array, 2),
-        (Int32, Binary) if (eval_mode == Legacy) => cast_whole_num_to_binary!(&array, Int32Array, 4),
-        (Int64, Binary) if (eval_mode == Legacy) => cast_whole_num_to_binary!(&array, Int64Array, 8),
+        (Int16, Binary) if (eval_mode == Legacy) => {
+            cast_whole_num_to_binary!(&array, Int16Array, 2)
+        }
+        (Int32, Binary) if (eval_mode == Legacy) => {
+            cast_whole_num_to_binary!(&array, Int32Array, 4)
+        }
+        (Int64, Binary) if (eval_mode == Legacy) => {
+            cast_whole_num_to_binary!(&array, Int64Array, 8)
+        }
+        (Boolean, Decimal128(precision, scale)) => {
+            cast_boolean_to_decimal(&array, *precision, *scale)
+        }
         _ if cast_options.is_adapting_schema
             || is_datafusion_spark_compatible(from_type, to_type) =>
         {
@@ -1146,6 +1158,16 @@ fn cast_array(
     Ok(spark_cast_postprocess(cast_result?, from_type, to_type))
 }
 
+fn cast_boolean_to_decimal(array: &ArrayRef, precision: u8, scale: i8) -> SparkResult<ArrayRef> {
+    let bool_array = array.as_boolean();
+    let scale_factor = 10_i128.pow(scale as u32);
+    let result: Decimal128Array = bool_array
+        .iter()
+        .map(|v| v.map(|b| if b { scale_factor } else { 0 }))
+        .collect();
+    Ok(Arc::new(result.with_precision_and_scale(precision, scale)?))
+}
+
 fn cast_string_to_float(
     array: &ArrayRef,
     to_type: &DataType,

spark/src/main/scala/org/apache/comet/expressions/CometCast.scala

@@ -21,7 +21,7 @@ package org.apache.comet.expressions
 
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, Literal}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, DataTypes, DecimalType, NullType, StructType}
+import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, NullType, StructType}
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
@@ -263,7 +263,7 @@ object CometCast extends CometExpressionSerde[Cast] with CometExprShim {
 
   private def canCastFromBoolean(toType: DataType): SupportLevel = toType match {
     case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType | DataTypes.LongType |
-        DataTypes.FloatType | DataTypes.DoubleType =>
+        DataTypes.FloatType | DataTypes.DoubleType | _: DecimalType =>
       Compatible()
     case _ => unsupported(DataTypes.BooleanType, toType)
   }

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

@@ -134,11 +134,18 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     castTest(generateBools(), DataTypes.DoubleType)
   }
 
-  ignore("cast BooleanType to DecimalType(10,2)") {
-    // Arrow error: Cast error: Casting from Boolean to Decimal128(10, 2) not supported
+  test("cast BooleanType to DecimalType(10,2)") {
     castTest(generateBools(), DataTypes.createDecimalType(10, 2))
   }
 
+  test("cast BooleanType to DecimalType(14,4)") {
+    castTest(generateBools(), DataTypes.createDecimalType(14, 4))
+  }
+
+  test("cast BooleanType to DecimalType(30,0)") {
+    castTest(generateBools(), DataTypes.createDecimalType(30, 0))
+  }
+
   test("cast BooleanType to StringType") {
     castTest(generateBools(), DataTypes.StringType)
   }
@@ -1366,9 +1373,11 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         }
 
         if (testTry) {
+          data.createOrReplaceTempView("t")
           // try_cast() should always return null for invalid inputs
+//          not using spark DSL since it `try_cast` is only available from Spark 4x
           val df2 =
-            data.select(col("a"), col("a").try_cast(toType)).orderBy(col("a"))
+            spark.sql(s"select a, try_cast(a as ${toType.sql}) from t order by a")
           if (hasIncompatibleType) {
             checkSparkAnswer(df2)
           } else {
@@ -1432,8 +1441,9 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
 
           // try_cast() should always return null for invalid inputs
           if (testTry) {
+            data.createOrReplaceTempView("t")
             val df2 =
-              data.select(col("a"), col("a").try_cast(toType)).orderBy(col("a"))
+              spark.sql(s"select a, try_cast(a as ${toType.sql}) from t order by a")
             if (hasIncompatibleType) {
               checkSparkAnswer(df2)
             } else {