ARROW-11822: [Rust][Datafusion] Support case sensitive for function

rust/datafusion/src/execution/context.rs

@@ -498,13 +498,34 @@ impl QueryPlanner for DefaultQueryPlanner {
     }
 }
 
+/// The style of case display
+#[derive(Clone, PartialEq)]
+pub enum CaseStyle {
+    /// function names in the output are displayed as they were provided:
+    /// 'mD5("a")' will appear in the results as `mD5("a")`. This
+    /// mimics MySQL behavior
+    PreserveCase,
+    /// function names in the output are displayed in lower case:
+    /// `mD5("a")` will appear in the results as  `md5("a")`This
+    /// mimics PostgreSQL behavior
+    LikePostgreSQL,
+}
+
 /// Configuration options for execution context
 #[derive(Clone)]
 pub struct ExecutionConfig {
     /// Number of concurrent threads for query execution.
     pub concurrency: usize,
     /// Default batch size when reading data sources
     pub batch_size: usize,
+    /// Whether to use case-sensitive matching for function names.
+    /// If `false` both `"SELECT COUNT(*) FROM t;` and "`SELECT count(*) FROM t;`
+    /// can be used to compute the `COUNT` aggregate. If `true` then only
+    /// `"SELECT count(*) FROM t"` can be used.
+    /// Defaults to `true`
+    pub case_sensitive: bool,
+    /// Default to like PostgreSQL
+    pub case_style: CaseStyle,
     /// Responsible for optimizing a logical plan
     optimizers: Vec<Arc<dyn OptimizerRule + Send + Sync>>,
     /// Responsible for planning `LogicalPlan`s, and `ExecutionPlan`
@@ -517,6 +538,8 @@ impl ExecutionConfig {
         Self {
             concurrency: num_cpus::get(),
             batch_size: 32768,
+            case_sensitive: true,
+            case_style: CaseStyle::LikePostgreSQL,
             optimizers: vec![
                 Arc::new(ConstantFolding::new()),
                 Arc::new(ProjectionPushDown::new()),
@@ -544,6 +567,18 @@ impl ExecutionConfig {
         self
     }
 
+    /// Customize case sensitive
+    pub fn with_case_sensitive(mut self, cs: bool) -> Self {
+        self.case_sensitive = cs;
+        self
+    }
+
+    /// Customize case style
+    pub fn with_case_style(mut self, cs: CaseStyle) -> Self {
+        self.case_style = cs;
+        self
+    }
+
     /// Replace the default query planner
     pub fn with_query_planner(
         mut self,
@@ -593,6 +628,10 @@ impl ContextProvider for ExecutionContextState {
     fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
         self.aggregate_functions.get(name).cloned()
     }
+
+    fn get_config(&self) -> ExecutionConfig {
+        self.config.clone()
+    }
 }
 
 impl FunctionRegistry for ExecutionContextState {

rust/datafusion/src/logical_plan/expr.rs

@@ -160,20 +160,26 @@ pub enum Expr {
     },
     /// Represents the call of a built-in scalar function with a set of arguments.
     ScalarFunction {
+        /// The input name of the function
+        input_name: String,
         /// The function
         fun: functions::BuiltinScalarFunction,
         /// List of expressions to feed to the functions as arguments
         args: Vec<Expr>,
     },
     /// Represents the call of a user-defined scalar function with arguments.
     ScalarUDF {
+        /// The input name of the function
+        input_name: String,
         /// The function
         fun: Arc<ScalarUDF>,
         /// List of expressions to feed to the functions as arguments
         args: Vec<Expr>,
     },
     /// Represents the call of an aggregate built-in function with arguments.
     AggregateFunction {
+        /// The input name of the function
+        input_name: String,
         /// Name of the function
         fun: aggregates::AggregateFunction,
         /// List of expressions to feed to the functions as arguments
@@ -183,6 +189,8 @@ pub enum Expr {
     },
     /// aggregate function
     AggregateUDF {
+        /// The input name of the function
+        input_name: String,
         /// The function
         fun: Arc<AggregateUDF>,
         /// List of expressions to feed to the functions as arguments
@@ -220,14 +228,14 @@ impl Expr {
             Expr::Literal(l) => Ok(l.get_datatype()),
             Expr::Case { when_then_expr, .. } => when_then_expr[0].1.get_type(schema),
             Expr::Cast { data_type, .. } => Ok(data_type.clone()),
-            Expr::ScalarUDF { fun, args } => {
+            Expr::ScalarUDF { fun, args, .. } => {
                 let data_types = args
                     .iter()
                     .map(|e| e.get_type(schema))
                     .collect::<Result<Vec<_>>>()?;
                 Ok((fun.return_type)(&data_types)?.as_ref().clone())
             }
-            Expr::ScalarFunction { fun, args } => {
+            Expr::ScalarFunction { fun, args, .. } => {
                 let data_types = args
                     .iter()
                     .map(|e| e.get_type(schema))
@@ -680,24 +688,42 @@ impl Expr {
                 asc,
                 nulls_first,
             },
-            Expr::ScalarFunction { args, fun } => Expr::ScalarFunction {
+            Expr::ScalarFunction {
+                input_name,
+                args,
+                fun,
+            } => Expr::ScalarFunction {
+                input_name,
                 args: rewrite_vec(args, rewriter)?,
                 fun,
             },
-            Expr::ScalarUDF { args, fun } => Expr::ScalarUDF {
+            Expr::ScalarUDF {
+                input_name,
+                args,
+                fun,
+                ..
+            } => Expr::ScalarUDF {
+                input_name,
                 args: rewrite_vec(args, rewriter)?,
                 fun,
             },
             Expr::AggregateFunction {
+                input_name,
                 args,
                 fun,
                 distinct,
             } => Expr::AggregateFunction {
+                input_name,
                 args: rewrite_vec(args, rewriter)?,
                 fun,
                 distinct,
             },
-            Expr::AggregateUDF { args, fun } => Expr::AggregateUDF {
+            Expr::AggregateUDF {
+                input_name,
+                args,
+                fun,
+            } => Expr::AggregateUDF {
+                input_name,
                 args: rewrite_vec(args, rewriter)?,
                 fun,
             },
@@ -928,6 +954,7 @@ pub fn col(name: &str) -> Expr {
 /// Create an expression to represent the min() aggregate function
 pub fn min(expr: Expr) -> Expr {
     Expr::AggregateFunction {
+        input_name: aggregates::AggregateFunction::Min.to_string(),
         fun: aggregates::AggregateFunction::Min,
         distinct: false,
         args: vec![expr],
@@ -937,6 +964,7 @@ pub fn min(expr: Expr) -> Expr {
 /// Create an expression to represent the max() aggregate function
 pub fn max(expr: Expr) -> Expr {
     Expr::AggregateFunction {
+        input_name: aggregates::AggregateFunction::Max.to_string(),
         fun: aggregates::AggregateFunction::Max,
         distinct: false,
         args: vec![expr],
@@ -946,6 +974,7 @@ pub fn max(expr: Expr) -> Expr {
 /// Create an expression to represent the sum() aggregate function
 pub fn sum(expr: Expr) -> Expr {
     Expr::AggregateFunction {
+        input_name: aggregates::AggregateFunction::Sum.to_string(),
         fun: aggregates::AggregateFunction::Sum,
         distinct: false,
         args: vec![expr],
@@ -955,6 +984,7 @@ pub fn sum(expr: Expr) -> Expr {
 /// Create an expression to represent the avg() aggregate function
 pub fn avg(expr: Expr) -> Expr {
     Expr::AggregateFunction {
+        input_name: aggregates::AggregateFunction::Avg.to_string(),
         fun: aggregates::AggregateFunction::Avg,
         distinct: false,
         args: vec![expr],
@@ -964,6 +994,7 @@ pub fn avg(expr: Expr) -> Expr {
 /// Create an expression to represent the count() aggregate function
 pub fn count(expr: Expr) -> Expr {
     Expr::AggregateFunction {
+        input_name: aggregates::AggregateFunction::Count.to_string(),
         fun: aggregates::AggregateFunction::Count,
         distinct: false,
         args: vec![expr],
@@ -973,6 +1004,7 @@ pub fn count(expr: Expr) -> Expr {
 /// Create an expression to represent the count(distinct) aggregate function
 pub fn count_distinct(expr: Expr) -> Expr {
     Expr::AggregateFunction {
+        input_name: aggregates::AggregateFunction::Count.to_string(),
         fun: aggregates::AggregateFunction::Count,
         distinct: true,
         args: vec![expr],
@@ -1046,6 +1078,7 @@ macro_rules! unary_scalar_expr {
         #[allow(missing_docs)]
         pub fn $FUNC(e: Expr) -> Expr {
             Expr::ScalarFunction {
+                input_name: functions::BuiltinScalarFunction::$ENUM.to_string(),
                 fun: functions::BuiltinScalarFunction::$ENUM,
                 args: vec![e],
             }
@@ -1113,6 +1146,7 @@ unary_scalar_expr!(Upper, upper);
 /// returns an array of fixed size with each argument on it.
 pub fn array(args: Vec<Expr>) -> Expr {
     Expr::ScalarFunction {
+        input_name: functions::BuiltinScalarFunction::Array.to_string(),
         fun: functions::BuiltinScalarFunction::Array,
         args,
     }
@@ -1330,24 +1364,26 @@ fn create_name(e: &Expr, input_schema: &DFSchema) -> Result<String> {
             let expr = create_name(expr, input_schema)?;
             Ok(format!("{} IS NOT NULL", expr))
         }
-        Expr::ScalarFunction { fun, args, .. } => {
-            create_function_name(&fun.to_string(), false, args, input_schema)
-        }
-        Expr::ScalarUDF { fun, args, .. } => {
-            create_function_name(&fun.name, false, args, input_schema)
-        }
+        Expr::ScalarFunction {
+            input_name, args, ..
+        } => create_function_name(input_name, false, args, input_schema),
+        Expr::ScalarUDF {
+            input_name, args, ..
+        } => create_function_name(input_name, false, args, input_schema),
         Expr::AggregateFunction {
-            fun,
+            input_name,
             distinct,
             args,
             ..
-        } => create_function_name(&fun.to_string(), *distinct, args, input_schema),
-        Expr::AggregateUDF { fun, args } => {
+        } => create_function_name(input_name, *distinct, args, input_schema),
+        Expr::AggregateUDF {
+            input_name, args, ..
+        } => {
             let mut names = Vec::with_capacity(args.len());
             for e in args {
                 names.push(create_name(e, input_schema)?);
             }
-            Ok(format!("{}({})", fun.name, names.join(",")))
+            Ok(format!("{}({})", input_name, names.join(",")))
         }
         Expr::InList {
             expr,

rust/datafusion/src/optimizer/utils.rs

@@ -295,20 +295,35 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> {
         }),
         Expr::IsNull(_) => Ok(Expr::IsNull(Box::new(expressions[0].clone()))),
         Expr::IsNotNull(_) => Ok(Expr::IsNotNull(Box::new(expressions[0].clone()))),
-        Expr::ScalarFunction { fun, .. } => Ok(Expr::ScalarFunction {
+        Expr::ScalarFunction {
+            input_name, fun, ..
+        } => Ok(Expr::ScalarFunction {
+            input_name: input_name.to_string(),
             fun: fun.clone(),
             args: expressions.to_vec(),
         }),
-        Expr::ScalarUDF { fun, .. } => Ok(Expr::ScalarUDF {
+        Expr::ScalarUDF {
+            input_name, fun, ..
+        } => Ok(Expr::ScalarUDF {
+            input_name: input_name.clone(),
             fun: fun.clone(),
             args: expressions.to_vec(),
         }),
-        Expr::AggregateFunction { fun, distinct, .. } => Ok(Expr::AggregateFunction {
+        Expr::AggregateFunction {
+            input_name,
+            fun,
+            distinct,
+            ..
+        } => Ok(Expr::AggregateFunction {
+            input_name: input_name.clone(),
             fun: fun.clone(),
             args: expressions.to_vec(),
             distinct: *distinct,
         }),
-        Expr::AggregateUDF { fun, .. } => Ok(Expr::AggregateUDF {
+        Expr::AggregateUDF {
+            input_name, fun, ..
+        } => Ok(Expr::AggregateUDF {
+            input_name: input_name.clone(),
             fun: fun.clone(),
             args: expressions.to_vec(),
         }),

rust/datafusion/src/physical_plan/planner.rs

@@ -554,14 +554,14 @@ impl DefaultPhysicalPlanner {
             Expr::IsNotNull(expr) => expressions::is_not_null(
                 self.create_physical_expr(expr, input_schema, ctx_state)?,
             ),
-            Expr::ScalarFunction { fun, args } => {
+            Expr::ScalarFunction { fun, args, .. } => {
                 let physical_args = args
                     .iter()
                     .map(|e| self.create_physical_expr(e, input_schema, ctx_state))
                     .collect::<Result<Vec<_>>>()?;
                 functions::create_physical_expr(fun, &physical_args, input_schema)
             }
-            Expr::ScalarUDF { fun, args } => {
+            Expr::ScalarUDF { fun, args, .. } => {
                 let mut physical_args = vec![];
                 for e in args {
                     physical_args.push(self.create_physical_expr(

rust/datafusion/src/physical_plan/udaf.rs

@@ -93,6 +93,7 @@ impl AggregateUDF {
     /// This utility allows using the UDAF without requiring access to the registry.
     pub fn call(&self, args: Vec<Expr>) -> Expr {
         Expr::AggregateUDF {
+            input_name: self.name.to_string(),
             fun: Arc::new(self.clone()),
             args,
         }

rust/datafusion/src/physical_plan/udf.rs

@@ -82,6 +82,7 @@ impl ScalarUDF {
     /// This utility allows using the UDF without requiring access to the registry.
     pub fn call(&self, args: Vec<Expr>) -> Expr {
         Expr::ScalarUDF {
+            input_name: self.name.to_string(),
             fun: Arc::new(self.clone()),
             args,
         }