Skip to content

Commit

Permalink
feat: Improved error message DSL -> IR resolving (#19032)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ritchie46
ritchie46 authored Sep 30, 2024
1 parent 83f0127 commit f3578c4
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 35 deletions.
2 changes: 1 addition & 1 deletion crates/polars-error/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ impl PolarsError {
StructFieldNotFound(msg) => StructFieldNotFound(func(msg).into()),
SQLInterface(msg) => SQLInterface(func(msg).into()),
SQLSyntax(msg) => SQLSyntax(func(msg).into()),
_ => unreachable!(),
Context { error, .. } => error.wrap_msg(func),
}
}

Expand Down
70 changes: 40 additions & 30 deletions crates/polars-plan/src/plans/conversion/dsl_to_ir.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,23 +28,12 @@ fn empty_df() -> IR {
}
}

macro_rules! failed_input {
($($t:tt)*) => {
failed_input_args!(stringify!($($t)*))
}
}
macro_rules! failed_input_args {
($name:expr) => {
format!("'{}' input failed to resolve", $name).into()
};
}

macro_rules! failed_here {
($($t:tt)*) => {
format!("'{}' failed", stringify!($($t)*)).into()
format!("'{}'", stringify!($($t)*)).into()
}
}
pub(super) use {failed_here, failed_input, failed_input_args};
pub(super) use failed_here;

pub fn to_alp(
lp: DslPlan,
Expand All @@ -65,7 +54,29 @@ pub fn to_alp(
opt_flags,
};

to_alp_impl(lp, &mut ctxt)
match to_alp_impl(lp, &mut ctxt) {
Ok(out) => Ok(out),
Err(err) => {
if let Some(ir_until_then) = lp_arena.last_node() {
let node_name = if let PolarsError::Context { msg, .. } = &err {
msg
} else {
"THIS_NODE"
};
let plan = IRPlan::new(
ir_until_then,
std::mem::take(lp_arena),
std::mem::take(expr_arena),
);
let location = format!("{}", plan.display());
Err(err.wrap_msg(|msg| {
format!("{msg}\n\nResolved plan until failure:\n\n{location}\n\t---> FAILED HERE RESOLVING {node_name} <---")
}))
} else {
Err(err)
}
},
}
}

pub(super) struct DslConversionContext<'a> {
Expand Down Expand Up @@ -284,7 +295,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
.into_iter()
.map(|lp| to_alp_impl(lp, ctxt))
.collect::<PolarsResult<Vec<_>>>()
.map_err(|e| e.context(failed_input!(vertical concat)))?;
.map_err(|e| e.context(failed_here!(vertical concat)))?;

if args.diagonal {
inputs =
Expand All @@ -293,7 +304,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult

if args.to_supertypes {
convert_utils::convert_st_union(&mut inputs, ctxt.lp_arena, ctxt.expr_arena)
.map_err(|e| e.context(failed_input!(vertical concat)))?;
.map_err(|e| e.context(failed_here!(vertical concat)))?;
}
let options = args.into();
IR::Union { inputs, options }
Expand All @@ -303,7 +314,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
.into_iter()
.map(|lp| to_alp_impl(lp, ctxt))
.collect::<PolarsResult<Vec<_>>>()
.map_err(|e| e.context(failed_input!(horizontal concat)))?;
.map_err(|e| e.context(failed_here!(horizontal concat)))?;

let schema = convert_utils::h_concat_schema(&inputs, ctxt.lp_arena)?;

Expand All @@ -315,7 +326,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
},
DslPlan::Filter { input, predicate } => {
let mut input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(filter)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(filter)))?;
let predicate = expand_filter(predicate, input, ctxt.lp_arena, ctxt.opt_flags)
.map_err(|e| e.context(failed_here!(filter)))?;

Expand Down Expand Up @@ -365,7 +376,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
},
DslPlan::Slice { input, offset, len } => {
let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(slice)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(slice)))?;
IR::Slice { input, offset, len }
},
DslPlan::DataFrameScan { df, schema } => IR::DataFrameScan {
Expand All @@ -380,7 +391,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
options,
} => {
let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(select)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(select)))?;
let schema = ctxt.lp_arena.get(input).schema(ctxt.lp_arena);
let (exprs, schema) = prepare_projection(expr, &schema, ctxt.opt_flags)
.map_err(|e| e.context(failed_here!(select)))?;
Expand Down Expand Up @@ -430,7 +441,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
);

let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(sort)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(sort)))?;

let mut expanded_cols = Vec::new();
let mut nulls_last = Vec::new();
Expand Down Expand Up @@ -477,7 +488,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
},
DslPlan::Cache { input, id } => {
let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(cache)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(cache)))?;
IR::Cache {
input,
id,
Expand All @@ -493,7 +504,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
options,
} => {
let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(group_by)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(group_by)))?;

let (keys, aggs, schema) = resolve_group_by(
input,
Expand Down Expand Up @@ -555,7 +566,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
options,
} => {
let input = to_alp_impl(owned(input), ctxt)
.map_err(|e| e.context(failed_input!(with_columns)))?;
.map_err(|e| e.context(failed_here!(with_columns)))?;
let (exprs, schema) =
resolve_with_columns(exprs, input, ctxt.lp_arena, ctxt.expr_arena, ctxt.opt_flags)
.map_err(|e| e.context(failed_here!(with_columns)))?;
Expand All @@ -572,7 +583,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
},
DslPlan::Distinct { input, options } => {
let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(unique)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(unique)))?;
let input_schema = ctxt.lp_arena.get(input).schema(ctxt.lp_arena);

let subset = options
Expand All @@ -589,9 +600,8 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
IR::Distinct { input, options }
},
DslPlan::MapFunction { input, function } => {
let input = to_alp_impl(owned(input), ctxt).map_err(|e| {
e.context(failed_input_args!(format!("{}", function).to_lowercase()))
})?;
let input = to_alp_impl(owned(input), ctxt)
.map_err(|e| e.context(failed_here!(format!("{}", function).to_lowercase())))?;
let input_schema = ctxt.lp_arena.get(input).schema(ctxt.lp_arena);

match function {
Expand Down Expand Up @@ -755,7 +765,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
},
DslPlan::ExtContext { input, contexts } => {
let input = to_alp_impl(owned(input), ctxt)
.map_err(|e| e.context(failed_input!(with_context)))?;
.map_err(|e| e.context(failed_here!(with_context)))?;
let contexts = contexts
.into_iter()
.map(|lp| to_alp_impl(lp, ctxt))
Expand All @@ -780,7 +790,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
},
DslPlan::Sink { input, payload } => {
let input =
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(sink)))?;
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(sink)))?;
IR::Sink { input, payload }
},
DslPlan::IR { node, dsl, version } => {
Expand Down
8 changes: 4 additions & 4 deletions crates/polars-plan/src/plans/conversion/join.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,10 +75,10 @@ pub fn resolve_join(
}

let input_left = input_left.map_right(Ok).right_or_else(|input| {
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(join left)))
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(join left)))
})?;
let input_right = input_right.map_right(Ok).right_or_else(|input| {
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_input!(join right)))
to_alp_impl(owned(input), ctxt).map_err(|e| e.context(failed_here!(join right)))
})?;

let schema_left = ctxt.lp_arena.get(input_left).schema(ctxt.lp_arena);
Expand Down Expand Up @@ -153,9 +153,9 @@ fn resolve_join_where(
) -> PolarsResult<Node> {
check_join_keys(&predicates)?;
let input_left = to_alp_impl(Arc::unwrap_or_clone(input_left), ctxt)
.map_err(|e| e.context(failed_input!(join left)))?;
.map_err(|e| e.context(failed_here!(join left)))?;
let input_right = to_alp_impl(Arc::unwrap_or_clone(input_right), ctxt)
.map_err(|e| e.context(failed_input!(join left)))?;
.map_err(|e| e.context(failed_here!(join left)))?;

let schema_left = ctxt.lp_arena.get(input_left).schema(ctxt.lp_arena);
let schema_right = ctxt
Expand Down
8 changes: 8 additions & 0 deletions crates/polars-utils/src/arena.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,14 @@ impl<T> Arena<T> {
self.items.pop()
}

pub fn last_node(&mut self) -> Option<Node> {
if self.is_empty() {
None
} else {
Some(Node(self.items.len() - 1))
}
}

pub fn len(&self) -> usize {
self.items.len()
}
Expand Down

0 comments on commit f3578c4

Please sign in to comment.