add optimizer rule test

Author: JanKaul

datafusion/optimizer/src/reorder_join/mod.rs

@@ -176,3 +176,162 @@ fn optimize_join_subtree(
     // Use the Ibaraki-Kameda algorithm to find the optimal join ordering
     optimal_left_deep_join_plan(query_graph, cost_estimator)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test::*;
+    use crate::OptimizerContext;
+    use datafusion_expr::logical_plan::JoinType;
+    use datafusion_expr::{col, LogicalPlanBuilder};
+
+    /// Test that the JoinReorder rule correctly optimizes a three-way join
+    #[test]
+    fn test_join_reorder_three_way_join() -> Result<()> {
+        // Create three-way join: customer JOIN orders JOIN lineitem
+        let customer = scan_tpch_table("customer");
+        let orders = scan_tpch_table("orders");
+        let lineitem = scan_tpch_table("lineitem");
+
+        let plan = LogicalPlanBuilder::from(customer)
+            .join(
+                orders,
+                JoinType::Inner,
+                (vec!["c_custkey"], vec!["o_custkey"]),
+                None,
+            )?
+            .join(
+                lineitem,
+                JoinType::Inner,
+                (vec!["o_orderkey"], vec!["l_orderkey"]),
+                None,
+            )?
+            .build()?;
+
+        // Apply the optimizer rule
+        let optimizer = JoinReorder::default();
+        let config = OptimizerContext::new();
+        let result = optimizer.rewrite(plan.clone(), &config)?;
+
+        // Verify that the plan was transformed
+        assert!(result.transformed);
+        assert!(matches!(result.data, LogicalPlan::Join(_)));
+
+        Ok(())
+    }
+
+    /// Test that the JoinReorder rule doesn't transform non-join plans
+    #[test]
+    fn test_join_reorder_non_join_plan() -> Result<()> {
+        // Create a simple table scan (no joins)
+        let plan = scan_tpch_table("customer");
+
+        // Apply the optimizer rule
+        let optimizer = JoinReorder::default();
+        let config = OptimizerContext::new();
+        let result = optimizer.rewrite(plan.clone(), &config)?;
+
+        // Verify that the plan was NOT transformed
+        assert!(!result.transformed);
+
+        Ok(())
+    }
+
+    /// Test that the JoinReorder rule handles join with filter
+    #[test]
+    fn test_join_reorder_with_filter() -> Result<()> {
+        // Create three-way join with a filter on one of the joins
+        let customer = scan_tpch_table("customer");
+        let orders = scan_tpch_table("orders");
+        let lineitem = scan_tpch_table("lineitem");
+
+        let plan = LogicalPlanBuilder::from(customer)
+            .join(
+                orders,
+                JoinType::Inner,
+                (vec!["c_custkey"], vec!["o_custkey"]),
+                None,
+            )?
+            .join_with_expr_keys(
+                lineitem,
+                JoinType::Inner,
+                (vec![col("o_orderkey")], vec![col("l_orderkey")]),
+                Some(col("l_quantity").gt(datafusion_expr::lit(10.0))),
+            )?
+            .build()?;
+
+        // Apply the optimizer rule
+        let optimizer = JoinReorder::default();
+        let config = OptimizerContext::new();
+        let result = optimizer.rewrite(plan, &config)?;
+
+        // Verify that the plan was transformed and is a valid join
+        assert!(result.transformed);
+        assert!(matches!(result.data, LogicalPlan::Join(_)));
+
+        // The filter should still be present in the optimized plan
+        if let LogicalPlan::Join(join) = &result.data {
+            // Check that at least one join in the tree has a filter
+            let has_filter = join.filter.is_some()
+                || matches!(join.left.as_ref(), LogicalPlan::Join(j) if j.filter.is_some())
+                || matches!(join.right.as_ref(), LogicalPlan::Join(j) if j.filter.is_some());
+            assert!(has_filter, "Filter should be preserved in optimized plan");
+        }
+
+        Ok(())
+    }
+
+    /// Test is_join_subtree_root correctly identifies consecutive joins
+    #[test]
+    fn test_is_join_subtree_root() -> Result<()> {
+        // Create a consecutive join subtree: should be identified as root
+        let customer = scan_tpch_table("customer");
+        let orders = scan_tpch_table("orders");
+        let lineitem = scan_tpch_table("lineitem");
+
+        let consecutive_join = LogicalPlanBuilder::from(customer)
+            .join(
+                orders,
+                JoinType::Inner,
+                (vec!["c_custkey"], vec!["o_custkey"]),
+                None,
+            )?
+            .join(
+                lineitem,
+                JoinType::Inner,
+                (vec!["o_orderkey"], vec!["l_orderkey"]),
+                None,
+            )?
+            .build()?;
+
+        assert!(is_join_subtree_root(&consecutive_join));
+
+        // Create a non-consecutive join (join with filter between joins)
+        // This creates: Join -> Filter -> Join -> ...
+        let customer = scan_tpch_table("customer");
+        let orders = scan_tpch_table("orders");
+        let lineitem = scan_tpch_table("lineitem");
+
+        let non_consecutive_join = LogicalPlanBuilder::from(customer)
+            .join(
+                orders,
+                JoinType::Inner,
+                (vec!["c_custkey"], vec!["o_custkey"]),
+                None,
+            )?
+            .filter(col("o_totalprice").gt(datafusion_expr::lit(1000.0)))?
+            .join(
+                lineitem,
+                JoinType::Inner,
+                (vec!["o_orderkey"], vec!["l_orderkey"]),
+                None,
+            )?
+            .build()?;
+
+        // The top-level is a Join, but its left child is a Filter that contains a nested Join.
+        // This breaks the consecutive join pattern, so it should NOT be identified as a join subtree root.
+        assert!(!is_join_subtree_root(&non_consecutive_join));
+
+        Ok(())
+    }
+}