Skip to content
This repository was archived by the owner on Jan 7, 2025. It is now read-only.

Commit 214848c

Browse files
Gun9niRGun9niR
authored
feat: support tpch q5 and q6 (#69)
- Support gt and lt for bin op type - Support `LikeExpr` (it's for Q2, but it turns out Q2 requires more work) - Support for cast types
1 parent d34c22c commit 214848c

File tree

6 files changed

+365
-8
lines changed

6 files changed

+365
-8
lines changed

‎optd-datafusion-bridge/src/from_optd.rs‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use datafusion::{
2323
use optd_datafusion_repr::{
2424
plan_nodes::{
2525
BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, ConstantType,
26-
Expr, ExprList, FuncExpr, FuncType, JoinType, LogOpExpr, LogOpType, OptRelNode,
26+
Expr, ExprList, FuncExpr, FuncType, JoinType, LikeExpr, LogOpExpr, LogOpType, OptRelNode,
2727
OptRelNodeRef, OptRelNodeTyp, PhysicalAgg, PhysicalEmptyRelation, PhysicalFilter,
2828
PhysicalHashJoin, PhysicalLimit, PhysicalNestedLoopJoin, PhysicalProjection, PhysicalScan,
2929
PhysicalSort, PlanNode, SortOrderExpr, SortOrderType,
@@ -49,6 +49,7 @@ fn from_optd_schema(optd_schema: OptdSchema) -> Schema {
4949
ConstantType::Int16 => DataType::Int16,
5050
ConstantType::Int32 => DataType::Int32,
5151
ConstantType::Int64 => DataType::Int64,
52+
ConstantType::Float64 => DataType::Float64,
5253
ConstantType::Date => DataType::Date32,
5354
ConstantType::Decimal => DataType::Float64,
5455
ConstantType::Utf8String => DataType::Utf8,
@@ -145,6 +146,7 @@ impl OptdPlanContext<'_> {
145146
ConstantType::Int16 => ScalarValue::Int16(Some(value.as_i16())),
146147
ConstantType::Int32 => ScalarValue::Int32(Some(value.as_i32())),
147148
ConstantType::Int64 => ScalarValue::Int64(Some(value.as_i64())),
149+
ConstantType::Float64 => ScalarValue::Float64(Some(value.as_f64())),
148150
ConstantType::Decimal => {
149151
ScalarValue::Decimal128(Some(value.as_f64() as i128), 20, 0)
150152
// TODO(chi): no hard code decimal
@@ -214,7 +216,9 @@ impl OptdPlanContext<'_> {
214216
BinOpType::Eq => Operator::Eq,
215217
BinOpType::Neq => Operator::NotEq,
216218
BinOpType::Leq => Operator::LtEq,
219+
BinOpType::Lt => Operator::Lt,
217220
BinOpType::Geq => Operator::GtEq,
221+
BinOpType::Gt => Operator::Gt,
218222
BinOpType::And => Operator::And,
219223
BinOpType::Add => Operator::Plus,
220224
BinOpType::Sub => Operator::Minus,
@@ -256,6 +260,19 @@ impl OptdPlanContext<'_> {
256260
datafusion::physical_plan::expressions::CastExpr::new(child, data_type, None),
257261
))
258262
}
263+
OptRelNodeTyp::Like => {
264+
let expr = LikeExpr::from_rel_node(expr.into_rel_node()).unwrap();
265+
let child = Self::conv_from_optd_expr(expr.child(), context)?;
266+
let pattern = Self::conv_from_optd_expr(expr.pattern(), context)?;
267+
Ok(Arc::new(
268+
datafusion::physical_plan::expressions::LikeExpr::new(
269+
expr.negated(),
270+
expr.case_insensitive(),
271+
child,
272+
pattern,
273+
),
274+
))
275+
}
259276
_ => unimplemented!("{}", expr.into_rel_node()),
260277
}
261278
}

‎optd-datafusion-bridge/src/into_optd.rs‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,10 @@ use optd_core::rel_node::RelNode;
99
use optd_datafusion_repr::{
1010
plan_nodes::{
1111
BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, Expr, ExprList,
12-
FuncExpr, FuncType, JoinType, LogOpExpr, LogOpType, LogicalAgg, LogicalEmptyRelation,
13-
LogicalFilter, LogicalJoin, LogicalLimit, LogicalProjection, LogicalScan, LogicalSort,
14-
OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode, SortOrderExpr, SortOrderType,
12+
FuncExpr, FuncType, JoinType, LikeExpr, LogOpExpr, LogOpType, LogicalAgg,
13+
LogicalEmptyRelation, LogicalFilter, LogicalJoin, LogicalLimit, LogicalProjection,
14+
LogicalScan, LogicalSort, OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode,
15+
SortOrderExpr, SortOrderType,
1516
},
1617
Value,
1718
};
@@ -54,7 +55,9 @@ impl OptdPlanContext<'_> {
5455
Operator::Eq => BinOpType::Eq,
5556
Operator::NotEq => BinOpType::Neq,
5657
Operator::LtEq => BinOpType::Leq,
58+
Operator::Lt => BinOpType::Lt,
5759
Operator::GtEq => BinOpType::Geq,
60+
Operator::Gt => BinOpType::Gt,
5861
Operator::And => BinOpType::And,
5962
Operator::Plus => BinOpType::Add,
6063
Operator::Minus => BinOpType::Sub,
@@ -101,6 +104,10 @@ impl OptdPlanContext<'_> {
101104
let x = x.as_ref().unwrap();
102105
Ok(ConstantExpr::int64(*x).into_expr())
103106
}
107+
ScalarValue::Float64(x) => {
108+
let x = x.as_ref().unwrap();
109+
Ok(ConstantExpr::float64(*x).into_expr())
110+
}
104111
ScalarValue::Utf8(x) => {
105112
let x = x.as_ref().unwrap();
106113
Ok(ConstantExpr::string(x).into_expr())
@@ -166,12 +173,24 @@ impl OptdPlanContext<'_> {
166173
let data_type = x.data_type.clone();
167174
let val = match data_type {
168175
arrow_schema::DataType::Int8 => Value::Int8(0),
176+
arrow_schema::DataType::Int16 => Value::Int16(0),
177+
arrow_schema::DataType::Int32 => Value::Int32(0),
178+
arrow_schema::DataType::Int64 => Value::Int64(0),
179+
arrow_schema::DataType::UInt8 => Value::UInt8(0),
180+
arrow_schema::DataType::UInt16 => Value::UInt16(0),
181+
arrow_schema::DataType::UInt32 => Value::UInt32(0),
182+
arrow_schema::DataType::UInt64 => Value::UInt64(0),
169183
arrow_schema::DataType::Date32 => Value::Date32(0),
170184
arrow_schema::DataType::Decimal128(_, _) => Value::Decimal128(0),
171185
other => unimplemented!("unimplemented datatype {:?}", other),
172186
};
173187
Ok(CastExpr::new(expr, val).into_expr())
174188
}
189+
Expr::Like(x) => {
190+
let expr = self.conv_into_optd_expr(x.expr.as_ref(), context)?;
191+
let pattern = self.conv_into_optd_expr(x.pattern.as_ref(), context)?;
192+
Ok(LikeExpr::new(x.negated, x.case_insensitive, expr, pattern).into_expr())
193+
}
175194
_ => bail!("Unsupported expression: {:?}", expr),
176195
}
177196
}

‎optd-datafusion-repr/src/plan_nodes.rs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@ pub use apply::{ApplyType, LogicalApply};
2424
pub use empty_relation::{LogicalEmptyRelation, PhysicalEmptyRelation};
2525
pub use expr::{
2626
BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, ConstantType,
27-
ExprList, FuncExpr, FuncType, LogOpExpr, LogOpType, SortOrderExpr, SortOrderType, UnOpExpr,
28-
UnOpType,
27+
ExprList, FuncExpr, FuncType, LikeExpr, LogOpExpr, LogOpType, SortOrderExpr, SortOrderType,
28+
UnOpExpr, UnOpType,
2929
};
3030
pub use filter::{LogicalFilter, PhysicalFilter};
3131
pub use join::{JoinType, LogicalJoin, PhysicalHashJoin, PhysicalNestedLoopJoin};
@@ -76,6 +76,7 @@ pub enum OptRelNodeTyp {
7676
SortOrder(SortOrderType),
7777
Between,
7878
Cast,
79+
Like,
7980
}
8081

8182
impl OptRelNodeTyp {
@@ -116,6 +117,7 @@ impl OptRelNodeTyp {
116117
| Self::LogOp(_)
117118
| Self::Between
118119
| Self::Cast
120+
| Self::Like
119121
)
120122
}
121123
}
@@ -382,6 +384,9 @@ pub fn explain(rel_node: OptRelNodeRef) -> Pretty<'static> {
382384
OptRelNodeTyp::Cast => CastExpr::from_rel_node(rel_node)
383385
.unwrap()
384386
.dispatch_explain(),
387+
OptRelNodeTyp::Like => LikeExpr::from_rel_node(rel_node)
388+
.unwrap()
389+
.dispatch_explain(),
385390
}
386391
}
387392

‎optd-datafusion-repr/src/plan_nodes/expr.rs‎

Lines changed: 72 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::fmt::Display;
1+
use std::{fmt::Display, sync::Arc};
22

33
use itertools::Itertools;
44
use pretty_xmlish::Pretty;
@@ -76,6 +76,7 @@ pub enum ConstantType {
7676
Int16,
7777
Int32,
7878
Int64,
79+
Float64,
7980
Date,
8081
Decimal,
8182
Any,
@@ -97,7 +98,7 @@ impl ConstantExpr {
9798
Value::Int16(_) => ConstantType::Int16,
9899
Value::Int32(_) => ConstantType::Int32,
99100
Value::Int64(_) => ConstantType::Int64,
100-
Value::Float(_) => ConstantType::Decimal,
101+
Value::Float(_) => ConstantType::Float64,
101102
_ => unimplemented!(),
102103
};
103104
Self::new_with_type(value, typ)
@@ -157,6 +158,10 @@ impl ConstantExpr {
157158
Self::new_with_type(Value::Int64(value), ConstantType::Int64)
158159
}
159160

161+
pub fn float64(value: f64) -> Self {
162+
Self::new_with_type(Value::Float(value.into()), ConstantType::Float64)
163+
}
164+
160165
pub fn date(value: i64) -> Self {
161166
Self::new_with_type(Value::Int64(value), ConstantType::Date)
162167
}
@@ -681,3 +686,68 @@ impl OptRelNode for CastExpr {
681686
)
682687
}
683688
}
689+
690+
#[derive(Clone, Debug)]
691+
pub struct LikeExpr(pub Expr);
692+
693+
impl LikeExpr {
694+
pub fn new(negated: bool, case_insensitive: bool, expr: Expr, pattern: Expr) -> Self {
695+
// TODO: support multiple values in data.
696+
let negated = if negated { 1 } else { 0 };
697+
let case_insensitive = if case_insensitive { 1 } else { 0 };
698+
LikeExpr(Expr(
699+
RelNode {
700+
typ: OptRelNodeTyp::Like,
701+
children: vec![expr.into_rel_node(), pattern.into_rel_node()],
702+
data: Some(Value::Serialized(Arc::new([negated, case_insensitive]))),
703+
}
704+
.into(),
705+
))
706+
}
707+
708+
pub fn child(&self) -> Expr {
709+
Expr(self.0.child(0))
710+
}
711+
712+
pub fn pattern(&self) -> Expr {
713+
Expr(self.0.child(1))
714+
}
715+
716+
pub fn negated(&self) -> bool {
717+
match self.0 .0.data.as_ref().unwrap() {
718+
Value::Serialized(data) => data[0] != 0,
719+
_ => panic!("not a serialized value"),
720+
}
721+
}
722+
723+
pub fn case_insensitive(&self) -> bool {
724+
match self.0 .0.data.as_ref().unwrap() {
725+
Value::Serialized(data) => data[1] != 0,
726+
_ => panic!("not a serialized value"),
727+
}
728+
}
729+
}
730+
731+
impl OptRelNode for LikeExpr {
732+
fn into_rel_node(self) -> OptRelNodeRef {
733+
self.0.into_rel_node()
734+
}
735+
736+
fn from_rel_node(rel_node: OptRelNodeRef) -> Option<Self> {
737+
if !matches!(rel_node.typ, OptRelNodeTyp::Like) {
738+
return None;
739+
}
740+
Expr::from_rel_node(rel_node).map(Self)
741+
}
742+
743+
fn dispatch_explain(&self) -> Pretty<'static> {
744+
Pretty::simple_record(
745+
"Like",
746+
vec![
747+
("expr", self.child().explain()),
748+
("pattern", self.pattern().explain()),
749+
],
750+
vec![],
751+
)
752+
}
753+
}

0 commit comments

Comments
 (0)