From 9ff972e314f1f6397660e3bd5d54bef719d010ed Mon Sep 17 00:00:00 2001 From: ritchie Date: Sat, 13 Jul 2024 11:49:50 +0200 Subject: [PATCH] fix: Don't needlessly unwrap in `pivot_schema` --- crates/polars-plan/src/plans/functions/schema.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/polars-plan/src/plans/functions/schema.rs b/crates/polars-plan/src/plans/functions/schema.rs index fbb83c137733..9d33f12e737b 100644 --- a/crates/polars-plan/src/plans/functions/schema.rs +++ b/crates/polars-plan/src/plans/functions/schema.rs @@ -156,8 +156,8 @@ fn unpivot_schema<'a>( let mut new_schema = args .index .iter() - .map(|id| Field::new(id, input_schema.get(id).unwrap().clone())) - .collect::(); + .map(|id| Ok(Field::new(id, input_schema.try_get(id)?.clone()))) + .collect::>()?; let variable_name = args .variable_name .as_ref() @@ -179,13 +179,13 @@ fn unpivot_schema<'a>( let id_vars = PlHashSet::from_iter(&args.index); for (name, dtype) in input_schema.iter() { if !id_vars.contains(name) { - supertype = try_get_supertype(&supertype, dtype).unwrap(); + supertype = try_get_supertype(&supertype, dtype)?; } } } else { for name in &args.on { - let dtype = input_schema.get(name).unwrap(); - supertype = try_get_supertype(&supertype, dtype).unwrap(); + let dtype = input_schema.try_get(name)?; + supertype = try_get_supertype(&supertype, dtype)?; } } new_schema.with_column(value_name, supertype);