Merge pull request #52 from AveryQi115/empty_relation

AveryQi115 · web-flow · commit 2694b9ae927b · 2024-02-10T23:26:14.000-05:00
feat: [LogicalOptimizer] add support for empty relation
diff --git a/datafusion-optd-cli/src/main.rs b/datafusion-optd-cli/src/main.rs
@@ -158,11 +158,10 @@ 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);
diff --git a/optd-datafusion-bridge/src/from_optd.rs b/optd-datafusion-bridge/src/from_optd.rs
@@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc};
 use anyhow::{bail, Context, Result};
 use async_recursion::async_recursion;
 use datafusion::{
-    arrow::datatypes::{Schema, SchemaRef},
+    arrow::datatypes::{DataType, Field, Schema, SchemaRef},
     datasource::source_as_provider,
     logical_expr::Operator,
     physical_expr,
@@ -12,7 +12,8 @@ use datafusion::{
         aggregates::AggregateMode,
         expressions::create_aggregate_expr,
         joins::{
-            utils::{ColumnIndex, JoinFilter}, CrossJoinExec, PartitionMode
+            utils::{ColumnIndex, JoinFilter},
+            CrossJoinExec, PartitionMode,
         },
         projection::ProjectionExec,
         AggregateExpr, ExecutionPlan, PhysicalExpr,
@@ -23,14 +24,36 @@ use optd_datafusion_repr::{
     plan_nodes::{
         BinOpExpr, BinOpType, ColumnRefExpr, ConstantExpr, ConstantType, Expr, FuncExpr, FuncType,
         JoinType, LogOpExpr, LogOpType, OptRelNode, OptRelNodeRef, OptRelNodeTyp, PhysicalAgg,
-        PhysicalFilter, PhysicalHashJoin, PhysicalNestedLoopJoin, PhysicalProjection, PhysicalScan,
-        PhysicalSort, PlanNode, SortOrderExpr, SortOrderType,
+        PhysicalEmptyRelation, PhysicalFilter, PhysicalHashJoin, PhysicalNestedLoopJoin,
+        PhysicalProjection, PhysicalScan, PhysicalSort, PlanNode, SortOrderExpr, SortOrderType,
     },
+    properties::schema::Schema as OptdSchema,
     PhysicalCollector,
 };
 
 use crate::{physical_collector::CollectorExec, OptdPlanContext};
 
+// TODO: current DataType and ConstantType are not 1 to 1 mapping
+// optd schema stores constantType from data type in catalog.get
+// for decimal128, the precision is lost
+fn from_optd_schema(optd_schema: &OptdSchema) -> Schema {
+    let match_type = |typ: &ConstantType| match typ {
+        ConstantType::Any => unimplemented!(),
+        ConstantType::Bool => DataType::Boolean,
+        ConstantType::Int => DataType::Int64,
+        ConstantType::Date => DataType::Date32,
+        ConstantType::Decimal => DataType::Float64,
+        ConstantType::Utf8String => DataType::Utf8,
+    };
+    let fields: Vec<_> = optd_schema
+        .0
+        .iter()
+        .enumerate()
+        .map(|(i, typ)| Field::new(&format!("c{}", i), match_type(typ), false))
+        .collect();
+    Schema::new(fields)
+}
+
 impl OptdPlanContext<'_> {
     #[async_recursion]
     async fn from_optd_table_scan(
@@ -317,7 +340,8 @@ 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>);
+            return Ok(Arc::new(CrossJoinExec::new(left_exec, right_exec))
+                as Arc<dyn ExecutionPlan + 'static>);
         }
 
         let join_type = match node.join_type() {
@@ -398,6 +422,10 @@ impl OptdPlanContext<'_> {
 
     #[async_recursion]
     async fn from_optd_plan_node(&mut self, node: PlanNode) -> Result<Arc<dyn ExecutionPlan>> {
+        let mut schema = OptdSchema(vec![]);
+        if node.typ() == OptRelNodeTyp::PhysicalEmptyRelation {
+            schema = node.schema(self.optimizer.unwrap().optd_optimizer());
+        }
         let rel_node = node.into_rel_node();
         let rel_node_dbg = rel_node.clone();
         let result = match &rel_node.typ {
@@ -440,6 +468,14 @@ impl OptdPlanContext<'_> {
                     self.optimizer.as_ref().unwrap().runtime_statistics.clone(),
                 )) as Arc<dyn ExecutionPlan>)
             }
+            OptRelNodeTyp::PhysicalEmptyRelation => {
+                let physical_node = PhysicalEmptyRelation::from_rel_node(rel_node).unwrap();
+                let datafusion_schema: Schema = from_optd_schema(&schema);
+                Ok(Arc::new(datafusion::physical_plan::empty::EmptyExec::new(
+                    physical_node.produce_one_row(),
+                    Arc::new(datafusion_schema),
+                )) as Arc<dyn ExecutionPlan>)
+            }
             typ => unimplemented!("{}", typ),
         };
         result.with_context(|| format!("when processing {}", rel_node_dbg))
diff --git a/optd-datafusion-bridge/src/into_optd.rs b/optd-datafusion-bridge/src/into_optd.rs
@@ -7,9 +7,9 @@ use datafusion::{
 use optd_core::rel_node::RelNode;
 use optd_datafusion_repr::plan_nodes::{
     BinOpExpr, BinOpType, ColumnRefExpr, ConstantExpr, Expr, ExprList, FuncExpr, FuncType,
-    JoinType, LogOpExpr, LogOpType, LogicalAgg, LogicalFilter, LogicalJoin, LogicalProjection,
-    LogicalScan, LogicalSort, OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode, SortOrderExpr,
-    SortOrderType,
+    JoinType, LogOpExpr, LogOpType, LogicalAgg, LogicalEmptyRelation, LogicalFilter, LogicalJoin,
+    LogicalProjection, LogicalScan, LogicalSort, OptRelNode, OptRelNodeRef, OptRelNodeTyp,
+    PlanNode, SortOrderExpr, SortOrderType,
 };
 
 use crate::OptdPlanContext;
@@ -69,6 +69,10 @@ impl OptdPlanContext<'_> {
                     let x = x.as_ref().unwrap();
                     Ok(ConstantExpr::string(x).into_expr())
                 }
+                ScalarValue::Int64(x) => {
+                    let x = x.as_ref().unwrap();
+                    Ok(ConstantExpr::int(*x as i64).into_expr())
+                }
                 ScalarValue::Date32(x) => {
                     let x = x.as_ref().unwrap();
                     Ok(ConstantExpr::date(*x as i64).into_expr())
@@ -113,7 +117,7 @@ impl OptdPlanContext<'_> {
                     expr,
                 )
                 .into_expr())
-            } 
+            }
             _ => bail!("Unsupported expression: {:?}", expr),
         }
     }
@@ -233,7 +237,19 @@ 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))
+        Ok(LogicalJoin::new(
+            left,
+            right,
+            ConstantExpr::bool(true).into_expr(),
+            JoinType::Cross,
+        ))
+    }
+
+    fn into_optd_empty_relation(
+        &mut self,
+        node: &logical_plan::EmptyRelation,
+    ) -> Result<LogicalEmptyRelation> {
+        Ok(LogicalEmptyRelation::new(node.produce_one_row))
     }
 
     fn into_optd_plan_node(&mut self, node: &LogicalPlan) -> Result<PlanNode> {
@@ -246,6 +262,9 @@ impl OptdPlanContext<'_> {
             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(),
+            LogicalPlan::EmptyRelation(node) => {
+                self.into_optd_empty_relation(node)?.into_plan_node()
+            }
             _ => bail!(
                 "unsupported plan node: {}",
                 format!("{:?}", node).split('\n').next().unwrap()
diff --git a/optd-datafusion-repr/src/cost/base_cost.rs b/optd-datafusion-repr/src/cost/base_cost.rs
@@ -106,6 +106,7 @@ impl CostModel<OptRelNodeTyp> for OptCostModel {
                     .unwrap_or(1) as f64;
                 Self::cost(row_cnt, 0.0, row_cnt)
             }
+            OptRelNodeTyp::PhysicalEmptyRelation => Self::cost(0.5, 0.01, 0.0),
             OptRelNodeTyp::PhysicalFilter => {
                 let (row_cnt, _, _) = Self::cost_tuple(&children[0]);
                 let (_, compute_cost, _) = Self::cost_tuple(&children[1]);
diff --git a/optd-datafusion-repr/src/plan_nodes.rs b/optd-datafusion-repr/src/plan_nodes.rs
@@ -2,6 +2,7 @@
 
 mod agg;
 mod apply;
+mod empty_relation;
 mod expr;
 mod filter;
 mod join;
@@ -19,6 +20,7 @@ use optd_core::{
 
 pub use agg::{LogicalAgg, PhysicalAgg};
 pub use apply::{ApplyType, LogicalApply};
+pub use empty_relation::{LogicalEmptyRelation, PhysicalEmptyRelation};
 pub use expr::{
     BinOpExpr, BinOpType, ColumnRefExpr, ConstantExpr, ConstantType, ExprList, FuncExpr, FuncType,
     LogOpExpr, LogOpType, SortOrderExpr, SortOrderType, UnOpExpr, UnOpType,
@@ -48,6 +50,7 @@ pub enum OptRelNodeTyp {
     Sort,
     Agg,
     Apply(ApplyType),
+    EmptyRelation,
     // Physical plan nodes
     PhysicalProjection,
     PhysicalFilter,
@@ -56,6 +59,7 @@ pub enum OptRelNodeTyp {
     PhysicalAgg,
     PhysicalHashJoin(JoinType),
     PhysicalNestedLoopJoin(JoinType),
+    PhysicalEmptyRelation,
     PhysicalCollector(GroupId), // only produced after optimization is done
     // Expressions
     Constant(ConstantType),
@@ -78,6 +82,7 @@ impl OptRelNodeTyp {
                 | Self::Apply(_)
                 | Self::Sort
                 | Self::Agg
+                | Self::EmptyRelation
                 | Self::PhysicalProjection
                 | Self::PhysicalFilter
                 | Self::PhysicalNestedLoopJoin(_)
@@ -86,6 +91,7 @@ impl OptRelNodeTyp {
                 | Self::PhysicalAgg
                 | Self::PhysicalHashJoin(_)
                 | Self::PhysicalCollector(_)
+                | Self::PhysicalEmptyRelation
         )
     }
 
@@ -120,6 +126,7 @@ impl RelNodeTyp for OptRelNodeTyp {
                 | Self::Apply(_)
                 | Self::Sort
                 | Self::Agg
+                | Self::EmptyRelation
         )
     }
 
@@ -194,7 +201,7 @@ impl PlanNode {
         self.0.typ.clone()
     }
 
-    pub fn schema(&self, optimizer: CascadesOptimizer<OptRelNodeTyp>) -> Schema {
+    pub fn schema(&self, optimizer: &CascadesOptimizer<OptRelNodeTyp>) -> Schema {
         let group_id = optimizer.resolve_group_id(self.0.clone());
         optimizer.get_property_by_group::<SchemaPropertyBuilder>(group_id, 0)
     }
@@ -300,6 +307,9 @@ pub fn explain(rel_node: OptRelNodeRef) -> Pretty<'static> {
         OptRelNodeTyp::Apply(_) => LogicalApply::from_rel_node(rel_node)
             .unwrap()
             .dispatch_explain(),
+        OptRelNodeTyp::EmptyRelation => LogicalEmptyRelation::from_rel_node(rel_node)
+            .unwrap()
+            .dispatch_explain(),
         OptRelNodeTyp::PhysicalFilter => PhysicalFilter::from_rel_node(rel_node)
             .unwrap()
             .dispatch_explain(),
@@ -345,6 +355,9 @@ pub fn explain(rel_node: OptRelNodeRef) -> Pretty<'static> {
         OptRelNodeTyp::PhysicalCollector(group_id) => PhysicalCollector::from_rel_node(rel_node)
             .unwrap()
             .dispatch_explain(),
+        OptRelNodeTyp::PhysicalEmptyRelation => PhysicalEmptyRelation::from_rel_node(rel_node)
+            .unwrap()
+            .dispatch_explain(),
     }
 }
 
diff --git a/optd-datafusion-repr/src/plan_nodes/empty_relation.rs b/optd-datafusion-repr/src/plan_nodes/empty_relation.rs
@@ -0,0 +1,88 @@
+use pretty_xmlish::Pretty;
+
+use optd_core::rel_node::{RelNode, Value};
+
+use super::{replace_typ, OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode};
+
+#[derive(Clone, Debug)]
+pub struct LogicalEmptyRelation(pub PlanNode);
+
+impl OptRelNode for LogicalEmptyRelation {
+    fn into_rel_node(self) -> OptRelNodeRef {
+        self.0.into_rel_node()
+    }
+
+    fn from_rel_node(rel_node: OptRelNodeRef) -> Option<Self> {
+        if rel_node.typ != OptRelNodeTyp::EmptyRelation {
+            return None;
+        }
+        PlanNode::from_rel_node(rel_node).map(Self)
+    }
+
+    fn dispatch_explain(&self) -> Pretty<'static> {
+        Pretty::childless_record(
+            "LogicalEmptyRelation",
+            vec![("produce_one_row", self.produce_one_row().to_string().into())],
+        )
+    }
+}
+
+impl LogicalEmptyRelation {
+    pub fn new(produce_one_row: bool) -> LogicalEmptyRelation {
+        LogicalEmptyRelation(PlanNode(
+            RelNode {
+                typ: OptRelNodeTyp::EmptyRelation,
+                children: vec![],
+                data: Some(Value::Bool(produce_one_row)),
+            }
+            .into(),
+        ))
+    }
+
+    pub fn produce_one_row(&self) -> bool {
+        self.clone()
+            .into_rel_node()
+            .data
+            .as_ref()
+            .unwrap()
+            .as_bool()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct PhysicalEmptyRelation(pub PlanNode);
+
+impl OptRelNode for PhysicalEmptyRelation {
+    fn into_rel_node(self) -> OptRelNodeRef {
+        replace_typ(self.0.into_rel_node(), OptRelNodeTyp::PhysicalEmptyRelation)
+    }
+
+    fn from_rel_node(rel_node: OptRelNodeRef) -> Option<Self> {
+        if rel_node.typ != OptRelNodeTyp::PhysicalEmptyRelation {
+            return None;
+        }
+        PlanNode::from_rel_node(rel_node).map(Self)
+    }
+
+    fn dispatch_explain(&self) -> Pretty<'static> {
+        Pretty::childless_record(
+            "PhysicalEmptyRelation",
+            vec![("produce_one_row", self.produce_one_row().to_string().into())],
+        )
+    }
+}
+
+impl PhysicalEmptyRelation {
+    pub fn new(node: PlanNode) -> PhysicalEmptyRelation {
+        Self(node)
+    }
+
+    pub fn produce_one_row(&self) -> bool {
+        self.clone()
+            .into_rel_node()
+            .data
+            .as_ref()
+            .unwrap()
+            .as_bool()
+    }
+}
diff --git a/optd-datafusion-repr/src/properties/schema.rs b/optd-datafusion-repr/src/properties/schema.rs
@@ -1,15 +1,11 @@
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex},
-};
-
 use optd_core::property::PropertyBuilder;
 
 use crate::plan_nodes::{ConstantType, OptRelNodeTyp};
 
 #[derive(Clone, Debug)]
 pub struct Schema(pub Vec<ConstantType>);
 
+// TODO: add names, nullable to schema
 impl Schema {
     pub fn len(&self) -> usize {
         self.0.len()
diff --git a/optd-datafusion-repr/src/rules/physical.rs b/optd-datafusion-repr/src/rules/physical.rs
@@ -40,6 +40,7 @@ impl PhysicalConversionRule {
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Filter)),
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Sort)),
             Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Agg)),
+            Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::EmptyRelation)),
         ]
     }
 }
@@ -117,6 +118,14 @@ impl<O: Optimizer<OptRelNodeTyp>> Rule<OptRelNodeTyp, O> for PhysicalConversionR
                 };
                 vec![node]
             }
+            OptRelNodeTyp::EmptyRelation => {
+                let node = RelNode {
+                    typ: OptRelNodeTyp::PhysicalEmptyRelation,
+                    children,
+                    data,
+                };
+                vec![node]
+            }
             _ => vec![],
         }
     }
diff --git a/optd-sqlplannertest/tests/empty_relation.planner.sql b/optd-sqlplannertest/tests/empty_relation.planner.sql
diff --git a/optd-sqlplannertest/tests/empty_relation.yml b/optd-sqlplannertest/tests/empty_relation.yml
diff --git a/optd-sqlplannertest/tests/join_enumerate.planner.sql b/optd-sqlplannertest/tests/join_enumerate.planner.sql

Original file line number	Diff line number	Diff line change
`@@ -106,6 +106,7 @@ impl CostModel<OptRelNodeTyp> for OptCostModel {`
`106`	`106`	`.unwrap_or(1) as f64;`
`107`	`107`	`Self::cost(row_cnt, 0.0, row_cnt)`
`108`	`108`	`}`
	`109`	`+ OptRelNodeTyp::PhysicalEmptyRelation => Self::cost(0.5, 0.01, 0.0),`
`109`	`110`	`OptRelNodeTyp::PhysicalFilter => {`
`110`	`111`	`let (row_cnt, _, _) = Self::cost_tuple(&children[0]);`
`111`	`112`	`let (_, compute_cost, _) = Self::cost_tuple(&children[1]);`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ impl PhysicalConversionRule {`
`40`	`40`	`Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Filter)),`
`41`	`41`	`Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Sort)),`
`42`	`42`	`Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::Agg)),`
	`43`	`+ Arc::new(PhysicalConversionRule::new(OptRelNodeTyp::EmptyRelation)),`
`43`	`44`	`]`
`44`	`45`	`}`
`45`	`46`	`}`
`@@ -117,6 +118,14 @@ impl<O: Optimizer<OptRelNodeTyp>> Rule<OptRelNodeTyp, O> for PhysicalConversionR`
`117`	`118`	`};`
`118`	`119`	`vec![node]`
`119`	`120`	`}`
	`121`	`+ OptRelNodeTyp::EmptyRelation => {`
	`122`	`+ let node = RelNode {`
	`123`	`+ typ: OptRelNodeTyp::PhysicalEmptyRelation,`
	`124`	`+ children,`
	`125`	`+ data,`
	`126`	`+ };`
	`127`	`+ vec![node]`
	`128`	`+ }`
`120`	`129`	`_ => vec![],`
`121`	`130`	`}`
`122`	`131`	`}`