use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
use datafusion_common::{DFSchema, DataFusionError, Result, ScalarValue};
use datafusion_expr::{BuiltinScalarFunction, Expr};
use sqlparser::ast::Expr as SQLExpr;
impl<'a, S: ContextProvider> SqlToRel<'a, S> {
pub(super) fn sql_substring_to_expr(
&self,
expr: Box<SQLExpr>,
substring_from: Option<Box<SQLExpr>>,
substring_for: Option<Box<SQLExpr>>,
schema: &DFSchema,
planner_context: &mut PlannerContext,
) -> Result<Expr> {
let args = match (substring_from, substring_for) {
(Some(from_expr), Some(for_expr)) => {
let arg =
self.sql_expr_to_logical_expr(*expr, schema, planner_context)?;
let from_logic =
self.sql_expr_to_logical_expr(*from_expr, schema, planner_context)?;
let for_logic =
self.sql_expr_to_logical_expr(*for_expr, schema, planner_context)?;
vec![arg, from_logic, for_logic]
}
(Some(from_expr), None) => {
let arg =
self.sql_expr_to_logical_expr(*expr, schema, planner_context)?;
let from_logic =
self.sql_expr_to_logical_expr(*from_expr, schema, planner_context)?;
vec![arg, from_logic]
}
(None, Some(for_expr)) => {
let arg =
self.sql_expr_to_logical_expr(*expr, schema, planner_context)?;
let from_logic = Expr::Literal(ScalarValue::Int64(Some(1)));
let for_logic =
self.sql_expr_to_logical_expr(*for_expr, schema, planner_context)?;
vec![arg, from_logic, for_logic]
}
(None, None) => {
let orig_sql = SQLExpr::Substring {
expr,
substring_from: None,
substring_for: None,
};
return Err(DataFusionError::Plan(format!(
"Substring without for/from is not valid {orig_sql:?}"
)));
}
};
Ok(Expr::ScalarFunction {
fun: BuiltinScalarFunction::Substr,
args,
})
}
}