feat(rust,python): better-characterise NotFound exceptions (#6670)

pola-rs · Feb 4, 2023 · 3a6c999 · 3a6c999
1 parent c7db9e0
commit 3a6c999
Show file tree

Hide file tree

Showing 21 changed files with 219 additions and 163 deletions.
diff --git a/polars/polars-core/src/chunked_array/logical/struct_/mod.rs b/polars/polars-core/src/chunked_array/logical/struct_/mod.rs
@@ -147,7 +147,7 @@ impl StructChunked {
         self.fields
             .iter()
             .find(|s| s.name() == name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+            .ok_or_else(|| PolarsError::StructFieldNotFound(name.to_string().into()))
             .map(|s| s.clone())
     }
 

diff --git a/polars/polars-core/src/error.rs b/polars/polars-core/src/error.rs
@@ -55,22 +55,26 @@ impl Display for ErrString {
 pub enum PolarsError {
     #[error(transparent)]
     ArrowError(Box<ArrowError>),
+    #[error("Not found: {0}")]
+    ColumnNotFound(ErrString),
+    #[error("{0}")]
+    ComputeError(ErrString),
+    #[error("DuplicateError: {0}")]
+    Duplicate(ErrString),
     #[error("Invalid operation {0}")]
     InvalidOperation(ErrString),
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+    #[error("Such empty...: {0}")]
+    NoData(ErrString),
+    #[error("Not found: {0}")]
+    SchemaFieldNotFound(ErrString),
     #[error("Data types don't match: {0}")]
     SchemaMisMatch(ErrString),
-    #[error("Not found: {0}")]
-    NotFound(ErrString),
     #[error("Lengths don't match: {0}")]
     ShapeMisMatch(ErrString),
-    #[error("{0}")]
-    ComputeError(ErrString),
-    #[error("Such empty...: {0}")]
-    NoData(ErrString),
-    #[error(transparent)]
-    Io(#[from] std::io::Error),
-    #[error("DuplicateError: {0}")]
-    Duplicate(ErrString),
+    #[error("Not found: {0}")]
+    StructFieldNotFound(ErrString),
 }
 
 impl From<ArrowError> for PolarsError {
@@ -105,15 +109,17 @@ impl PolarsError {
     pub fn wrap_msg(&self, func: &dyn Fn(&str) -> String) -> Self {
         use PolarsError::*;
         match self {
-            ComputeError(msg) => ComputeError(func(msg).into()),
             ArrowError(err) => ComputeError(func(&format!("ArrowError: {err}")).into()),
+            ColumnNotFound(msg) => ColumnNotFound(func(msg).into()),
+            ComputeError(msg) => ComputeError(func(msg).into()),
+            Duplicate(msg) => Duplicate(func(msg).into()),
             InvalidOperation(msg) => InvalidOperation(func(msg).into()),
+            Io(err) => ComputeError(func(&format!("IO: {err}")).into()),
+            NoData(msg) => NoData(func(msg).into()),
+            SchemaFieldNotFound(msg) => SchemaFieldNotFound(func(msg).into()),
             SchemaMisMatch(msg) => SchemaMisMatch(func(msg).into()),
-            NotFound(msg) => NotFound(func(msg).into()),
             ShapeMisMatch(msg) => ShapeMisMatch(func(msg).into()),
-            NoData(msg) => NoData(func(msg).into()),
-            Io(err) => ComputeError(func(&format!("IO: {err}")).into()),
-            Duplicate(msg) => Duplicate(func(msg).into()),
+            StructFieldNotFound(msg) => StructFieldNotFound(func(msg).into()),
         }
     }
 }
diff --git a/polars/polars-core/src/frame/explode.rs b/polars/polars-core/src/frame/explode.rs
@@ -253,7 +253,7 @@ impl DataFrame {
         let mut iter = value_vars.iter().map(|v| {
             schema
                 .get(v)
-                .ok_or_else(|| PolarsError::NotFound(v.to_string().into()))
+                .ok_or_else(|| PolarsError::ColumnNotFound(v.to_string().into()))
         });
         let mut st = iter.next().unwrap()?.clone();
         for dt in iter {

diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -186,7 +186,7 @@ impl DataFrame {
     /// Get the index of the column.
     fn check_name_to_idx(&self, name: &str) -> PolarsResult<usize> {
         self.find_idx_by_name(name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+            .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into()))
     }
 
     fn check_already_present(&self, name: &str) -> PolarsResult<()> {
@@ -1338,7 +1338,7 @@ impl DataFrame {
     /// Get column index of a `Series` by name.
     pub fn try_find_idx_by_name(&self, name: &str) -> PolarsResult<usize> {
         self.find_idx_by_name(name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+            .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into()))
     }
 
     /// Select a single column by name.
@@ -1357,7 +1357,7 @@ impl DataFrame {
     pub fn column(&self, name: &str) -> PolarsResult<&Series> {
         let idx = self
             .find_idx_by_name(name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))?;
+            .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into()))?;
         Ok(self.select_at_idx(idx).unwrap())
     }
 
@@ -1478,7 +1478,7 @@ impl DataFrame {
                 .map(|name| {
                     let idx = *name_to_idx
                         .get(name.as_str())
-                        .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))?;
+                        .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into()))?;
                     Ok(self
                         .select_at_idx(idx)
                         .unwrap()
@@ -1506,7 +1506,7 @@ impl DataFrame {
                 .map(|name| {
                     let idx = *name_to_idx
                         .get(name.as_str())
-                        .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))?;
+                        .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into()))?;
                     Ok(self.select_at_idx(idx).unwrap().clone())
                 })
                 .collect::<PolarsResult<Vec<_>>>()?
@@ -1745,7 +1745,7 @@ impl DataFrame {
     /// ```
     pub fn rename(&mut self, column: &str, name: &str) -> PolarsResult<&mut Self> {
         self.select_mut(column)
-            .ok_or_else(|| PolarsError::NotFound(column.to_string().into()))
+            .ok_or_else(|| PolarsError::ColumnNotFound(column.to_string().into()))
             .map(|s| s.rename(name))?;
 
         let unique_names: AHashSet<&str, ahash::RandomState> =
@@ -2197,7 +2197,7 @@ impl DataFrame {
     {
         let idx = self
             .find_idx_by_name(column)
-            .ok_or_else(|| PolarsError::NotFound(column.to_string().into()))?;
+            .ok_or_else(|| PolarsError::ColumnNotFound(column.to_string().into()))?;
         self.try_apply_at_idx(idx, f)
     }
 
@@ -3332,7 +3332,7 @@ impl DataFrame {
             for col in cols {
                 let _ = schema
                     .get(&col)
-                    .ok_or_else(|| PolarsError::NotFound(col.into()))?;
+                    .ok_or_else(|| PolarsError::ColumnNotFound(col.into()))?;
             }
         }
         DataFrame::new(new_cols)

diff --git a/polars/polars-core/src/schema.rs b/polars/polars-core/src/schema.rs
@@ -122,13 +122,13 @@ impl Schema {
 
     pub fn try_get(&self, name: &str) -> PolarsResult<&DataType> {
         self.get(name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+            .ok_or_else(|| PolarsError::SchemaFieldNotFound(name.to_string().into()))
     }
 
     pub fn try_get_full(&self, name: &str) -> PolarsResult<(usize, &String, &DataType)> {
         self.inner
             .get_full(name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+            .ok_or_else(|| PolarsError::SchemaFieldNotFound(name.to_string().into()))
     }
 
     pub fn remove(&mut self, name: &str) -> Option<DataType> {
@@ -148,7 +148,7 @@ impl Schema {
     pub fn try_get_field(&self, name: &str) -> PolarsResult<Field> {
         self.inner
             .get(name)
-            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+            .ok_or_else(|| PolarsError::SchemaFieldNotFound(name.to_string().into()))
             .map(|dtype| Field::new(name, dtype.clone()))
     }
 

diff --git a/polars/polars-io/src/parquet/predicates.rs b/polars/polars-io/src/parquet/predicates.rs
@@ -108,7 +108,7 @@ pub(super) fn read_this_row_group(
                 // a parquet file may not have statistics of all columns
                 if matches!(should_read, Ok(false)) {
                     return Ok(false);
-                } else if !matches!(should_read, Err(PolarsError::NotFound(_))) {
+                } else if !matches!(should_read, Err(PolarsError::ColumnNotFound(_))) {
                     let _ = should_read?;
                 }
             }

diff --git a/polars/polars-io/src/utils.rs b/polars/polars-io/src/utils.rs
@@ -52,10 +52,9 @@ pub(crate) fn columns_to_projection(
     use ahash::AHashMap;
 
     let err = |column: &str| {
-        let valid_fields: Vec<String> = schema.fields.iter().map(|f| f.name.clone()).collect();
-        PolarsError::NotFound(
-            format!("Unable to get field named \"{column}\". Valid fields: {valid_fields:?}",)
-                .into(),
+        let valid_columns: Vec<String> = schema.fields.iter().map(|f| f.name.clone()).collect();
+        PolarsError::ColumnNotFound(
+            format!("Unable to find \"{column}\". Valid columns: {valid_columns:?}",).into(),
         )
     };
 

diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs
@@ -214,11 +214,15 @@ impl FunctionExpr {
                                 .iter()
                                 .find(|fld| fld.name() == name.as_ref())
                                 .ok_or_else(|| {
-                                    PolarsError::NotFound(name.as_ref().to_string().into())
+                                    PolarsError::StructFieldNotFound(
+                                        name.as_ref().to_string().into(),
+                                    )
                                 })?;
                             Ok(fld.clone())
                         } else {
-                            Err(PolarsError::NotFound(name.as_ref().to_string().into()))
+                            Err(PolarsError::StructFieldNotFound(
+                                name.as_ref().to_string().into(),
+                            ))
                         }
                     }
                 }

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/aexpr/schema.rs b/polars/polars-lazy/polars-plan/src/logical_plan/aexpr/schema.rs
@@ -32,7 +32,7 @@ impl AExpr {
             Column(name) => {
                 let field = schema
                     .get_field(name)
-                    .ok_or_else(|| PolarsError::NotFound(name.to_string().into()));
+                    .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into()));
 
                 match ctxt {
                     Context::Default => field,

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/builder.rs b/polars/polars-lazy/polars-plan/src/logical_plan/builder.rs
@@ -502,7 +502,7 @@ impl LogicalPlanBuilder {
                 let dtype = try_delayed!(
                     current_schema
                         .get(name)
-                        .ok_or_else(|| PolarsError::NotFound(name.to_string().into())),
+                        .ok_or_else(|| PolarsError::ColumnNotFound(name.to_string().into())),
                     self.0,
                     into
                 );

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/functions/mod.rs b/polars/polars-lazy/polars-plan/src/logical_plan/functions/mod.rs
@@ -102,9 +102,9 @@ impl FunctionNode {
                     .iter()
                     .map(|name| {
                         let name = name.as_ref();
-                        input_schema
-                            .get_field(name)
-                            .ok_or_else(|| PolarsError::NotFound(name.to_string().into()))
+                        input_schema.get_field(name).ok_or_else(|| {
+                            PolarsError::SchemaFieldNotFound(name.to_string().into())
+                        })
                     })
                     .collect::<PolarsResult<Schema>>()?;
                 Ok(Cow::Owned(Arc::new(schema)))

diff --git a/polars/polars-lazy/src/physical_plan/expressions/column.rs b/polars/polars-lazy/src/physical_plan/expressions/column.rs
@@ -169,9 +169,9 @@ impl PhysicalExpr for ColumnExpr {
 
     fn to_field(&self, input_schema: &Schema) -> PolarsResult<Field> {
         let field = input_schema.get_field(&self.name).ok_or_else(|| {
-            PolarsError::NotFound(
+            PolarsError::ColumnNotFound(
                 format!(
-                    "could not find column: {} in schema: {:?}",
+                    "could not find: {} in schema: {:?}",
                     self.name, &input_schema
                 )
                 .into(),

diff --git a/polars/tests/it/lazy/projection_queries.rs b/polars/tests/it/lazy/projection_queries.rs
@@ -89,7 +89,7 @@ fn test_err_no_found() {
 
     assert!(matches!(
         df.lazy().filter(col("nope").gt(lit(2))).collect(),
-        Err(PolarsError::NotFound(_))
+        Err(PolarsError::ColumnNotFound(_))
     ));
 }
 

diff --git a/py-polars/docs/source/reference/exceptions.rst b/py-polars/docs/source/reference/exceptions.rst
@@ -8,14 +8,16 @@ Exceptions
     :nosignatures:
 
     ArrowError
+    ColumnNotFoundError
     ComputeError
     DuplicateError
     InvalidOperationError
     NoDataError
-    NotFoundError
     NoRowsReturned
     PanicException
     RowsException
     SchemaError
+    SchemaFieldNotFoundError
     ShapeError
+    StructFieldNotFoundError
     TooManyRowsReturned