Skip to content

Commit

Permalink
Merge pull request #5362 from leiysky/negate-unary-functions
Browse files Browse the repository at this point in the history
feature(planner): Support some scalar expressions in new planner
  • Loading branch information
BohuTANG authored May 13, 2022
2 parents b8e5194 + b6fc1a1 commit cda4e21
Show file tree
Hide file tree
Showing 8 changed files with 167 additions and 26 deletions.
17 changes: 14 additions & 3 deletions query/src/sql/exec/expression_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.

use common_datavalues::DataSchemaRef;
use common_datavalues::DataTypeImpl;
use common_datavalues::DataValue;
use common_exception::ErrorCode;
use common_exception::Result;
Expand Down Expand Up @@ -67,7 +68,9 @@ impl<'a> ExpressionBuilder<'a> {
Scalar::BoundColumnRef(BoundColumnRef { column }) => {
self.build_column_ref(column.index)
}
Scalar::ConstantExpr(ConstantExpr { value }) => self.build_literal(value),
Scalar::ConstantExpr(ConstantExpr { value, data_type }) => {
self.build_literal(value, data_type)
}
Scalar::ComparisonExpr(ComparisonExpr { op, left, right }) => {
self.build_binary_operator(left, right, op.to_func_name())
}
Expand Down Expand Up @@ -134,8 +137,16 @@ impl<'a> ExpressionBuilder<'a> {
)))
}

pub fn build_literal(&self, data_value: &DataValue) -> Result<Expression> {
Ok(Expression::create_literal(data_value.clone()))
pub fn build_literal(
&self,
data_value: &DataValue,
data_type: &DataTypeImpl,
) -> Result<Expression> {
Ok(Expression::Literal {
value: data_value.clone(),
column_name: None,
data_type: data_type.clone(),
})
}

pub fn build_binary_operator(
Expand Down
15 changes: 11 additions & 4 deletions query/src/sql/planner/binder/limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.

use common_ast::ast::Expr;
use common_datavalues::DataType;
use common_exception::ErrorCode;
use common_exception::Result;

Expand All @@ -34,8 +35,11 @@ impl<'a> Binder {

let limit_cnt = match limit {
Some(Expr::Literal { span: _, lit: x }) => {
let data = type_checker.resolve_literal(x, None)?.as_u64()?;
Some(data as usize)
let (value, data_type) = type_checker.resolve_literal(x, None)?;
if !data_type.data_type_id().is_integer() {
return Err(ErrorCode::IllegalDataType("Unsupported limit type"));
}
Some(value.as_u64()? as usize)
}
Some(_) => {
return Err(ErrorCode::IllegalDataType("Unsupported limit type"));
Expand All @@ -44,8 +48,11 @@ impl<'a> Binder {
};

let offset_cnt = if let Some(Expr::Literal { span: _, lit: x }) = offset {
let data = type_checker.resolve_literal(x, None)?.as_u64()?;
data as usize
let (value, data_type) = type_checker.resolve_literal(x, None)?;
if !data_type.data_type_id().is_integer() {
return Err(ErrorCode::IllegalDataType("Unsupported limit type"));
}
value.as_u64()? as usize
} else {
0
};
Expand Down
12 changes: 7 additions & 5 deletions query/src/sql/planner/binder/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,11 +224,13 @@ impl<'a> Binder {
let expressions = args
.into_iter()
.map(|(scalar, _)| match scalar {
Scalar::ConstantExpr(ConstantExpr { value }) => Ok(Expression::Literal {
value: value.clone(),
column_name: None,
data_type: value.data_type(),
}),
Scalar::ConstantExpr(ConstantExpr { value, data_type }) => {
Ok(Expression::Literal {
value,
column_name: None,
data_type,
})
}
_ => Err(ErrorCode::UnImplement(format!(
"Unsupported table argument type: {:?}",
scalar
Expand Down
4 changes: 3 additions & 1 deletion query/src/sql/planner/plans/scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,11 +72,13 @@ impl ScalarExpr for BoundColumnRef {
#[derive(Clone, PartialEq, Debug)]
pub struct ConstantExpr {
pub value: DataValue,

pub data_type: DataTypeImpl,
}

impl ScalarExpr for ConstantExpr {
fn data_type(&self) -> DataTypeImpl {
self.value.data_type()
self.data_type.clone()
}

fn used_columns(&self) -> ColumnSet {
Expand Down
137 changes: 124 additions & 13 deletions query/src/sql/planner/semantic/type_check.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
use std::sync::Arc;

use common_ast::ast::BinaryOperator;
use common_ast::ast::DateTimeField;
use common_ast::ast::Expr;
use common_ast::ast::Literal;
use common_ast::ast::MapAccessor;
Expand All @@ -25,6 +26,9 @@ use common_datavalues::BooleanType;
use common_datavalues::DataField;
use common_datavalues::DataTypeImpl;
use common_datavalues::DataValue;
use common_datavalues::IntervalKind;
use common_datavalues::IntervalType;
use common_datavalues::TimestampType;
use common_exception::ErrorCode;
use common_exception::Result;
use common_functions::aggregates::AggregateFunctionFactory;
Expand Down Expand Up @@ -184,11 +188,8 @@ impl<'a> TypeChecker<'a> {
expr, target_type, ..
} => {
let (scalar, data_type) = self.resolve(expr, required_type).await?;
let cast_func = CastFunction::create_try(
"",
target_type.to_string().as_str(),
data_type.clone(),
)?;
let cast_func =
CastFunction::create("", target_type.to_string().as_str(), data_type.clone())?;
Ok((
CastExpr {
argument: Box::new(scalar),
Expand Down Expand Up @@ -223,9 +224,15 @@ impl<'a> TypeChecker<'a> {
}

Expr::Literal { lit, .. } => {
let value = self.resolve_literal(lit, required_type)?;
let data_type = value.data_type();
Ok((ConstantExpr { value }.into(), data_type))
let (value, data_type) = self.resolve_literal(lit, required_type)?;
Ok((
ConstantExpr {
value,
data_type: data_type.clone(),
}
.into(),
data_type,
))
}

Expr::FunctionCall {
Expand All @@ -250,7 +257,7 @@ impl<'a> TypeChecker<'a> {
// Check aggregate function
let params = params
.iter()
.map(|literal| self.resolve_literal(literal, None))
.map(|literal| self.resolve_literal(literal, None).map(|(value, _)| value))
.collect::<Result<Vec<DataValue>>>()?;

let mut arguments = vec![];
Expand Down Expand Up @@ -333,6 +340,30 @@ impl<'a> TypeChecker<'a> {
Ok(self.resolve_function("get", &[&**expr, &arg], None).await?)
}

Expr::TryCast {
expr, target_type, ..
} => {
let (scalar, data_type) = self.resolve(expr, required_type).await?;
let cast_func = CastFunction::create_try(
"",
target_type.to_string().as_str(),
data_type.clone(),
)?;
Ok((
CastExpr {
argument: Box::new(scalar),
from_type: data_type,
target_type: cast_func.return_type(),
}
.into(),
cast_func.return_type(),
))
}

Expr::Extract { field, expr, .. } => {
self.resolve_extract_expr(field, expr, required_type).await
}

_ => Err(ErrorCode::UnImplement(format!(
"Unsupported expr: {:?}",
expr
Expand Down Expand Up @@ -458,8 +489,65 @@ impl<'a> TypeChecker<'a> {
child: &Expr<'a>,
required_type: Option<DataTypeImpl>,
) -> Result<(Scalar, DataTypeImpl)> {
self.resolve_function(op.to_string().as_str(), &[child], required_type)
.await
match op {
UnaryOperator::Plus => {
// Omit unary + operator
self.resolve(child, required_type).await
}

UnaryOperator::Minus => {
self.resolve_function("negate", &[child], required_type)
.await
}

UnaryOperator::Not => self.resolve_function("not", &[child], required_type).await,
}
}

pub async fn resolve_extract_expr(
&mut self,
date_field: &DateTimeField,
arg: &Expr<'a>,
_required_type: Option<DataTypeImpl>,
) -> Result<(Scalar, DataTypeImpl)> {
match date_field {
DateTimeField::Year => {
self.resolve_function("toYear", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Month => {
self.resolve_function("toMonth", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Day => {
self.resolve_function("toDayOfMonth", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Hour => {
self.resolve_function("toHour", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Minute => {
self.resolve_function("toMinute", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Second => {
self.resolve_function("toSecond", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Dow => {
self.resolve_function("toDayOfWeek", &[arg], Some(TimestampType::new_impl(0)))
.await
}
DateTimeField::Doy => {
self.resolve_function("toDayOfYear", &[arg], Some(TimestampType::new_impl(0)))
.await
}
_ => Err(ErrorCode::SemanticError(format!(
"Invalid time field: {}",
date_field
))),
}
}

pub async fn resolve_subquery(
Expand Down Expand Up @@ -497,18 +585,41 @@ impl<'a> TypeChecker<'a> {
&self,
literal: &Literal,
_required_type: Option<DataTypeImpl>,
) -> Result<DataValue> {
) -> Result<(DataValue, DataTypeImpl)> {
// TODO(leiysky): try cast value to required type
let value = match literal {
Literal::Number(string) => DataValue::try_from_literal(string, None)?,
Literal::String(string) => DataValue::String(string.as_bytes().to_vec()),
Literal::Boolean(boolean) => DataValue::Boolean(*boolean),
Literal::Null => DataValue::Null,
Literal::Interval(interval) => {
let num = interval.value.parse::<i64>()?;
DataValue::Int64(num)
}
_ => Err(ErrorCode::SemanticError(format!(
"Unsupported literal value: {literal}"
)))?,
};

Ok(value)
let data_type = if let Literal::Interval(interval) = literal {
match &interval.field {
DateTimeField::Year => IntervalType::new_impl(IntervalKind::Year),
DateTimeField::Month => IntervalType::new_impl(IntervalKind::Month),
DateTimeField::Day => IntervalType::new_impl(IntervalKind::Day),
DateTimeField::Hour => IntervalType::new_impl(IntervalKind::Hour),
DateTimeField::Minute => IntervalType::new_impl(IntervalKind::Minute),
DateTimeField::Second => IntervalType::new_impl(IntervalKind::Second),
_ => {
return Err(ErrorCode::SemanticError(format!(
"Invalid interval kind: {}",
interval.field
)));
}
}
} else {
value.data_type()
};

Ok((value, data_type))
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ insert into table tt values ('2021-04-30 22:48:00'), (to_timestamp('2021-04-30 2
select * from tt;
set timezone = 'Asia/Shanghai';
select * from tt;
drop table tt;
-- number function
-- 1619820000000000 = 2021-04-30 22:00:00
select "====NUMBER_FUNCTION====";
Expand Down
3 changes: 3 additions & 0 deletions tests/suites/0_stateless/20+_others/20_0001_planner_v2.result
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,13 @@
====ALIAS====
0 1
0 1
====SCALAR_EXPRESSION====
1 13
====COMPARISON====
5
====CAST====
5
5
====BINARY_OPERATOR====
-0.75
====FUNCTIONS====
Expand Down
4 changes: 4 additions & 0 deletions tests/suites/0_stateless/20+_others/20_0001_planner_v2.sql
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,17 @@ select '====ALIAS====';
select number as a, number + 1 as b from numbers(1);
select number as a, number + 1 as b from numbers(1) group by a, number order by number;

select '====SCALAR_EXPRESSION====';
select interval '1' day, extract(day from to_date('2022-05-13'));

-- Comparison expressions
select '====COMPARISON====';
select * from numbers(10) where number between 1 and 9 and number > 2 and number < 8 and number is not null and number = 5 and number >= 5 and number <= 5;

-- Cast expression
select '====CAST====';
select * from numbers(10) where cast(number as string) = '5';
select * from numbers(10) where try_cast(number as string) = '5';

-- Binary operator
select '====BINARY_OPERATOR====';
Expand Down

0 comments on commit cda4e21

Please sign in to comment.