adding ansi checks and code refactor

apache · May 7, 2024 · 71501f5 · 71501f5
1 parent fd84f72
commit 71501f5
Showing 1 changed file with 49 additions and 39 deletions.
diff --git a/core/src/execution/datafusion/expressions/cast.rs b/core/src/execution/datafusion/expressions/cast.rs
@@ -334,45 +334,7 @@ impl Cast {
             }
 
             (DataType::Float64, DataType::Decimal128(precision, scale)) => {
-                let input = array.as_any().downcast_ref::<Float64Array>().unwrap();
-                let mut cast_array = PrimitiveArray::<Decimal128Type>::builder(input.len());
-
-                let mul = (*precision as f64).powi(*scale as i32);
-
-                for i in 0..input.len() {
-                    if input.is_null(i) {
-                        cast_array.append_null();
-                    } else {
-                        let input_value = input.value(i);
-                        let value = (input_value * mul).round().to_i128();
-
-                        if let Some(value) = value {
-                            if Decimal128Type::validate_decimal_precision(value, *precision).is_err() {
-                                return Err(
-                                    CometError::NumericValueOutOfRange {
-                                        value: input_value.to_string(),
-                                        precision: *precision,
-                                        scale: *scale,
-                                    }
-                                    .into(),
-                                );
-                            }
-
-                            cast_array.append_value(value);
-                        } else {
-                            return Err(
-                                CometError::NumericValueOutOfRange {
-                                    value: input_value.to_string(),
-                                    precision: *precision,
-                                    scale: *scale,
-                                }
-                                .into(),
-                            );
-                        }
-                    }
-                }
-
-                Arc::new(cast_array.with_precision_and_scale(*precision, *scale)?.finish()) as ArrayRef
+                Self::cast_float64_to_decimal128(&array, *precision, *scale, self.eval_mode)?
             }            
             _ => {
                 // when we have no Spark-specific casting we delegate to DataFusion
@@ -437,6 +399,54 @@ impl Cast {
         Ok(cast_array)
     }
 
+    fn cast_float64_to_decimal128 (
+        array: &dyn Array,
+        precision: u8,
+        scale: i8,
+        eval_mode: EvalMode,
+    ) -> CometResult<ArrayRef> {
+        let input = array.as_any().downcast_ref::<Float64Array>().unwrap();
+        let mut cast_array = PrimitiveArray::<Decimal128Type>::builder(input.len());
+
+        let mul = (precision as f64).powi(scale as i32);
+
+        for i in 0..input.len() {
+            if input.is_null(i) {
+                cast_array.append_null();
+            } else {
+                let input_value = input.value(i);    
+                let value = (input_value * mul).round().to_i128();
+
+                match value {
+                    Some(v) => {
+                        if Decimal128Type::validate_decimal_precision(v, precision).is_err() {
+                            return Err(CometError::NumericValueOutOfRange {
+                                value: input_value.to_string(),
+                                precision,
+                                scale,
+                            });
+                        }
+                        cast_array.append_value(v);
+                    }
+                    None => {
+                        if eval_mode == EvalMode::Ansi {
+                            return Err(CometError::NumericValueOutOfRange {
+                                value: input_value.to_string(),
+                                precision,
+                                scale,
+                            });
+                        } else {
+                            cast_array.append_null();
+                        }
+                    }
+                }
+            }
+        }        
+
+        let res = Arc::new(cast_array.with_precision_and_scale(precision, scale)?.finish()) as ArrayRef;
+        Ok(res)
+    }        
+
     fn spark_cast_float64_to_utf8<OffsetSize>(
         from: &dyn Array,
         _eval_mode: EvalMode,