refactor_boolean_cast_ops_add_benchmarks_rebase_main

Author: coderfender

native/spark-expr/Cargo.toml

@@ -97,5 +97,5 @@ name = "test_udf_registration"
 path = "tests/spark_expr_reg.rs"
 
 [[bench]]
-name = "cast_boolean"
+name = "cast_from_boolean"
 harness = false

native/spark-expr/benches/cast_boolean.rs …/spark-expr/benches/cast_from_boolean.rsnative/spark-expr/benches/cast_boolean.rs renamed to native/spark-expr/benches/cast_from_boolean.rs native/spark-expr/benches/cast_boolean.rs renamed to native/spark-expr/benches/cast_from_boolean.rs

@@ -34,9 +34,10 @@ fn criterion_benchmark(c: &mut Criterion) {
     let cast_to_i64 = Cast::new(expr.clone(), DataType::Int64, spark_cast_options.clone());
     let cast_to_f32 = Cast::new(expr.clone(), DataType::Float32, spark_cast_options.clone());
     let cast_to_f64 = Cast::new(expr.clone(), DataType::Float64, spark_cast_options.clone());
-    let cast_to_str = Cast::new(expr, DataType::Utf8, spark_cast_options);
+    let cast_to_str = Cast::new(expr.clone(), DataType::Utf8, spark_cast_options.clone());
+    let cast_to_decimal = Cast::new(expr, DataType::Decimal128(10, 4), spark_cast_options);
 
-    let mut group = c.benchmark_group("cast_bool_to_int".to_string());
+    let mut group = c.benchmark_group("cast_bool".to_string());
     group.bench_function("i8", |b| {
         b.iter(|| cast_to_i8.evaluate(&boolean_batch).unwrap());
     });
@@ -58,6 +59,9 @@ fn criterion_benchmark(c: &mut Criterion) {
     group.bench_function("str", |b| {
         b.iter(|| cast_to_str.evaluate(&boolean_batch).unwrap());
     });
+    group.bench_function("decimal", |b| {
+        b.iter(|| cast_to_decimal.evaluate(&boolean_batch).unwrap());
+    });
 }
 
 fn create_boolean_batch() -> RecordBatch {

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

@@ -15,7 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::SparkResult;
+use arrow::array::{ArrayRef, AsArray, Decimal128Array};
 use arrow::datatypes::DataType;
+use std::sync::Arc;
 
 pub fn can_cast_from_boolean(to_type: &DataType) -> bool {
     use DataType::*;
@@ -25,6 +28,21 @@ pub fn can_cast_from_boolean(to_type: &DataType) -> bool {
     )
 }
 
+// only incompatible boolean cast
+pub fn cast_boolean_to_decimal(
+    array: &ArrayRef,
+    precision: u8,
+    scale: i8,
+) -> SparkResult<ArrayRef> {
+    let bool_array = array.as_boolean();
+    let scaled_val = 10_i128.pow(scale as u32);
+    let result: Decimal128Array = bool_array
+        .iter()
+        .map(|v| v.map(|b| if b { scaled_val } else { 0 }))
+        .collect();
+    Ok(Arc::new(result.with_precision_and_scale(precision, scale)?))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -34,6 +52,7 @@ mod tests {
         Array, ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
         Int64Array, Int8Array, StringArray,
     };
+    use arrow::datatypes::DataType::Decimal128;
     use std::sync::Arc;
 
     fn test_input_bool_array() -> ArrayRef {
@@ -54,6 +73,7 @@ mod tests {
         assert!(can_cast_from_boolean(&DataType::Float32));
         assert!(can_cast_from_boolean(&DataType::Float64));
         assert!(can_cast_from_boolean(&DataType::Utf8));
+        assert!(can_cast_from_boolean(&DataType::Decimal128(10, 4)));
         assert!(!can_cast_from_boolean(&DataType::Null));
     }
 
@@ -154,4 +174,21 @@ mod tests {
         assert_eq!(arr.value(1), "false");
         assert!(arr.is_null(2));
     }
+
+    #[test]
+    fn test_bool_to_decimal_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &Decimal128(10, 4),
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let expected_arr = Decimal128Array::from(vec![10000_i128, 0_i128])
+            .with_precision_and_scale(10, 4)
+            .unwrap();
+        let arr = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
+        assert_eq!(arr.value(0), expected_arr.value(0));
+        assert_eq!(arr.value(1), expected_arr.value(1));
+        assert!(arr.is_null(2));
+    }
 }

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

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::conversion_funcs::boolean::can_cast_from_boolean;
+use crate::conversion_funcs::boolean::{can_cast_from_boolean, cast_boolean_to_decimal};
 use crate::conversion_funcs::utils::{cast_overflow, invalid_value};
 use crate::conversion_funcs::utils::{is_identity_cast, spark_cast_postprocess};
 use crate::utils::array_with_timezone;
@@ -1187,16 +1187,6 @@ fn cast_date_to_timestamp(
     ))
 }
 
-fn cast_boolean_to_decimal(array: &ArrayRef, precision: u8, scale: i8) -> SparkResult<ArrayRef> {
-    let bool_array = array.as_boolean();
-    let scaled_val = 10_i128.pow(scale as u32);
-    let result: Decimal128Array = bool_array
-        .iter()
-        .map(|v| v.map(|b| if b { scaled_val } else { 0 }))
-        .collect();
-    Ok(Arc::new(result.with_precision_and_scale(precision, scale)?))
-}
-
 fn cast_string_to_float(
     array: &ArrayRef,
     to_type: &DataType,