refactored nth_value

Author: buraksenn

datafusion/core/src/dataframe/mod.rs

@@ -1946,9 +1946,8 @@ mod tests {
     use datafusion_common_runtime::SpawnedTask;
     use datafusion_expr::expr::WindowFunction;
     use datafusion_expr::{
-        cast, create_udf, lit, BuiltInWindowFunction, ExprFunctionExt,
-        ScalarFunctionImplementation, Volatility, WindowFrame, WindowFrameBound,
-        WindowFrameUnits, WindowFunctionDefinition,
+        cast, create_udf, lit, ExprFunctionExt, ScalarFunctionImplementation, Volatility,
+        WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
     };
     use datafusion_functions_aggregate::expr_fn::{array_agg, count_distinct};
     use datafusion_functions_window::expr_fn::row_number;
@@ -2172,31 +2171,6 @@ mod tests {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn select_with_window_exprs() -> Result<()> {
-        // build plan using Table API
-        let t = test_table().await?;
-        let first_row = Expr::WindowFunction(WindowFunction::new(
-            WindowFunctionDefinition::BuiltInWindowFunction(
-                BuiltInWindowFunction::FirstValue,
-            ),
-            vec![col("aggregate_test_100.c1")],
-        ))
-        .partition_by(vec![col("aggregate_test_100.c2")])
-        .build()
-        .unwrap();
-        let t2 = t.select(vec![col("c1"), first_row])?;
-        let plan = t2.plan.clone();
-
-        let sql_plan = create_plan(
-            "select c1, first_value(c1) over (partition by c2) from aggregate_test_100",
-        )
-        .await?;
-
-        assert_same_plan(&plan, &sql_plan);
-        Ok(())
-    }
-
     #[tokio::test]
     async fn select_with_periods() -> Result<()> {
         // define data with a column name that has a "." in it:

datafusion/core/tests/fuzz_cases/window_fuzz.rs

@@ -34,8 +34,7 @@ use datafusion_common::{Result, ScalarValue};
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_expr::type_coercion::functions::data_types_with_aggregate_udf;
 use datafusion_expr::{
-    BuiltInWindowFunction, WindowFrame, WindowFrameBound, WindowFrameUnits,
-    WindowFunctionDefinition,
+    WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
 };
 use datafusion_functions_aggregate::count::count_udaf;
 use datafusion_functions_aggregate::min_max::{max_udaf, min_udaf};
@@ -414,36 +413,6 @@ fn get_random_function(
             ),
         );
     }
-    window_fn_map.insert(
-        "first_value",
-        (
-            WindowFunctionDefinition::BuiltInWindowFunction(
-                BuiltInWindowFunction::FirstValue,
-            ),
-            vec![arg.clone()],
-        ),
-    );
-    window_fn_map.insert(
-        "last_value",
-        (
-            WindowFunctionDefinition::BuiltInWindowFunction(
-                BuiltInWindowFunction::LastValue,
-            ),
-            vec![arg.clone()],
-        ),
-    );
-    window_fn_map.insert(
-        "nth_value",
-        (
-            WindowFunctionDefinition::BuiltInWindowFunction(
-                BuiltInWindowFunction::NthValue,
-            ),
-            vec![
-                arg.clone(),
-                lit(ScalarValue::Int64(Some(rng.gen_range(1..10)))),
-            ],
-        ),
-    );
 
     let rand_fn_idx = rng.gen_range(0..window_fn_map.len());
     let fn_name = window_fn_map.keys().collect::<Vec<_>>()[rand_fn_idx];

datafusion/expr/src/expr.rs

@@ -21,17 +21,13 @@ use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Display, Formatter, Write};
 use std::hash::{Hash, Hasher};
 use std::mem;
-use std::str::FromStr;
 use std::sync::Arc;
 
 use crate::expr_fn::binary_expr;
 use crate::logical_plan::Subquery;
 use crate::utils::expr_to_columns;
 use crate::Volatility;
-use crate::{
-    udaf, BuiltInWindowFunction, ExprSchemable, Operator, Signature, WindowFrame,
-    WindowUDF,
-};
+use crate::{udaf, ExprSchemable, Operator, Signature, WindowFrame, WindowUDF};
 
 use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::cse::HashNode;
@@ -693,9 +689,6 @@ impl AggregateFunction {
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
 /// Defines which implementation of an aggregate function DataFusion should call.
 pub enum WindowFunctionDefinition {
-    /// A built in aggregate function that leverages an aggregate function
-    /// A a built-in window function
-    BuiltInWindowFunction(BuiltInWindowFunction),
     /// A user defined aggregate function
     AggregateUDF(Arc<crate::AggregateUDF>),
     /// A user defined aggregate function
@@ -711,9 +704,6 @@ impl WindowFunctionDefinition {
         display_name: &str,
     ) -> Result<DataType> {
         match self {
-            WindowFunctionDefinition::BuiltInWindowFunction(fun) => {
-                fun.return_type(input_expr_types)
-            }
             WindowFunctionDefinition::AggregateUDF(fun) => {
                 fun.return_type(input_expr_types)
             }
@@ -726,7 +716,6 @@ impl WindowFunctionDefinition {
     /// The signatures supported by the function `fun`.
     pub fn signature(&self) -> Signature {
         match self {
-            WindowFunctionDefinition::BuiltInWindowFunction(fun) => fun.signature(),
             WindowFunctionDefinition::AggregateUDF(fun) => fun.signature().clone(),
             WindowFunctionDefinition::WindowUDF(fun) => fun.signature().clone(),
         }
@@ -735,7 +724,6 @@ impl WindowFunctionDefinition {
     /// Function's name for display
     pub fn name(&self) -> &str {
         match self {
-            WindowFunctionDefinition::BuiltInWindowFunction(fun) => fun.name(),
             WindowFunctionDefinition::WindowUDF(fun) => fun.name(),
             WindowFunctionDefinition::AggregateUDF(fun) => fun.name(),
         }
@@ -745,19 +733,12 @@ impl WindowFunctionDefinition {
 impl Display for WindowFunctionDefinition {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         match self {
-            WindowFunctionDefinition::BuiltInWindowFunction(fun) => Display::fmt(fun, f),
             WindowFunctionDefinition::AggregateUDF(fun) => Display::fmt(fun, f),
             WindowFunctionDefinition::WindowUDF(fun) => Display::fmt(fun, f),
         }
     }
 }
 
-impl From<BuiltInWindowFunction> for WindowFunctionDefinition {
-    fn from(value: BuiltInWindowFunction) -> Self {
-        Self::BuiltInWindowFunction(value)
-    }
-}
-
 impl From<Arc<crate::AggregateUDF>> for WindowFunctionDefinition {
     fn from(value: Arc<crate::AggregateUDF>) -> Self {
         Self::AggregateUDF(value)
@@ -783,9 +764,10 @@ impl From<Arc<WindowUDF>> for WindowFunctionDefinition {
 /// ```
 /// # use datafusion_expr::{Expr, BuiltInWindowFunction, col, ExprFunctionExt};
 /// # use datafusion_expr::expr::WindowFunction;
+/// use datafusion_expr::WindowFunctionDefinition::WindowUDF;
 /// // Create FIRST_VALUE(a) OVER (PARTITION BY b ORDER BY c)
 /// let expr = Expr::WindowFunction(
-///     WindowFunction::new(BuiltInWindowFunction::FirstValue, vec![col("a")])
+///     WindowFunction::new(WindowUDF::, vec![col("a")])
 /// )
 ///   .partition_by(vec![col("b")])
 ///   .order_by(vec![col("b").sort(true, true)])
@@ -823,23 +805,6 @@ impl WindowFunction {
     }
 }
 
-/// Find DataFusion's built-in window function by name.
-pub fn find_df_window_func(name: &str) -> Option<WindowFunctionDefinition> {
-    let name = name.to_lowercase();
-    // Code paths for window functions leveraging ordinary aggregators and
-    // built-in window functions are quite different, and the same function
-    // may have different implementations for these cases. If the sought
-    // function is not found among built-in window functions, we search for
-    // it among aggregate functions.
-    if let Ok(built_in_function) = BuiltInWindowFunction::from_str(name.as_str()) {
-        Some(WindowFunctionDefinition::BuiltInWindowFunction(
-            built_in_function,
-        ))
-    } else {
-        None
-    }
-}
-
 /// EXISTS expression
 #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
 pub struct Exists {
@@ -2525,77 +2490,6 @@ mod test {
 
     use super::*;
 
-    #[test]
-    fn test_first_value_return_type() -> Result<()> {
-        let fun = find_df_window_func("first_value").unwrap();
-        let observed = fun.return_type(&[DataType::Utf8], &[true], "")?;
-        assert_eq!(DataType::Utf8, observed);
-
-        let observed = fun.return_type(&[DataType::UInt64], &[true], "")?;
-        assert_eq!(DataType::UInt64, observed);
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_last_value_return_type() -> Result<()> {
-        let fun = find_df_window_func("last_value").unwrap();
-        let observed = fun.return_type(&[DataType::Utf8], &[true], "")?;
-        assert_eq!(DataType::Utf8, observed);
-
-        let observed = fun.return_type(&[DataType::Float64], &[true], "")?;
-        assert_eq!(DataType::Float64, observed);
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_nth_value_return_type() -> Result<()> {
-        let fun = find_df_window_func("nth_value").unwrap();
-        let observed =
-            fun.return_type(&[DataType::Utf8, DataType::UInt64], &[true, true], "")?;
-        assert_eq!(DataType::Utf8, observed);
-
-        let observed =
-            fun.return_type(&[DataType::Float64, DataType::UInt64], &[true, true], "")?;
-        assert_eq!(DataType::Float64, observed);
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_window_function_case_insensitive() -> Result<()> {
-        let names = vec!["first_value", "last_value", "nth_value"];
-        for name in names {
-            let fun = find_df_window_func(name).unwrap();
-            let fun2 = find_df_window_func(name.to_uppercase().as_str()).unwrap();
-            assert_eq!(fun, fun2);
-            if fun.to_string() == "first_value" || fun.to_string() == "last_value" {
-                assert_eq!(fun.to_string(), name);
-            } else {
-                assert_eq!(fun.to_string(), name.to_uppercase());
-            }
-        }
-        Ok(())
-    }
-
-    #[test]
-    fn test_find_df_window_function() {
-        assert_eq!(
-            find_df_window_func("first_value"),
-            Some(WindowFunctionDefinition::BuiltInWindowFunction(
-                BuiltInWindowFunction::FirstValue
-            ))
-        );
-        assert_eq!(
-            find_df_window_func("LAST_value"),
-            Some(WindowFunctionDefinition::BuiltInWindowFunction(
-                BuiltInWindowFunction::LastValue
-            ))
-        );
-        assert_eq!(find_df_window_func("not_exist"), None)
-    }
-
     #[test]
     fn test_display_wildcard() {
         assert_eq!(format!("{}", wildcard()), "*");

datafusion/expr/src/expr_schema.rs

@@ -478,12 +478,6 @@ impl Expr {
             .map(|e| e.get_type(schema))
             .collect::<Result<Vec<_>>>()?;
         match fun {
-            WindowFunctionDefinition::BuiltInWindowFunction(window_fun) => {
-                let return_type = window_fun.return_type(&data_types)?;
-                let nullable =
-                    !["RANK", "NTILE", "CUME_DIST"].contains(&window_fun.name());
-                Ok((return_type, nullable))
-            }
             WindowFunctionDefinition::AggregateUDF(udaf) => {
                 let new_types = data_types_with_aggregate_udf(&data_types, udaf)
                     .map_err(|err| {

datafusion/expr/src/lib.rs

@@ -64,7 +64,6 @@ pub mod type_coercion;
 pub mod utils;
 pub mod var_provider;
 pub mod window_frame;
-pub mod window_function;
 pub mod window_state;
 
 pub use built_in_window_function::BuiltInWindowFunction;

datafusion/expr/src/window_function.rs

@@ -34,6 +34,7 @@ pub mod macros;
 
 pub mod cume_dist;
 pub mod lead_lag;
+pub mod nth_value;
 pub mod ntile;
 pub mod rank;
 pub mod row_number;
@@ -44,6 +45,7 @@ pub mod expr_fn {
     pub use super::cume_dist::cume_dist;
     pub use super::lead_lag::lag;
     pub use super::lead_lag::lead;
+    pub use super::nth_value::{first_value, last_value, nth_value};
     pub use super::ntile::ntile;
     pub use super::rank::{dense_rank, percent_rank, rank};
     pub use super::row_number::row_number;
@@ -60,6 +62,9 @@ pub fn all_default_window_functions() -> Vec<Arc<WindowUDF>> {
         rank::dense_rank_udwf(),
         rank::percent_rank_udwf(),
         ntile::ntile_udwf(),
+        nth_value::first_value_udwf(),
+        nth_value::last_value_udwf(),
+        nth_value::nth_value_udwf(),
     ]
 }
 /// Registers all enabled packages with a [`FunctionRegistry`]