Migrate core test to insta, part1

datafusion/core/tests/expr_api/simplification.rs

@@ -17,6 +17,8 @@
 
 //! This program demonstrates the DataFusion expression simplification API.
 
+use insta::assert_snapshot;
+
 use arrow::array::types::IntervalDayTime;
 use arrow::array::{ArrayRef, Int32Array};
 use arrow::datatypes::{DataType, Field, Schema};
@@ -237,11 +239,15 @@ fn to_timestamp_expr_folded() -> Result<()> {
         .project(proj)?
         .build()?;
 
-    let expected = "Projection: TimestampNanosecond(1599566400000000000, None) AS to_timestamp(Utf8(\"2020-09-08T12:00:00+00:00\"))\
-            \n  TableScan: test"
-        .to_string();
-    let actual = get_optimized_plan_formatted(plan, &Utc::now());
-    assert_eq!(expected, actual);
+    let formatted = get_optimized_plan_formatted(plan, &Utc::now());
+    let actual = formatted.trim();
+    assert_snapshot!(
+        actual,
+        @r###"
+    Projection: TimestampNanosecond(1599566400000000000, None) AS to_timestamp(Utf8("2020-09-08T12:00:00+00:00"))
+      TableScan: test
+    "###
+    );
     Ok(())
 }
 
@@ -262,11 +268,16 @@ fn now_less_than_timestamp() -> Result<()> {
 
     // Note that constant folder runs and folds the entire
     // expression down to a single constant (true)
-    let expected = "Filter: Boolean(true)\
-                        \n  TableScan: test";
-    let actual = get_optimized_plan_formatted(plan, &time);
-
-    assert_eq!(expected, actual);
+    let formatted = get_optimized_plan_formatted(plan, &time);
+    let actual = formatted.trim();
+
+    assert_snapshot!(
+        actual,
+        @r###"
+    Filter: Boolean(true)
+      TableScan: test
+    "###
+    );
     Ok(())
 }
 
@@ -296,11 +307,16 @@ fn select_date_plus_interval() -> Result<()> {
 
     // Note that constant folder runs and folds the entire
     // expression down to a single constant (true)
-    let expected = r#"Projection: Date32("2021-01-09") AS to_timestamp(Utf8("2020-09-08T12:05:00+00:00")) + IntervalDayTime("IntervalDayTime { days: 123, milliseconds: 0 }")
-  TableScan: test"#;
-    let actual = get_optimized_plan_formatted(plan, &time);
-
-    assert_eq!(expected, actual);
+    let formatted = get_optimized_plan_formatted(plan, &time);
+    let actual = formatted.trim();
+
+    assert_snapshot!(
+        actual,
+        @r###"
+    Projection: Date32("2021-01-09") AS to_timestamp(Utf8("2020-09-08T12:05:00+00:00")) + IntervalDayTime("IntervalDayTime { days: 123, milliseconds: 0 }")
+      TableScan: test
+    "###
+    );
     Ok(())
 }
 
@@ -314,10 +330,15 @@ fn simplify_project_scalar_fn() -> Result<()> {
 
     // before simplify: power(t.f, 1.0)
     // after simplify:  t.f as "power(t.f, 1.0)"
-    let expected = "Projection: test.f AS power(test.f,Float64(1))\
-                      \n  TableScan: test";
-    let actual = get_optimized_plan_formatted(plan, &Utc::now());
-    assert_eq!(expected, actual);
+    let formatter = get_optimized_plan_formatted(plan, &Utc::now());
+    let actual = formatter.trim();
+    assert_snapshot!(
+        actual,
+        @r###"
+    Projection: test.f AS power(test.f,Float64(1))
+      TableScan: test
+    "###
+    );
     Ok(())
 }
 
@@ -337,9 +358,9 @@ fn simplify_scan_predicate() -> Result<()> {
 
     // before simplify: t.g = power(t.f, 1.0)
     // after simplify:  t.g = t.f"
-    let expected = "TableScan: test, full_filters=[g = f]";
-    let actual = get_optimized_plan_formatted(plan, &Utc::now());
-    assert_eq!(expected, actual);
+    let formatted = get_optimized_plan_formatted(plan, &Utc::now());
+    let actual = formatted.trim();
+    assert_snapshot!(actual, @"TableScan: test, full_filters=[g = f]");
     Ok(())
 }
 

datafusion/core/tests/optimizer/mod.rs

@@ -18,6 +18,7 @@
 //! Tests for the DataFusion SQL query planner that require functions from the
 //! datafusion-functions crate.
 
+use insta::assert_snapshot;
 use std::any::Any;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -56,9 +57,14 @@ fn init() {
 #[test]
 fn select_arrow_cast() {
     let sql = "SELECT arrow_cast(1234, 'Float64') as f64, arrow_cast('foo', 'LargeUtf8') as large";
-    let expected = "Projection: Float64(1234) AS f64, LargeUtf8(\"foo\") AS large\
-        \n  EmptyRelation";
-    quick_test(sql, expected);
+    let plan = test_sql(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: Float64(1234) AS f64, LargeUtf8("foo") AS large
+      EmptyRelation
+    "#
+    );
 }
 #[test]
 fn timestamp_nano_ts_none_predicates() -> Result<()> {
@@ -68,11 +74,15 @@ fn timestamp_nano_ts_none_predicates() -> Result<()> {
     // a scan should have the now()... predicate folded to a single
     // constant and compared to the column without a cast so it can be
     // pushed down / pruned
-    let expected =
-        "Projection: test.col_int32\
-         \n  Filter: test.col_ts_nano_none < TimestampNanosecond(1666612093000000000, None)\
-         \n    TableScan: test projection=[col_int32, col_ts_nano_none]";
-    quick_test(sql, expected);
+    let plan = test_sql(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r"
+    Projection: test.col_int32
+      Filter: test.col_ts_nano_none < TimestampNanosecond(1666612093000000000, None)
+        TableScan: test projection=[col_int32, col_ts_nano_none]
+    "
+    );
     Ok(())
 }
 
@@ -84,21 +94,30 @@ fn timestamp_nano_ts_utc_predicates() {
     // a scan should have the now()... predicate folded to a single
     // constant and compared to the column without a cast so it can be
     // pushed down / pruned
-    let expected =
-        "Projection: test.col_int32\n  Filter: test.col_ts_nano_utc < TimestampNanosecond(1666612093000000000, Some(\"+00:00\"))\
-         \n    TableScan: test projection=[col_int32, col_ts_nano_utc]";
-    quick_test(sql, expected);
+    let plan = test_sql(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: test.col_int32
+      Filter: test.col_ts_nano_utc < TimestampNanosecond(1666612093000000000, Some("+00:00"))
+        TableScan: test projection=[col_int32, col_ts_nano_utc]
+    "#
+    );
 }
 
 #[test]
 fn concat_literals() -> Result<()> {
     let sql = "SELECT concat(true, col_int32, false, null, 'hello', col_utf8, 12, 3.4) \
         AS col
         FROM test";
-    let expected =
-        "Projection: concat(Utf8(\"true\"), CAST(test.col_int32 AS Utf8), Utf8(\"falsehello\"), test.col_utf8, Utf8(\"123.4\")) AS col\
-        \n  TableScan: test projection=[col_int32, col_utf8]";
-    quick_test(sql, expected);
+    let plan = test_sql(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: concat(Utf8("true"), CAST(test.col_int32 AS Utf8), Utf8("falsehello"), test.col_utf8, Utf8("123.4")) AS col
+      TableScan: test projection=[col_int32, col_utf8]
+    "#
+    );
     Ok(())
 }
 
@@ -107,16 +126,15 @@ fn concat_ws_literals() -> Result<()> {
     let sql = "SELECT concat_ws('-', true, col_int32, false, null, 'hello', col_utf8, 12, '', 3.4) \
         AS col
         FROM test";
-    let expected =
-        "Projection: concat_ws(Utf8(\"-\"), Utf8(\"true\"), CAST(test.col_int32 AS Utf8), Utf8(\"false-hello\"), test.col_utf8, Utf8(\"12--3.4\")) AS col\
-        \n  TableScan: test projection=[col_int32, col_utf8]";
-    quick_test(sql, expected);
-    Ok(())
-}
-
-fn quick_test(sql: &str, expected_plan: &str) {
     let plan = test_sql(sql).unwrap();
-    assert_eq!(expected_plan, format!("{plan}"));
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: concat_ws(Utf8("-"), Utf8("true"), CAST(test.col_int32 AS Utf8), Utf8("false-hello"), test.col_utf8, Utf8("12--3.4")) AS col
+      TableScan: test projection=[col_int32, col_utf8]
+    "#
+    );
+    Ok(())
 }
 
 fn test_sql(sql: &str) -> Result<LogicalPlan> {

datafusion/core/tests/sql/aggregates.rs

@@ -19,6 +19,7 @@ use super::*;
 use datafusion::common::test_util::batches_to_string;
 use datafusion_catalog::MemTable;
 use datafusion_common::ScalarValue;
+use insta::assert_snapshot;
 
 #[tokio::test]
 async fn csv_query_array_agg_distinct() -> Result<()> {