feat(bridge): cross join support (#46)

* feat: support cross join

Signed-off-by: Yuchen Liang <[email protected]>

* feat: disable datafusion optimizer

Signed-off-by: Yuchen Liang <[email protected]>

* fix

Signed-off-by: Yuchen Liang <[email protected]>

* add implementation rule and use CrossExec

Signed-off-by: Yuchen Liang <[email protected]>

---------

Signed-off-by: Yuchen Liang <[email protected]>

Author: yliang412

datafusion-optd-cli/src/main.rs

@@ -134,6 +134,8 @@ struct Args {
         default_value = "40"
     )]
     maxrows: MaxRows,
+    #[clap(long, help = "Turn on datafusion logical optimizer before optd")]
+    enable_logical: bool,
 }
 
 #[tokio::main]
@@ -156,6 +158,11 @@ pub async fn main() -> Result<()> {
     };
 
     let mut session_config = SessionConfig::from_env()?.with_information_schema(true);
+    
+    if !args.enable_logical {
+        session_config.options_mut().optimizer.max_passes = 0;
+    }
+    
 
     if let Some(batch_size) = args.batch_size {
         session_config = session_config.with_batch_size(batch_size);
@@ -187,9 +194,11 @@ pub async fn main() -> Result<()> {
     let mut ctx = {
         let mut state =
             SessionState::new_with_config_rt(session_config.clone(), Arc::new(runtime_env));
-        // clean up optimizer rules so that we can plug in our own optimizer
-        // state = state.with_optimizer_rules(vec![]);
-        // state = state.with_physical_optimizer_rules(vec![]);
+        if !args.enable_logical {
+            // clean up optimizer rules so that we can plug in our own optimizer
+            state = state.with_optimizer_rules(vec![]);
+            state = state.with_physical_optimizer_rules(vec![]);
+        }
         // use optd-bridge query planner
         let optimizer = DatafusionOptimizer::new_physical(Box::new(DatafusionCatalog::new(
             state.catalog_list(),

optd-adaptive-demo/src/bin/optd-adaptive-tpch-q8.rs

@@ -19,7 +19,8 @@ static GLOBAL: MiMalloc = MiMalloc;
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let session_config = SessionConfig::from_env()?.with_information_schema(true);
+    let mut session_config = SessionConfig::from_env()?.with_information_schema(true);
+    session_config.options_mut().optimizer.max_passes = 0;
 
     let rn_config = RuntimeConfig::new();
     let runtime_env = RuntimeEnv::new(rn_config.clone())?;
@@ -30,6 +31,10 @@ async fn main() -> Result<()> {
         let optimizer = DatafusionOptimizer::new_physical(Box::new(DatafusionCatalog::new(
             state.catalog_list(),
         )));
+        // clean up optimizer rules so that we can plug in our own optimizer
+        state = state.with_optimizer_rules(vec![]);
+        state = state.with_physical_optimizer_rules(vec![]);
+        // use optd-bridge query planner
         state = state.with_query_planner(Arc::new(OptdQueryPlanner::new(optimizer)));
         SessionContext::new_with_state(state)
     };

optd-datafusion-bridge/src/from_optd.rs

@@ -3,21 +3,16 @@ use std::{collections::HashMap, sync::Arc};
 use anyhow::{bail, Context, Result};
 use async_recursion::async_recursion;
 use datafusion::{
-    arrow::{
-        compute::kernels::filter,
-        datatypes::{Schema, SchemaRef},
-    },
+    arrow::datatypes::{Schema, SchemaRef},
     datasource::source_as_provider,
     logical_expr::Operator,
     physical_expr,
     physical_plan::{
         self,
         aggregates::AggregateMode,
-        explain::ExplainExec,
         expressions::create_aggregate_expr,
         joins::{
-            utils::{ColumnIndex, JoinFilter},
-            PartitionMode,
+            utils::{ColumnIndex, JoinFilter}, CrossJoinExec, PartitionMode
         },
         projection::ProjectionExec,
         AggregateExpr, ExecutionPlan, PhysicalExpr,
@@ -31,7 +26,7 @@ use optd_datafusion_repr::{
         PhysicalFilter, PhysicalHashJoin, PhysicalNestedLoopJoin, PhysicalProjection, PhysicalScan,
         PhysicalSort, PlanNode, SortOrderExpr, SortOrderType,
     },
-    PhysicalCollector, Value,
+    PhysicalCollector,
 };
 
 use crate::{physical_collector::CollectorExec, OptdPlanContext};
@@ -320,6 +315,11 @@ impl OptdPlanContext<'_> {
         };
 
         let physical_expr = self.from_optd_expr(node.cond(), &Arc::new(filter_schema.clone()))?;
+
+        if let JoinType::Cross = node.join_type() {
+            return Ok(Arc::new(CrossJoinExec::new(left_exec, right_exec)) as Arc<dyn ExecutionPlan + 'static>);
+        }
+
         let join_type = match node.join_type() {
             JoinType::Inner => datafusion::logical_expr::JoinType::Inner,
             JoinType::LeftOuter => datafusion::logical_expr::JoinType::Left,

optd-datafusion-bridge/src/into_optd.rs

@@ -113,8 +113,8 @@ impl OptdPlanContext<'_> {
                     expr,
                 )
                 .into_expr())
-            }
-            _ => bail!("{:?}", expr),
+            } 
+            _ => bail!("Unsupported expression: {:?}", expr),
         }
     }
 
@@ -230,6 +230,12 @@ impl OptdPlanContext<'_> {
         }
     }
 
+    fn into_optd_cross_join(&mut self, node: &logical_plan::CrossJoin) -> Result<LogicalJoin> {
+        let left = self.into_optd_plan_node(node.left.as_ref())?;
+        let right = self.into_optd_plan_node(node.right.as_ref())?;
+        Ok(LogicalJoin::new(left, right, ConstantExpr::bool(true).into_expr(), JoinType::Cross))
+    }
+
     fn into_optd_plan_node(&mut self, node: &LogicalPlan) -> Result<PlanNode> {
         let node = match node {
             LogicalPlan::TableScan(node) => self.into_optd_table_scan(node)?.into_plan_node(),
@@ -239,6 +245,7 @@ impl OptdPlanContext<'_> {
             LogicalPlan::SubqueryAlias(node) => self.into_optd_plan_node(node.input.as_ref())?,
             LogicalPlan::Join(node) => self.into_optd_join(node)?.into_plan_node(),
             LogicalPlan::Filter(node) => self.into_optd_filter(node)?.into_plan_node(),
+            LogicalPlan::CrossJoin(node) => self.into_optd_cross_join(node)?.into_plan_node(),
             _ => bail!(
                 "unsupported plan node: {}",
                 format!("{:?}", node).split('\n').next().unwrap()

optd-datafusion-repr/src/rules/physical.rs

@@ -34,6 +34,9 @@ impl PhysicalConversionRule {
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Join(
                 JoinType::LeftOuter,
             ))),
+            Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Join(
+                JoinType::Cross,
+            ))),
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Filter)),
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Sort)),
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Agg)),

optd-sqlplannertest/src/lib.rs

@@ -24,7 +24,8 @@ pub struct DatafusionDb {
 
 impl DatafusionDb {
     pub async fn new() -> Result<Self> {
-        let session_config = SessionConfig::from_env()?.with_information_schema(true);
+        let mut session_config = SessionConfig::from_env()?.with_information_schema(true);
+        session_config.options_mut().optimizer.max_passes = 0;
 
         let rn_config = RuntimeConfig::new();
         let runtime_env = RuntimeEnv::new(rn_config.clone())?;
@@ -35,6 +36,10 @@ impl DatafusionDb {
             let optimizer = DatafusionOptimizer::new_physical(Box::new(DatafusionCatalog::new(
                 state.catalog_list(),
             )));
+            // clean up optimizer rules so that we can plug in our own optimizer
+            state = state.with_optimizer_rules(vec![]);
+            state = state.with_physical_optimizer_rules(vec![]);
+            // use optd-bridge query planner
             state = state.with_query_planner(Arc::new(OptdQueryPlanner::new(optimizer)));
             SessionContext::new_with_state(state)
         };