Migrate physical plan tests to insta (Part-2) (#15364)

* Migrated tests to insta in aggregate/topk/heap.rs

* Migrated tests to insta in joins/nested_loop_join.rs

* Migrated tests to insta in joins/sort_merge_join.rs

* Remove snapshot uploaded by mistake

* Migrated tests to insta in repartition/mod.rs

* Fix batches_to_sort_string problem and formatting

Author: Shreyaskr1409

datafusion/physical-plan/src/aggregates/topk/heap.rs

@@ -485,6 +485,8 @@ pub fn new_heap(
 
 #[cfg(test)]
 mod tests {
+    use insta::assert_snapshot;
+
     use super::*;
 
     #[test]
@@ -494,10 +496,9 @@ mod tests {
         heap.append_or_replace(1, 1, &mut map);
 
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=1 idx=0, bucket=1
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         Ok(())
     }
@@ -514,11 +515,10 @@ val=1 idx=0, bucket=1
         assert_eq!(map, vec![(2, 0), (1, 1)]);
 
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=2 idx=0, bucket=2
 └── val=1 idx=1, bucket=1
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         Ok(())
     }
@@ -532,22 +532,20 @@ val=2 idx=0, bucket=2
         heap.append_or_replace(2, 2, &mut map);
         heap.append_or_replace(3, 3, &mut map);
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=3 idx=0, bucket=3
 ├── val=1 idx=1, bucket=1
 └── val=2 idx=2, bucket=2
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         let mut map = vec![];
         heap.append_or_replace(0, 0, &mut map);
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=2 idx=0, bucket=2
 ├── val=1 idx=1, bucket=1
 └── val=0 idx=2, bucket=0
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
         assert_eq!(map, vec![(2, 0), (0, 2)]);
 
         Ok(())
@@ -563,24 +561,22 @@ val=2 idx=0, bucket=2
         heap.append_or_replace(3, 3, &mut map);
         heap.append_or_replace(4, 4, &mut map);
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=4 idx=0, bucket=4
 ├── val=3 idx=1, bucket=3
 │   └── val=1 idx=3, bucket=1
 └── val=2 idx=2, bucket=2
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         let mut map = vec![];
         heap.replace_if_better(1, 0, &mut map);
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=4 idx=0, bucket=4
 ├── val=1 idx=1, bucket=1
 │   └── val=0 idx=3, bucket=3
 └── val=2 idx=2, bucket=2
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
         assert_eq!(map, vec![(1, 1), (3, 3)]);
 
         Ok(())
@@ -595,11 +591,10 @@ val=4 idx=0, bucket=4
         heap.append_or_replace(2, 2, &mut map);
 
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=2 idx=0, bucket=2
 └── val=1 idx=1, bucket=1
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         assert_eq!(heap.worst_val(), Some(&2));
         assert_eq!(heap.worst_map_idx(), 2);
@@ -616,11 +611,10 @@ val=2 idx=0, bucket=2
         heap.append_or_replace(2, 2, &mut map);
 
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=2 idx=0, bucket=2
 └── val=1 idx=1, bucket=1
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         let (vals, map_idxs) = heap.drain();
         assert_eq!(vals, vec![1, 2]);
@@ -639,20 +633,18 @@ val=2 idx=0, bucket=2
         heap.append_or_replace(2, 2, &mut map);
 
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=2 idx=0, bucket=2
 └── val=1 idx=1, bucket=1
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         let numbers = vec![(0, 1), (1, 2)];
         heap.renumber(numbers.as_slice());
         let actual = heap.to_string();
-        let expected = r#"
+        assert_snapshot!(actual, @r#"
 val=2 idx=0, bucket=1
 └── val=1 idx=1, bucket=2
-        "#;
-        assert_eq!(actual.trim(), expected.trim());
+            "#);
 
         Ok(())
     }

datafusion/physical-plan/src/joins/nested_loop_join.rs

@@ -1047,13 +1047,15 @@ pub(crate) mod tests {
     use arrow::array::Int32Array;
     use arrow::compute::SortOptions;
     use arrow::datatypes::{DataType, Field};
-    use datafusion_common::{assert_batches_sorted_eq, assert_contains, ScalarValue};
+    use datafusion_common::test_util::batches_to_sort_string;
+    use datafusion_common::{assert_contains, ScalarValue};
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
     use datafusion_expr::Operator;
     use datafusion_physical_expr::expressions::{BinaryExpr, Literal};
     use datafusion_physical_expr::{Partitioning, PhysicalExpr};
     use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 
+    use insta::assert_snapshot;
     use rstest::rstest;
 
     fn build_table(
@@ -1216,15 +1218,13 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
-        let expected = [
-            "+----+----+----+----+----+----+",
-            "| a1 | b1 | c1 | a2 | b2 | c2 |",
-            "+----+----+----+----+----+----+",
-            "| 5  | 5  | 50 | 2  | 2  | 80 |",
-            "+----+----+----+----+----+----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+----+----+----+----+
+            | a1 | b1 | c1 | a2 | b2 | c2 |
+            +----+----+----+----+----+----+
+            | 5  | 5  | 50 | 2  | 2  | 80 |
+            +----+----+----+----+----+----+
+            "#);
 
         Ok(())
     }
@@ -1245,17 +1245,15 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
-        let expected = [
-            "+----+----+-----+----+----+----+",
-            "| a1 | b1 | c1  | a2 | b2 | c2 |",
-            "+----+----+-----+----+----+----+",
-            "| 11 | 8  | 110 |    |    |    |",
-            "| 5  | 5  | 50  | 2  | 2  | 80 |",
-            "| 9  | 8  | 90  |    |    |    |",
-            "+----+----+-----+----+----+----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+-----+----+----+----+
+            | a1 | b1 | c1  | a2 | b2 | c2 |
+            +----+----+-----+----+----+----+
+            | 11 | 8  | 110 |    |    |    |
+            | 5  | 5  | 50  | 2  | 2  | 80 |
+            | 9  | 8  | 90  |    |    |    |
+            +----+----+-----+----+----+----+
+            "#);
 
         Ok(())
     }
@@ -1276,17 +1274,15 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
-        let expected = [
-            "+----+----+----+----+----+-----+",
-            "| a1 | b1 | c1 | a2 | b2 | c2  |",
-            "+----+----+----+----+----+-----+",
-            "|    |    |    | 10 | 10 | 100 |",
-            "|    |    |    | 12 | 10 | 40  |",
-            "| 5  | 5  | 50 | 2  | 2  | 80  |",
-            "+----+----+----+----+----+-----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+----+----+----+-----+
+            | a1 | b1 | c1 | a2 | b2 | c2  |
+            +----+----+----+----+----+-----+
+            |    |    |    | 10 | 10 | 100 |
+            |    |    |    | 12 | 10 | 40  |
+            | 5  | 5  | 50 | 2  | 2  | 80  |
+            +----+----+----+----+----+-----+
+            "#);
 
         Ok(())
     }
@@ -1307,19 +1303,17 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]);
-        let expected = [
-            "+----+----+-----+----+----+-----+",
-            "| a1 | b1 | c1  | a2 | b2 | c2  |",
-            "+----+----+-----+----+----+-----+",
-            "|    |    |     | 10 | 10 | 100 |",
-            "|    |    |     | 12 | 10 | 40  |",
-            "| 11 | 8  | 110 |    |    |     |",
-            "| 5  | 5  | 50  | 2  | 2  | 80  |",
-            "| 9  | 8  | 90  |    |    |     |",
-            "+----+----+-----+----+----+-----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+-----+----+----+-----+
+            | a1 | b1 | c1  | a2 | b2 | c2  |
+            +----+----+-----+----+----+-----+
+            |    |    |     | 10 | 10 | 100 |
+            |    |    |     | 12 | 10 | 40  |
+            | 11 | 8  | 110 |    |    |     |
+            | 5  | 5  | 50  | 2  | 2  | 80  |
+            | 9  | 8  | 90  |    |    |     |
+            +----+----+-----+----+----+-----+
+            "#);
 
         Ok(())
     }
@@ -1340,15 +1334,13 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1"]);
-        let expected = [
-            "+----+----+----+",
-            "| a1 | b1 | c1 |",
-            "+----+----+----+",
-            "| 5  | 5  | 50 |",
-            "+----+----+----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+----+
+            | a1 | b1 | c1 |
+            +----+----+----+
+            | 5  | 5  | 50 |
+            +----+----+----+
+            "#);
 
         Ok(())
     }
@@ -1369,16 +1361,14 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1"]);
-        let expected = [
-            "+----+----+-----+",
-            "| a1 | b1 | c1  |",
-            "+----+----+-----+",
-            "| 11 | 8  | 110 |",
-            "| 9  | 8  | 90  |",
-            "+----+----+-----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+-----+
+            | a1 | b1 | c1  |
+            +----+----+-----+
+            | 11 | 8  | 110 |
+            | 9  | 8  | 90  |
+            +----+----+-----+
+            "#);
 
         Ok(())
     }
@@ -1399,15 +1389,13 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a2", "b2", "c2"]);
-        let expected = [
-            "+----+----+----+",
-            "| a2 | b2 | c2 |",
-            "+----+----+----+",
-            "| 2  | 2  | 80 |",
-            "+----+----+----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+----+
+            | a2 | b2 | c2 |
+            +----+----+----+
+            | 2  | 2  | 80 |
+            +----+----+----+
+            "#);
 
         Ok(())
     }
@@ -1428,16 +1416,14 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a2", "b2", "c2"]);
-        let expected = [
-            "+----+----+-----+",
-            "| a2 | b2 | c2  |",
-            "+----+----+-----+",
-            "| 10 | 10 | 100 |",
-            "| 12 | 10 | 40  |",
-            "+----+----+-----+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+-----+
+            | a2 | b2 | c2  |
+            +----+----+-----+
+            | 10 | 10 | 100 |
+            | 12 | 10 | 40  |
+            +----+----+-----+
+            "#);
 
         Ok(())
     }
@@ -1458,17 +1444,15 @@ pub(crate) mod tests {
         )
         .await?;
         assert_eq!(columns, vec!["a1", "b1", "c1", "mark"]);
-        let expected = [
-            "+----+----+-----+-------+",
-            "| a1 | b1 | c1  | mark  |",
-            "+----+----+-----+-------+",
-            "| 11 | 8  | 110 | false |",
-            "| 5  | 5  | 50  | true  |",
-            "| 9  | 8  | 90  | false |",
-            "+----+----+-----+-------+",
-        ];
-
-        assert_batches_sorted_eq!(expected, &batches);
+        assert_snapshot!(batches_to_sort_string(&batches), @r#"
+            +----+----+-----+-------+
+            | a1 | b1 | c1  | mark  |
+            +----+----+-----+-------+
+            | 11 | 8  | 110 | false |
+            | 5  | 5  | 50  | true  |
+            | 9  | 8  | 90  | false |
+            +----+----+-----+-------+
+            "#);
 
         Ok(())
     }