refactor_boolean_cast_ops_add_tests

Author: B Vadlamani

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

@@ -21,6 +21,137 @@ pub fn can_cast_from_boolean(to_type: &DataType) -> bool {
     use DataType::*;
     matches!(
         to_type,
-        Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Utf8
+        Boolean | Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Utf8
     )
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::cast::cast_array;
+    use crate::{EvalMode, SparkCastOptions};
+    use arrow::array::{
+        Array, ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
+        Int64Array, Int8Array, StringArray,
+    };
+    use std::sync::Arc;
+
+    fn test_input_bool_array() -> ArrayRef {
+        Arc::new(BooleanArray::from(vec![Some(true), Some(false), None]))
+    }
+
+    fn test_input_spark_opts() -> SparkCastOptions {
+        SparkCastOptions::new(EvalMode::Legacy, "Asia/Kolkata", false)
+    }
+
+    #[test]
+    fn test_can_cast_from_boolean() {
+        assert!(can_cast_from_boolean(&DataType::Boolean));
+        assert!(can_cast_from_boolean(&DataType::Int8));
+        assert!(can_cast_from_boolean(&DataType::Int16));
+        assert!(can_cast_from_boolean(&DataType::Int32));
+        assert!(can_cast_from_boolean(&DataType::Int64));
+        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::Null));
+    }
+
+    #[test]
+    fn test_bool_to_int8_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Int8,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<Int8Array>().unwrap();
+        assert_eq!(arr.value(0), 1);
+        assert_eq!(arr.value(1), 0);
+        assert!(arr.is_null(2));
+    }
+
+    #[test]
+    fn test_bool_to_int16_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Int16,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<Int16Array>().unwrap();
+        assert_eq!(arr.value(0), 1);
+        assert_eq!(arr.value(1), 0);
+        assert!(arr.is_null(2));
+    }
+
+    #[test]
+    fn test_bool_to_int32_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Int32,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(arr.value(0), 1);
+        assert_eq!(arr.value(1), 0);
+        assert!(arr.is_null(2));
+    }
+
+    #[test]
+    fn test_bool_to_int64_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Int64,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<Int64Array>().unwrap();
+        assert_eq!(arr.value(0), 1);
+        assert_eq!(arr.value(1), 0);
+        assert!(arr.is_null(2));
+    }
+
+    #[test]
+    fn test_bool_to_float32_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Float32,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<Float32Array>().unwrap();
+        assert_eq!(arr.value(0), 1.0);
+        assert_eq!(arr.value(1), 0.0);
+        assert!(arr.is_null(2));
+    }
+
+    #[test]
+    fn test_bool_to_float64_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Float64,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<Float64Array>().unwrap();
+        assert_eq!(arr.value(0), 1.0);
+        assert_eq!(arr.value(1), 0.0);
+        assert!(arr.is_null(2));
+    }
+
+    #[test]
+    fn test_bool_to_string_cast() {
+        let result = cast_array(
+            test_input_bool_array(),
+            &DataType::Utf8,
+            &test_input_spark_opts(),
+        )
+        .unwrap();
+        let arr = result.as_any().downcast_ref::<StringArray>().unwrap();
+        assert_eq!(arr.value(0), "true");
+        assert_eq!(arr.value(1), "false");
+        assert!(arr.is_null(2));
+    }
+}

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

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::conversion_funcs::boolean::can_cast_from_boolean;
-use crate::conversion_funcs::utils::spark_cast_postprocess;
+use crate::conversion_funcs::utils::{is_identity_cast, spark_cast_postprocess};
 use crate::utils::array_with_timezone;
 use crate::{timezone, BinaryOutputStyle};
 use crate::{EvalMode, SparkError, SparkResult};
@@ -176,34 +176,31 @@ pub fn cast_supported(
         to_type
     };
 
-    if from_type == to_type {
-        return true;
-    }
-
-    match (from_type, to_type) {
-        (Boolean, _) => can_cast_from_boolean(to_type),
-        (UInt8 | UInt16 | UInt32 | UInt64, Int8 | Int16 | Int32 | Int64)
-            if options.allow_cast_unsigned_ints =>
-        {
-            true
+    is_identity_cast(from_type, to_type)
+        || match (from_type, to_type) {
+            (Boolean, _) => can_cast_from_boolean(to_type),
+            (UInt8 | UInt16 | UInt32 | UInt64, Int8 | Int16 | Int32 | Int64)
+                if options.allow_cast_unsigned_ints =>
+            {
+                true
+            }
+            (Int8, _) => can_cast_from_byte(to_type, options),
+            (Int16, _) => can_cast_from_short(to_type, options),
+            (Int32, _) => can_cast_from_int(to_type, options),
+            (Int64, _) => can_cast_from_long(to_type, options),
+            (Float32, _) => can_cast_from_float(to_type, options),
+            (Float64, _) => can_cast_from_double(to_type, options),
+            (Decimal128(p, s), _) => can_cast_from_decimal(p, s, to_type, options),
+            (Timestamp(_, None), _) => can_cast_from_timestamp_ntz(to_type, options),
+            (Timestamp(_, Some(_)), _) => can_cast_from_timestamp(to_type, options),
+            (Utf8 | LargeUtf8, _) => can_cast_from_string(to_type, options),
+            (_, Utf8 | LargeUtf8) => can_cast_to_string(from_type, options),
+            (Struct(from_fields), Struct(to_fields)) => from_fields
+                .iter()
+                .zip(to_fields.iter())
+                .all(|(a, b)| cast_supported(a.data_type(), b.data_type(), options)),
+            _ => false,
         }
-        (Int8, _) => can_cast_from_byte(to_type, options),
-        (Int16, _) => can_cast_from_short(to_type, options),
-        (Int32, _) => can_cast_from_int(to_type, options),
-        (Int64, _) => can_cast_from_long(to_type, options),
-        (Float32, _) => can_cast_from_float(to_type, options),
-        (Float64, _) => can_cast_from_double(to_type, options),
-        (Decimal128(p, s), _) => can_cast_from_decimal(p, s, to_type, options),
-        (Timestamp(_, None), _) => can_cast_from_timestamp_ntz(to_type, options),
-        (Timestamp(_, Some(_)), _) => can_cast_from_timestamp(to_type, options),
-        (Utf8 | LargeUtf8, _) => can_cast_from_string(to_type, options),
-        (_, Utf8 | LargeUtf8) => can_cast_to_string(from_type, options),
-        (Struct(from_fields), Struct(to_fields)) => from_fields
-            .iter()
-            .zip(to_fields.iter())
-            .all(|(a, b)| cast_supported(a.data_type(), b.data_type(), options)),
-        _ => false,
-    }
 }
 
 fn can_cast_from_string(to_type: &DataType, options: &SparkCastOptions) -> bool {
@@ -947,7 +944,7 @@ fn dict_from_values<K: ArrowDictionaryKeyType>(
     Ok(Arc::new(dict_array))
 }
 
-fn cast_array(
+pub fn cast_array(
     array: ArrayRef,
     to_type: &DataType,
     cast_options: &SparkCastOptions,
@@ -1303,16 +1300,26 @@ fn cast_binary_formatter(value: &[u8]) -> String {
 /// Determines if DataFusion supports the given cast in a way that is
 /// compatible with Spark
 fn is_datafusion_spark_compatible(from_type: &DataType, to_type: &DataType) -> bool {
-    if from_type == to_type {
-        return true;
-    }
-    match from_type {
-        DataType::Null => {
-            matches!(to_type, DataType::List(_))
-        }
-        DataType::Boolean => can_cast_from_boolean(to_type),
-        DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
-            matches!(
+    is_identity_cast(from_type, to_type)
+        || match from_type {
+            DataType::Null => {
+                matches!(to_type, DataType::List(_))
+            }
+            DataType::Boolean => can_cast_from_boolean(to_type),
+            DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
+                matches!(
+                    to_type,
+                    DataType::Boolean
+                        | DataType::Int8
+                        | DataType::Int16
+                        | DataType::Int32
+                        | DataType::Int64
+                        | DataType::Float32
+                        | DataType::Float64
+                        | DataType::Utf8
+                )
+            }
+            DataType::Float32 | DataType::Float64 => matches!(
                 to_type,
                 DataType::Boolean
                     | DataType::Int8
@@ -1321,46 +1328,34 @@ fn is_datafusion_spark_compatible(from_type: &DataType, to_type: &DataType) -> b
                     | DataType::Int64
                     | DataType::Float32
                     | DataType::Float64
-                    | DataType::Utf8
-            )
-        }
-        DataType::Float32 | DataType::Float64 => matches!(
-            to_type,
-            DataType::Boolean
-                | DataType::Int8
-                | DataType::Int16
-                | DataType::Int32
-                | DataType::Int64
-                | DataType::Float32
-                | DataType::Float64
-        ),
-        DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => matches!(
-            to_type,
-            DataType::Int8
-                | DataType::Int16
-                | DataType::Int32
-                | DataType::Int64
-                | DataType::Float32
-                | DataType::Float64
-                | DataType::Decimal128(_, _)
-                | DataType::Decimal256(_, _)
-                | DataType::Utf8 // note that there can be formatting differences
-        ),
-        DataType::Utf8 => matches!(to_type, DataType::Binary),
-        DataType::Date32 => matches!(to_type, DataType::Int32 | DataType::Utf8),
-        DataType::Timestamp(_, _) => {
-            matches!(
+            ),
+            DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => matches!(
                 to_type,
-                DataType::Int64 | DataType::Date32 | DataType::Utf8 | DataType::Timestamp(_, _)
-            )
-        }
-        DataType::Binary => {
-            // note that this is not completely Spark compatible because
-            // DataFusion only supports binary data containing valid UTF-8 strings
-            matches!(to_type, DataType::Utf8)
+                DataType::Int8
+                    | DataType::Int16
+                    | DataType::Int32
+                    | DataType::Int64
+                    | DataType::Float32
+                    | DataType::Float64
+                    | DataType::Decimal128(_, _)
+                    | DataType::Decimal256(_, _)
+                    | DataType::Utf8 // note that there can be formatting differences
+            ),
+            DataType::Utf8 => matches!(to_type, DataType::Binary),
+            DataType::Date32 => matches!(to_type, DataType::Int32 | DataType::Utf8),
+            DataType::Timestamp(_, _) => {
+                matches!(
+                    to_type,
+                    DataType::Int64 | DataType::Date32 | DataType::Utf8 | DataType::Timestamp(_, _)
+                )
+            }
+            DataType::Binary => {
+                // note that this is not completely Spark compatible because
+                // DataFusion only supports binary data containing valid UTF-8 strings
+                matches!(to_type, DataType::Utf8)
+            }
+            _ => false,
         }
-        _ => false,
-    }
 }
 
 /// Cast between struct types based on logic in

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

@@ -107,3 +107,8 @@ fn trim_end(s: &str) -> &str {
         s
     }
 }
+
+#[inline]
+pub fn is_identity_cast(from_type: &DataType, to_type: &DataType) -> bool {
+    from_type == to_type
+}