Skip to content

Commit

Permalink
feat: improve datafusion SQL planning and parsing errors (#2959)
Browse files Browse the repository at this point in the history
Show as invalid input instead of server internal error
  • Loading branch information
LuQQiu authored Sep 30, 2024
1 parent c1edcae commit 8d58258
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 22 deletions.
38 changes: 32 additions & 6 deletions rust/lance-core/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// SPDX-FileCopyrightText: Copyright The Lance Authors

use arrow_schema::ArrowError;
use datafusion_common::DataFusionError;
use snafu::{Location, Snafu};

type BoxedError = Box<dyn std::error::Error + Send + Sync + 'static>;
Expand Down Expand Up @@ -305,20 +306,45 @@ impl From<datafusion_sql::sqlparser::tokenizer::TokenizerError> for Error {
}

#[cfg(feature = "datafusion")]
impl From<Error> for datafusion_common::DataFusionError {
impl From<Error> for DataFusionError {
#[track_caller]
fn from(e: Error) -> Self {
Self::Execution(e.to_string())
}
}

#[cfg(feature = "datafusion")]
impl From<datafusion_common::DataFusionError> for Error {
impl From<DataFusionError> for Error {
#[track_caller]
fn from(e: datafusion_common::DataFusionError) -> Self {
Self::IO {
source: box_error(e),
location: std::panic::Location::caller().to_snafu_location(),
fn from(e: DataFusionError) -> Self {
let location = std::panic::Location::caller().to_snafu_location();
match e {
DataFusionError::SQL(..)
| DataFusionError::Plan(..)
| DataFusionError::Configuration(..) => Self::InvalidInput {
source: box_error(e),
location,
},
DataFusionError::SchemaError(..) => Self::Schema {
message: e.to_string(),
location,
},
DataFusionError::ArrowError(..) => Self::Arrow {
message: e.to_string(),
location,
},
DataFusionError::NotImplemented(..) => Self::NotSupported {
source: box_error(e),
location,
},
DataFusionError::Execution(..) => Self::Execution {
message: e.to_string(),
location,
},
_ => Self::IO {
source: box_error(e),
location,
},
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion rust/lance-datafusion/src/logical_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use snafu::{location, Location};
fn resolve_value(expr: &Expr, data_type: &DataType) -> Result<Expr> {
match expr {
Expr::Literal(scalar_value) => {
Ok(Expr::Literal(safe_coerce_scalar(scalar_value, data_type).ok_or_else(|| Error::io(
Ok(Expr::Literal(safe_coerce_scalar(scalar_value, data_type).ok_or_else(|| Error::invalid_input(
format!("Received literal {expr} and could not convert to literal of type '{data_type:?}'"),
location!(),
))?))
Expand Down
30 changes: 15 additions & 15 deletions rust/lance-datafusion/src/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,7 @@ impl Planner {
BinaryOperator::And => Operator::And,
BinaryOperator::Or => Operator::Or,
_ => {
return Err(Error::io(
return Err(Error::invalid_input(
format!("Operator {op} is not supported"),
location!(),
));
Expand Down Expand Up @@ -306,7 +306,7 @@ impl Planner {
Err(_) => lit(-n
.parse::<f64>()
.map_err(|_e| {
Error::io(
Error::invalid_input(
format!("negative operator can be only applied to integer and float operands, got: {n}"),
location!(),
)
Expand All @@ -319,7 +319,7 @@ impl Planner {
}

_ => {
return Err(Error::io(
return Err(Error::invalid_input(
format!("Unary operator '{:?}' is not supported", op),
location!(),
));
Expand All @@ -339,7 +339,7 @@ impl Planner {
Ok(lit(n))
} else {
value.parse::<f64>().map(lit).map_err(|_| {
Error::io(
Error::invalid_input(
format!("'{value}' is not supported number value."),
location!(),
)
Expand All @@ -364,7 +364,7 @@ impl Planner {
fn parse_function_args(&self, func_args: &FunctionArg) -> Result<Expr> {
match func_args {
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => self.parse_sql_expr(expr),
_ => Err(Error::io(
_ => Err(Error::invalid_input(
format!("Unsupported function args: {:?}", func_args),
location!(),
)),
Expand All @@ -381,7 +381,7 @@ impl Planner {
match &func.args {
FunctionArguments::List(args) => {
if func.name.0.len() != 1 {
return Err(Error::io(
return Err(Error::invalid_input(
format!("Function name must have 1 part, got: {:?}", func.name.0),
location!(),
));
Expand All @@ -390,7 +390,7 @@ impl Planner {
self.parse_function_args(&args.args[0])?,
)))
}
_ => Err(Error::io(
_ => Err(Error::invalid_input(
format!("Unsupported function args: {:?}", &func.args),
location!(),
)),
Expand Down Expand Up @@ -456,7 +456,7 @@ impl Planner {
match tz {
TimezoneInfo::None => {}
_ => {
return Err(Error::io(
return Err(Error::invalid_input(
"Timezone not supported in timestamp".to_string(),
location!(),
));
Expand All @@ -470,7 +470,7 @@ impl Planner {
Some(6) => TimeUnit::Microsecond,
Some(9) => TimeUnit::Nanosecond,
_ => {
return Err(Error::io(
return Err(Error::invalid_input(
format!("Unsupported datetime resolution: {:?}", resolution),
location!(),
));
Expand All @@ -486,7 +486,7 @@ impl Planner {
Some(6) => TimeUnit::Microsecond,
Some(9) => TimeUnit::Nanosecond,
_ => {
return Err(Error::io(
return Err(Error::invalid_input(
format!("Unsupported datetime resolution: {:?}", resolution),
location!(),
));
Expand All @@ -498,15 +498,15 @@ impl Planner {
ExactNumberInfo::PrecisionAndScale(precision, scale) => {
Ok(ArrowDataType::Decimal128(*precision as u8, *scale as i8))
}
_ => Err(Error::io(
_ => Err(Error::invalid_input(
format!(
"Must provide precision and scale for decimal: {:?}",
number_info
),
location!(),
)),
},
_ => Err(Error::io(
_ => Err(Error::invalid_input(
format!(
"Unsupported data type: {:?}. Supported types: {:?}",
data_type, SUPPORTED_TYPES
Expand Down Expand Up @@ -539,7 +539,7 @@ impl Planner {
let mut values = vec![];

let array_literal_error = |pos: usize, value: &_| {
Err(Error::io(
Err(Error::invalid_input(
format!(
"Expected a literal value in array, instead got {} at position {}",
value, pos
Expand Down Expand Up @@ -582,7 +582,7 @@ impl Planner {

for value in &mut values {
if value.data_type() != data_type {
*value = safe_coerce_scalar(value, &data_type).ok_or_else(|| Error::io(
*value = safe_coerce_scalar(value, &data_type).ok_or_else(|| Error::invalid_input(
format!("Array expressions must have a consistent datatype. Expected: {}, got: {}", data_type, value.data_type()),
location!()
))?;
Expand Down Expand Up @@ -665,7 +665,7 @@ impl Planner {
expr: Box::new(self.parse_sql_expr(expr)?),
data_type: self.parse_type(data_type)?,
})),
_ => Err(Error::io(
_ => Err(Error::invalid_input(
format!("Expression '{expr}' is not supported SQL in lance"),
location!(),
)),
Expand Down

0 comments on commit 8d58258

Please sign in to comment.