Minor: Move test code from context.rs into sql_integration (#2473)

apache · May 6, 2022 · 32cf354 · 32cf354
1 parent 940fbf0
commit 32cf354
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 99 deletions.
diff --git a/datafusion/core/src/execution/context.rs b/datafusion/core/src/execution/context.rs
@@ -1603,14 +1603,13 @@ impl FunctionRegistry for TaskContext {
 mod tests {
     use super::*;
     use crate::execution::context::QueryPlanner;
-    use crate::logical_plan::{binary_expr, lit, Operator};
     use crate::physical_plan::functions::make_scalar_function;
     use crate::test;
     use crate::test_util::parquet_test_data;
     use crate::variable::VarType;
     use crate::{
-        assert_batches_eq, assert_batches_sorted_eq,
-        logical_plan::{col, create_udf, sum, Expr},
+        assert_batches_eq,
+        logical_plan::{create_udf, Expr},
     };
     use crate::{logical_plan::create_udaf, physical_plan::expressions::AvgAccumulator};
     use arrow::array::ArrayRef;
@@ -1701,56 +1700,6 @@ mod tests {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn unprojected_filter() {
-        let ctx = SessionContext::new();
-        let df = ctx
-            .read_table(test::table_with_sequence(1, 3).unwrap())
-            .unwrap();
-
-        let df = df
-            .select(vec![binary_expr(col("i"), Operator::Plus, col("i"))])
-            .unwrap()
-            .filter(col("i").gt(lit(2)))
-            .unwrap();
-        let results = df.collect().await.unwrap();
-
-        let expected = vec![
-            "+--------------------------+",
-            "| ?table?.i Plus ?table?.i |",
-            "+--------------------------+",
-            "| 6                        |",
-            "+--------------------------+",
-        ];
-        assert_batches_sorted_eq!(expected, &results);
-    }
-
-    #[tokio::test]
-    async fn aggregate_with_alias() -> Result<()> {
-        let tmp_dir = TempDir::new()?;
-        let ctx = create_ctx(&tmp_dir, 1).await?;
-
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("c1", DataType::Utf8, false),
-            Field::new("c2", DataType::UInt32, false),
-        ]));
-
-        let plan = LogicalPlanBuilder::scan_empty(None, schema.as_ref(), None)?
-            .aggregate(vec![col("c1")], vec![sum(col("c2"))])?
-            .project(vec![col("c1"), sum(col("c2")).alias("total_salary")])?
-            .build()?;
-
-        let plan = ctx.optimize(&plan)?;
-
-        let physical_plan = ctx.create_physical_plan(&Arc::new(plan)).await?;
-        assert_eq!("c1", physical_plan.schema().field(0).name().as_str());
-        assert_eq!(
-            "total_salary",
-            physical_plan.schema().field(1).name().as_str()
-        );
-        Ok(())
-    }
-
     #[tokio::test]
     async fn case_sensitive_identifiers_user_defined_functions() -> Result<()> {
         let mut ctx = SessionContext::new();
@@ -2097,51 +2046,6 @@ mod tests {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn union_test() -> Result<()> {
-        let ctx = SessionContext::with_config(
-            SessionConfig::new().with_information_schema(true),
-        );
-
-        let result = ctx
-            .sql("SELECT 1 A UNION ALL SELECT 2")
-            .await
-            .unwrap()
-            .collect()
-            .await
-            .unwrap();
-
-        #[rustfmt::skip]
-        let expected = vec![
-            "+---+",
-            "| a |",
-            "+---+",
-            "| 1 |",
-            "| 2 |",
-            "+---+"
-        ];
-        assert_batches_eq!(expected, &result);
-
-        let result = ctx
-            .sql("SELECT 1 UNION SELECT 2")
-            .await
-            .unwrap()
-            .collect()
-            .await
-            .unwrap();
-
-        let expected = vec![
-            "+----------+",
-            "| Int64(1) |",
-            "+----------+",
-            "| 1        |",
-            "| 2        |",
-            "+----------+",
-        ];
-        assert_batches_eq!(expected, &result);
-        Ok(())
-    }
-
     #[tokio::test]
     async fn read_with_glob_path() -> Result<()> {
         let ctx = SessionContext::new();

diff --git a/datafusion/core/tests/sql/aggregates.rs b/datafusion/core/tests/sql/aggregates.rs
@@ -16,7 +16,7 @@
 // under the License.
 
 use super::*;
-use datafusion::scalar::ScalarValue;
+use datafusion::{logical_plan::LogicalPlanBuilder, scalar::ScalarValue};
 
 #[tokio::test]
 async fn csv_query_avg_multi_batch() -> Result<()> {
@@ -1447,3 +1447,28 @@ async fn count_distinct_integers_aggregated_multiple_partitions() -> Result<()>
 
     Ok(())
 }
+
+#[tokio::test]
+async fn aggregate_with_alias() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("c1", DataType::Utf8, false),
+        Field::new("c2", DataType::UInt32, false),
+    ]));
+
+    let plan = LogicalPlanBuilder::scan_empty(None, schema.as_ref(), None)?
+        .aggregate(vec![col("c1")], vec![sum(col("c2"))])?
+        .project(vec![col("c1"), sum(col("c2")).alias("total_salary")])?
+        .build()?;
+
+    let plan = ctx.optimize(&plan)?;
+
+    let physical_plan = ctx.create_physical_plan(&Arc::new(plan)).await?;
+    assert_eq!("c1", physical_plan.schema().field(0).name().as_str());
+    assert_eq!(
+        "total_salary",
+        physical_plan.schema().field(1).name().as_str()
+    );
+    Ok(())
+}
diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs
@@ -1188,3 +1188,25 @@ async fn boolean_literal() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn unprojected_filter() {
+    let ctx = SessionContext::new();
+    let df = ctx.read_table(table_with_sequence(1, 3).unwrap()).unwrap();
+
+    let df = df
+        .select(vec![col("i") + col("i")])
+        .unwrap()
+        .filter(col("i").gt(lit(2)))
+        .unwrap();
+    let results = df.collect().await.unwrap();
+
+    let expected = vec![
+        "+--------------------------+",
+        "| ?table?.i Plus ?table?.i |",
+        "+--------------------------+",
+        "| 6                        |",
+        "+--------------------------+",
+    ];
+    assert_batches_sorted_eq!(expected, &results);
+}
diff --git a/datafusion/core/tests/sql/union.rs b/datafusion/core/tests/sql/union.rs
@@ -79,3 +79,47 @@ async fn union_all_with_aggregate() -> Result<()> {
     assert_batches_eq!(expected, &actual);
     Ok(())
 }
+
+#[tokio::test]
+async fn union_schemas() -> Result<()> {
+    let ctx =
+        SessionContext::with_config(SessionConfig::new().with_information_schema(true));
+
+    let result = ctx
+        .sql("SELECT 1 A UNION ALL SELECT 2")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+
+    #[rustfmt::skip]
+    let expected = vec![
+        "+---+",
+        "| a |",
+        "+---+",
+        "| 1 |",
+        "| 2 |",
+        "+---+"
+    ];
+    assert_batches_eq!(expected, &result);
+
+    let result = ctx
+        .sql("SELECT 1 UNION SELECT 2")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+
+    let expected = vec![
+        "+----------+",
+        "| Int64(1) |",
+        "+----------+",
+        "| 1        |",
+        "| 2        |",
+        "+----------+",
+    ];
+    assert_batches_eq!(expected, &result);
+    Ok(())
+}