adding from str

Author: Jiayu Liu

datafusion/src/physical_plan/windows.rs

@@ -38,19 +38,16 @@ pub enum WindowFunction {
 impl FromStr for WindowFunction {
     type Err = DataFusionError;
     fn from_str(name: &str) -> Result<WindowFunction> {
-        Ok(match name {
-            "min" => WindowFunction::AggregateFunction(AggregateFunction::Min),
-            "max" => WindowFunction::AggregateFunction(AggregateFunction::Max),
-            "count" => WindowFunction::AggregateFunction(AggregateFunction::Count),
-            "avg" => WindowFunction::AggregateFunction(AggregateFunction::Avg),
-            "sum" => WindowFunction::AggregateFunction(AggregateFunction::Sum),
-            _ => {
-                return Err(DataFusionError::Plan(format!(
-                    "There is no built-in function named {}",
-                    name
-                )))
-            }
-        })
+        if let Ok(aggregate) = AggregateFunction::from_str(name) {
+            Ok(WindowFunction::AggregateFunction(aggregate))
+        } else if let Ok(build_in_function) = BuiltInWindowFunction::from_str(name) {
+            Ok(WindowFunction::BuiltInWindowFunction(build_in_function))
+        } else {
+            Err(DataFusionError::Plan(format!(
+                "There is no built-in function named {}",
+                name
+            )))
+        }
     }
 }
 
@@ -80,6 +77,27 @@ pub enum BuiltInWindowFunction {
     LastValue,
 }
 
+impl FromStr for BuiltInWindowFunction {
+    type Err = DataFusionError;
+    fn from_str(name: &str) -> Result<BuiltInWindowFunction> {
+        Ok(match name {
+            "row_number" => BuiltInWindowFunction::RowNumber,
+            "rank" => BuiltInWindowFunction::Rank,
+            "dense_rank" => BuiltInWindowFunction::DenseRank,
+            "first_value" => BuiltInWindowFunction::FirstValue,
+            "last_value" => BuiltInWindowFunction::LastValue,
+            "lag" => BuiltInWindowFunction::Lag,
+            "lead" => BuiltInWindowFunction::Lead,
+            _ => {
+                return Err(DataFusionError::Plan(format!(
+                    "There is no built-in function named {}",
+                    name
+                )))
+            }
+        })
+    }
+}
+
 /// Returns the datatype of the scalar function
 pub fn return_type(fun: &WindowFunction, arg_types: &[DataType]) -> Result<DataType> {
     // Note that this function *must* return the same type that the respective physical expression returns
@@ -90,15 +108,15 @@ pub fn return_type(fun: &WindowFunction, arg_types: &[DataType]) -> Result<DataT
 
     match fun {
         WindowFunction::AggregateFunction(fun) => aggregates::return_type(fun, arg_types),
-        WindowFunction::BuiltInWindowFunction(fun) => Ok(match fun {
+        WindowFunction::BuiltInWindowFunction(fun) => match fun {
             BuiltInWindowFunction::RowNumber
             | BuiltInWindowFunction::Rank
-            | BuiltInWindowFunction::DenseRank => DataType::UInt64,
+            | BuiltInWindowFunction::DenseRank => Ok(DataType::UInt64),
             BuiltInWindowFunction::Lag
             | BuiltInWindowFunction::Lead
             | BuiltInWindowFunction::FirstValue
-            | BuiltInWindowFunction::LastValue => arg_types[0].clone(),
-        }),
+            | BuiltInWindowFunction::LastValue => Ok(arg_types[0].clone()),
+        },
     }
 }