fix

Author: xiedeyantu

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

@@ -983,25 +983,6 @@ async fn roundtrip_is_not_unknown() -> Result<()> {
     roundtrip("SELECT * FROM data WHERE d IS NOT UNKNOWN").await
 }
 
-#[tokio::test]
-async fn roundtrip_union() -> Result<()> {
-    assert_expected_plan(
-        "SELECT a, e FROM data UNION SELECT a, e FROM data",
-        "Aggregate: groupBy=[[data.a, data.e]], aggr=[[]]\n  TableScan: data projection=[a, e]",
-        false,
-    )
-    .await
-}
-
-#[tokio::test]
-async fn roundtrip_union2() -> Result<()> {
-    assert_expected_plan(
-        "SELECT a, b FROM data UNION SELECT a, b FROM data UNION SELECT a, b FROM data",
-        "Aggregate: groupBy=[[data.a, data.b]], aggr=[[]]\n  TableScan: data projection=[a, b]",
-        false,
-    )
-    .await
-}
 
 #[tokio::test]
 async fn roundtrip_union_all() -> Result<()> {
@@ -2058,6 +2039,26 @@ async fn generate_plan_from_sql(
     Ok(plan2)
 }
 
+async fn assert_expected_plan_with_ctx(
+    ctx: SessionContext,
+    sql: &str,
+    expected_plan_str: &str,
+    assert_schema: bool,
+) -> Result<()> {
+    let df = ctx.sql(sql).await?;
+    let plan = df.into_optimized_plan()?;
+    let plan2 = substrait_roundtrip(&plan, &ctx).await?;
+
+    if assert_schema {
+        assert_eq!(plan.schema(), plan2.schema());
+    }
+
+    let plan2str = format!("{plan2}");
+    assert_eq!(expected_plan_str, &plan2str);
+
+    Ok(())
+}
+
 async fn assert_expected_plan(
     sql: &str,
     expected_plan_str: &str,