test

Author: davidlghellin

src/max_min_by.rs

@@ -325,11 +325,15 @@ mod tests {
 
     #[cfg(test)]
     mod max_by {
+
         use super::*;
         async fn extract_string(df: prelude::DataFrame) -> error::Result<String> {
             let results = df.collect().await?;
             let col = results[0].column(0);
-            let arr = col.as_any().downcast_ref::<arrow::array::StringArray>().unwrap();
+            let arr = col
+                .as_any()
+                .downcast_ref::<arrow::array::StringArray>()
+                .unwrap();
             Ok(arr.value(0).to_string())
         }
 
@@ -338,6 +342,7 @@ mod tests {
             ctx.register_udaf(MaxByFunction::new().into());
             Ok(ctx)
         }
+
         #[tokio::test]
         async fn test_max_by_string_int() -> error::Result<()> {
             let query = format!(
@@ -350,41 +355,6 @@ mod tests {
             Ok(())
         }
 
-        #[tokio::test]
-        async fn test_max_by_ignores_nulls_in_ok() -> error::Result<()> {
-            let ctx = ctx_max()?;
-            let sql = r#"
-                SELECT max_by(v, k)
-                FROM (
-                    VALUES
-                        ('a', 1),
-                        ('b', CAST(NULL AS INT)),
-                        ('c', 2)
-                ) AS t(v, k)
-            "#;
-            let df = ctx.sql(sql).await?;
-            let got = extract_string(df).await?;
-            assert_eq!(got, "c", "max_by should ignore NULLs");
-            Ok(())
-        }
-        #[tokio::test]
-        async fn test_max_by_ignores_nulls_in_ko() -> error::Result<()> {
-            let ctx = ctx_max()?;
-            let sql = r#"
-                SELECT max_by(v, k)
-                FROM (
-                    VALUES
-                        ('a', 1),
-                        ('b', CAST(NULL AS INT)),
-                        ('c', 2)
-                ) AS t(v, k)
-            "#;
-            let df = ctx.sql(sql).await?;
-            let got = extract_string(df).await?;
-            assert_eq!(got, "b", "max_by should ignore NULLs");
-            Ok(())
-        }
-
         #[tokio::test]
         async fn test_max_by_string_float() -> error::Result<()> {
             let query = format!(
@@ -433,6 +403,43 @@ mod tests {
             Ok(())
         }
 
+        #[tokio::test]
+        async fn test_max_by_ignores_nulls() -> error::Result<()> {
+            let ctx = ctx_max()?;
+            let sql = r#"
+                SELECT max_by(v, k)
+                FROM (
+                    VALUES
+                        ('a', 1),
+                        ('b', CAST(NULL AS INT)),
+                        ('c', 2)
+                ) AS t(v, k)
+            "#;
+            let df = ctx.sql(sql).await?;
+            let got = extract_string(df).await?;
+            assert_eq!(got, "c", "max_by should ignore NULLs");
+            Ok(())
+        }
+
+        #[tokio::test]
+        async fn test_max_like_main_test() -> error::Result<()> {
+            let ctx = ctx_max()?;
+            let sql = r#"
+                SELECT max_by(v, k)
+                FROM (
+                    VALUES
+                        ('a', 10),
+                        ('b', 5),
+                        ('c', 15),
+                        ('d', 8)
+                ) AS t(v, k)
+            "#;
+            let df = ctx.sql(sql).await?;
+            let got = extract_string(df).await?;
+            assert_eq!(got, "c");
+            Ok(())
+        }
+
         fn ctx() -> error::Result<prelude::SessionContext> {
             let ctx = test_ctx()?;
             let max_by_udaf = MaxByFunction::new();
@@ -445,6 +452,21 @@ mod tests {
     mod min_by {
 
         use super::*;
+        async fn extract_string(df: prelude::DataFrame) -> error::Result<String> {
+            let results = df.collect().await?;
+            let col = results[0].column(0);
+            let arr = col
+                .as_any()
+                .downcast_ref::<arrow::array::StringArray>()
+                .unwrap();
+            Ok(arr.value(0).to_string())
+        }
+
+        fn ctx_min() -> error::Result<prelude::SessionContext> {
+            let ctx = prelude::SessionContext::new();
+            ctx.register_udaf(MinByFunction::new().into());
+            Ok(ctx)
+        }
 
         #[tokio::test]
         async fn test_min_by_string_int() -> error::Result<()> {
@@ -506,6 +528,43 @@ mod tests {
             Ok(())
         }
 
+        #[tokio::test]
+        async fn test_min_by_ignores_nulls() -> error::Result<()> {
+            let ctx = ctx_min()?;
+            let sql = r#"
+                SELECT min_by(v, k)
+                FROM (
+                    VALUES
+                        ('a', 1),
+                        ('b', CAST(NULL AS INT)),
+                        ('c', 2)
+                ) AS t(v, k)
+            "#;
+            let df = ctx.sql(sql).await?;
+            let got = extract_string(df).await?;
+            assert_eq!(got, "a", "max_by should ignore NULLs");
+            Ok(())
+        }
+
+        #[tokio::test]
+        async fn test_min_like_main_test() -> error::Result<()> {
+            let ctx = ctx_min()?;
+            let sql = r#"
+                SELECT min_by(v, k)
+                FROM (
+                    VALUES
+                        ('a', 10),
+                        ('b', 5),
+                        ('c', 15),
+                        ('d', 8)
+                ) AS t(v, k)
+            "#;
+            let df = ctx.sql(sql).await?;
+            let got = extract_string(df).await?;
+            assert_eq!(got, "b");
+            Ok(())
+        }
+
         fn ctx() -> error::Result<prelude::SessionContext> {
             let ctx = test_ctx()?;
             let min_by_udaf = MinByFunction::new();