fix some tests that were relying on legacy scan operator

Eventual-Inc · Dec 15, 2023 · 25b6808 · 25b6808
1 parent cfa53b1
commit 25b6808
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 41 deletions.
diff --git a/src/daft-plan/src/logical_ops/source.rs b/src/daft-plan/src/logical_ops/source.rs
@@ -53,8 +53,7 @@ impl Source {
                 partitioning_keys,
                 pushdowns,
             })) => {
-                res.push("Source:".to_string());
-                res.push(format!("Operator = {}", scan_op));
+                res.push(format!("Source: Operator = {}", scan_op));
                 res.push(format!("File schema = {}", source_schema.short_string()));
                 res.push(format!("Partitioning keys = {:?}", partitioning_keys));
                 res.extend(pushdowns.multiline_display());

diff --git a/src/daft-plan/src/optimization/logical_plan_tracker.rs b/src/daft-plan/src/optimization/logical_plan_tracker.rs
@@ -125,20 +125,18 @@ mod tests {
     #[test]
     fn same_plans_eq() -> DaftResult<()> {
         // Both plan1 and plan2 are Filter -> Project -> Source
-        let plan1 = dummy_scan_node(vec![
+        let scan_node = dummy_scan_node(vec![
             Field::new("a", DataType::Int64),
             Field::new("b", DataType::Utf8),
-        ])
-        .project(vec![col("a")], Default::default())?
-        .filter(col("a").lt(&lit(2)))?
-        .build();
-        let plan2 = dummy_scan_node(vec![
-            Field::new("a", DataType::Int64),
-            Field::new("b", DataType::Utf8),
-        ])
-        .project(vec![col("a")], Default::default())?
-        .filter(col("a").lt(&lit(2)))?
-        .build();
+        ]);
+        let plan1 = scan_node
+            .project(vec![col("a")], Default::default())?
+            .filter(col("a").lt(&lit(2)))?
+            .build();
+        let plan2 = scan_node
+            .project(vec![col("a")], Default::default())?
+            .filter(col("a").lt(&lit(2)))?
+            .build();
         // Double-check that logical plans are equal.
         assert_eq!(plan1, plan2);
 
@@ -157,20 +155,21 @@ mod tests {
     #[test]
     fn different_plans_not_eq_op_ordering() -> DaftResult<()> {
         // plan1 is Project -> Filter -> Source, while plan2 is Filter -> Project -> Source.
-        let plan1 = dummy_scan_node(vec![
-            Field::new("a", DataType::Int64),
-            Field::new("b", DataType::Utf8),
-        ])
-        .filter(col("a").lt(&lit(2)))?
-        .project(vec![col("a")], Default::default())?
-        .build();
-        let plan2 = dummy_scan_node(vec![
+
+        let scan_node = dummy_scan_node(vec![
             Field::new("a", DataType::Int64),
             Field::new("b", DataType::Utf8),
-        ])
-        .project(vec![col("a")], Default::default())?
-        .filter(col("a").lt(&lit(2)))?
-        .build();
+        ]);
+
+        let plan1 = scan_node
+            .filter(col("a").lt(&lit(2)))?
+            .project(vec![col("a")], Default::default())?
+            .build();
+
+        let plan2 = scan_node
+            .project(vec![col("a")], Default::default())?
+            .filter(col("a").lt(&lit(2)))?
+            .build();
         // Double-check that logical plans are NOT equal.
         assert_ne!(plan1, plan2);
 
@@ -189,20 +188,21 @@ mod tests {
     #[test]
     fn different_plans_not_eq_same_order_diff_config() -> DaftResult<()> {
         // Both plan1 and plan2 are Filter -> Project -> Source, but with different filter predicates.
-        let plan1 = dummy_scan_node(vec![
-            Field::new("a", DataType::Int64),
-            Field::new("b", DataType::Utf8),
-        ])
-        .project(vec![col("a")], Default::default())?
-        .filter(col("a").lt(&lit(2)))?
-        .build();
-        let plan2 = dummy_scan_node(vec![
+
+        let scan_node = dummy_scan_node(vec![
             Field::new("a", DataType::Int64),
             Field::new("b", DataType::Utf8),
-        ])
-        .project(vec![col("a")], Default::default())?
-        .filter(col("a").lt(&lit(4)))?
-        .build();
+        ]);
+
+        let plan1 = scan_node
+            .project(vec![col("a")], Default::default())?
+            .filter(col("a").lt(&lit(2)))?
+            .build();
+
+        let plan2 = scan_node
+            .project(vec![col("a")], Default::default())?
+            .filter(col("a").lt(&lit(4)))?
+            .build();
         // Double-check that logical plans are NOT equal.
         assert_ne!(plan1, plan2);
 

diff --git a/src/daft-plan/src/optimization/optimizer.rs b/src/daft-plan/src/optimization/optimizer.rs
@@ -538,7 +538,7 @@ mod tests {
         let expected = "\
         Filter: [[[col(a) < lit(2)] | lit(false)] | lit(false)] & lit(true)\
         \n  Project: col(a) + lit(3) AS c, col(a) + lit(1), col(a) + lit(2) AS b\
-        \n    Source: Json, File paths = [/foo], File schema = a (Int64), Format-specific config = Json(JsonSourceConfig { buffer_size: None, chunk_size: None }), Storage config = Native(NativeStorageConfig { io_config: None, multithreaded_io: true }), Output schema = a (Int64)";
+        \n    Source: Operator = AnonymousScanOperator: File paths=[/foo], Format-specific config = Json(JsonSourceConfig { buffer_size: None, chunk_size: None }), Storage config = Native(NativeStorageConfig { io_config: None, multithreaded_io: true }), File schema = a (Int64), Partitioning keys = [], Output schema = a (Int64)";
         assert_eq!(opt_plan.repr_indent(), expected);
         Ok(())
     }

diff --git a/src/daft-plan/src/optimization/rules/drop_repartition.rs b/src/daft-plan/src/optimization/rules/drop_repartition.rs
@@ -98,7 +98,7 @@ mod tests {
         .build();
         let expected = "\
         Repartition: Scheme = Hash, Number of partitions = 5, Partition by = col(a)\
-        \n  Source: Json, File paths = [/foo], File schema = a (Int64), b (Utf8), Format-specific config = Json(JsonSourceConfig { buffer_size: None, chunk_size: None }), Storage config = Native(NativeStorageConfig { io_config: None, multithreaded_io: true }), Output schema = a (Int64), b (Utf8)";
+        \n  Source: Operator = AnonymousScanOperator: File paths=[/foo], Format-specific config = Json(JsonSourceConfig { buffer_size: None, chunk_size: None }), Storage config = Native(NativeStorageConfig { io_config: None, multithreaded_io: true }), File schema = a (Int64), b (Utf8), Partitioning keys = [], Output schema = a (Int64), b (Utf8)";
         assert_optimized_plan_eq(plan, expected)?;
         Ok(())
     }

diff --git a/src/daft-plan/src/planner.rs b/src/daft-plan/src/planner.rs
@@ -662,7 +662,7 @@ mod tests {
             .repartition(Some(1), vec![col("a")], PartitionScheme::Hash)?
             .build();
         let physical_plan = plan(logical_plan.as_ref(), cfg.clone())?;
-        assert_matches!(physical_plan, PhysicalPlan::TabularScanJson(_));
+        assert_matches!(physical_plan, PhysicalPlan::TabularScan(_));
         Ok(())
     }