Introduce INFORMATION_SCHEMA.ROUTINES table (apache#13255)

* tmp

* introduce routines table

* add is_deterministic field

* cargo fmt

* rollback the session_state changed

Author: goldmedal

datafusion/physical-expr/src/expressions/case.rs

@@ -19,19 +19,19 @@ use std::borrow::Cow;
 use std::hash::Hash;
 use std::{any::Any, sync::Arc};
 
-use crate::expressions::try_cast;
+use crate::expressions::{try_cast, BinaryExpr, CastExpr};
 use crate::PhysicalExpr;
 
 use arrow::array::*;
+use arrow::compute::kernels::cmp::eq;
 use arrow::compute::kernels::zip::zip;
 use arrow::compute::{and, and_not, is_null, not, nullif, or, prep_null_mask_filter};
 use arrow::datatypes::{DataType, Schema};
 use datafusion_common::cast::as_boolean_array;
 use datafusion_common::{exec_err, internal_err, DataFusionError, Result, ScalarValue};
-use datafusion_expr::ColumnarValue;
+use datafusion_expr::{ColumnarValue, Operator};
 
 use super::{Column, Literal};
-use datafusion_physical_expr_common::datum::compare_with_eq;
 use itertools::Itertools;
 
 type WhenThen = (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>);
@@ -57,9 +57,14 @@ enum EvalMethod {
     InfallibleExprOrNull,
     /// This is a specialization for a specific use case where we can take a fast path
     /// if there is just one when/then pair and both the `then` and `else` expressions
-    /// are literal values
+    /// are literal values.
     /// CASE WHEN condition THEN literal ELSE literal END
     ScalarOrScalar,
+    /// This is a specialization for a sprcific use case where we can take a fast path
+    /// for the divide-by-zero expression when the divisor is zero.
+    ///
+    /// CASE WHEN y > 0 THEN x / y ELSE NULL END
+    DivideZeroExpression,
 }
 
 /// The CASE expression is similar to a series of nested if/else and there are two forms that
@@ -149,6 +154,51 @@ impl CaseExpr {
                 && else_expr.as_ref().unwrap().as_any().is::<Literal>()
             {
                 EvalMethod::ScalarOrScalar
+            } else if when_then_expr.len() == 1
+                && when_then_expr[0].0.as_any().is::<BinaryExpr>()
+            {
+                let b = when_then_expr[0]
+                    .0
+                    .as_any()
+                    .downcast_ref::<BinaryExpr>()
+                    .expect("expected binary expression");
+
+                if b.op().eq(&Operator::Gt) {
+                    if let Some(col) = b.left().as_any().downcast_ref::<Column>() {
+                        if let Some(lit) = b.right().as_any().downcast_ref::<Literal>() {
+                            if matches!(lit.value(), ScalarValue::Int32(Some(0))) {
+                                if let Some(b) = when_then_expr[0]
+                                    .1
+                                    .as_any()
+                                    .downcast_ref::<BinaryExpr>()
+                                {
+                                    if b.op().eq(&Operator::Divide) {
+                                        if let Some(cast) =
+                                            b.right().as_any().downcast_ref::<CastExpr>()
+                                        {
+                                            if let Some(col2) = cast
+                                                .expr()
+                                                .as_any()
+                                                .downcast_ref::<Column>()
+                                            {
+                                                if col.name() == col2.name() {
+                                                    return Ok(Self {
+                                                        expr: None,
+                                                        when_then_expr,
+                                                        else_expr,
+                                                        eval_method: EvalMethod::DivideZeroExpression,
+                                                    });
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+
+                EvalMethod::NoExpression
             } else {
                 EvalMethod::NoExpression
             };
@@ -203,13 +253,7 @@ impl CaseExpr {
                 .evaluate_selection(batch, &remainder)?;
             let when_value = when_value.into_array(batch.num_rows())?;
             // build boolean array representing which rows match the "when" value
-            let when_match = compare_with_eq(
-                &when_value,
-                &base_value,
-                // The types of case and when expressions will be coerced to match.
-                // We only need to check if the base_value is nested.
-                base_value.data_type().is_nested(),
-            )?;
+            let when_match = eq(&when_value, &base_value)?;
             // Treat nulls as false
             let when_match = match when_match.null_count() {
                 0 => Cow::Borrowed(&when_match),
@@ -385,12 +429,49 @@ impl CaseExpr {
 
         // keep `else_expr`'s data type and return type consistent
         let e = self.else_expr.as_ref().unwrap();
-        let expr = try_cast(Arc::clone(e), &batch.schema(), return_type)
+        let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
             .unwrap_or_else(|_| Arc::clone(e));
         let else_ = Scalar::new(expr.evaluate(batch)?.into_array(1)?);
 
         Ok(ColumnarValue::Array(zip(&when_value, &then_value, &else_)?))
     }
+
+    fn divide_by_zero_expr(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        let return_type = self.data_type(&batch.schema())?;
+
+        // start with nulls as default output
+        let mut current_value = new_null_array(&return_type, batch.num_rows());
+        let mut remainder = BooleanArray::from(vec![true; batch.num_rows()]);
+        let when_value = BooleanArray::from(vec![true; batch.num_rows()]);
+        let then_value = self.when_then_expr()[0].1.evaluate(batch)?;
+        current_value = match then_value {
+            ColumnarValue::Scalar(ScalarValue::Null) => {
+                nullif(current_value.as_ref(), &when_value)?
+            }
+            ColumnarValue::Scalar(then_value) => {
+                zip(&when_value, &then_value.to_scalar()?, &current_value)?
+            }
+            ColumnarValue::Array(then_value) => {
+                zip(&when_value, &then_value, &current_value)?
+            }
+        };
+
+        // Succeed tuples should be filtered out for short-circuit evaluation,
+        // null values for the current when expr should be kept
+        remainder = and_not(&remainder, &when_value)?;
+
+        if let Some(e) = &self.else_expr {
+            // keep `else_expr`'s data type and return type consistent
+            let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())
+                .unwrap_or_else(|_| Arc::clone(e));
+            let else_ = expr
+                .evaluate_selection(batch, &remainder)?
+                .into_array(batch.num_rows())?;
+            current_value = zip(&remainder, &else_, &current_value)?;
+        }
+
+        Ok(ColumnarValue::Array(current_value))
+    }
 }
 
 impl PhysicalExpr for CaseExpr {
@@ -454,6 +535,7 @@ impl PhysicalExpr for CaseExpr {
                 self.case_column_or_null(batch)
             }
             EvalMethod::ScalarOrScalar => self.scalar_or_scalar(batch),
+            EvalMethod::DivideZeroExpression => self.divide_by_zero_expr(batch),
         }
     }
 
@@ -741,6 +823,13 @@ mod tests {
         Ok(batch)
     }
 
+    fn case_test_batch2() -> Result<RecordBatch> {
+        let schema = Schema::new(vec![Field::new("y", DataType::Int32, true)]);
+        let a = Int32Array::from(vec![Some(1), Some(0), None, Some(5)]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
+        Ok(batch)
+    }
+
     #[test]
     fn case_without_expr_else() -> Result<()> {
         let batch = case_test_batch()?;
@@ -1212,4 +1301,58 @@ mod tests {
                 comparison_coercion(&left_type, right_type)
             })
     }
+
+    #[test]
+    fn gen_optimize_case_for_div_zero() -> Result<()> {
+        let batch = case_test_batch1()?;
+        let schema = batch.schema();
+
+        let batch2 = case_test_batch2()?;
+        let schema2 = batch2.schema();
+
+        // DivideZeroExpression: CASE WHEN a > 0 THEN 25.0 / cast(a, float64) ELSE float64(null) END
+        let when1 = binary(col("a", &schema)?, Operator::Gt, lit(0i32), &batch.schema())?;
+        let then1 = binary(
+            lit(25.0f64),
+            Operator::Divide,
+            cast(col("a", &schema)?, &batch.schema(), Float64)?,
+            &batch.schema(),
+        )?;
+        let x = lit(ScalarValue::Float64(None));
+        let expr = generate_case_when_with_type_coercion(
+            None,
+            vec![(when1, then1)],
+            Some(x),
+            schema.as_ref(),
+        )?;
+        let case = expr
+            .as_any()
+            .downcast_ref::<CaseExpr>()
+            .expect("expected case expression");
+        assert_eq!(case.eval_method, EvalMethod::DivideZeroExpression);
+
+        // NoExpression: CASE WHEN a > 0 THEN 25.0 / cast(y, float64) ELSE float64(null) END
+        let when1 = binary(col("a", &schema)?, Operator::Gt, lit(0i32), &batch.schema())?;
+        let then1 = binary(
+            lit(25.0f64),
+            Operator::Divide,
+            cast(col("y", &schema2)?, &batch2.schema(), Float64)?,
+            &batch2.schema(),
+        )?;
+        let x = lit(ScalarValue::Float64(None));
+
+        let expr = generate_case_when_with_type_coercion(
+            None,
+            vec![(when1, then1)],
+            Some(x),
+            schema.as_ref(),
+        )?;
+        let case = expr
+            .as_any()
+            .downcast_ref::<CaseExpr>()
+            .expect("expected case expression");
+        assert_eq!(case.eval_method, EvalMethod::NoExpression);
+
+        Ok(())
+    }
 }