improve error on missing column access (#4095)

pola-rs · Jul 20, 2022 · 547229d · 547229d
1 parent 453af3c
commit 547229d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 18 deletions.
diff --git a/polars/polars-lazy/src/frame/mod.rs b/polars/polars-lazy/src/frame/mod.rs
@@ -573,18 +573,14 @@ impl LazyFrame {
         if projection_pushdown {
             let projection_pushdown_opt = ProjectionPushDown {};
             let alp = lp_arena.take(lp_top);
-            let alp = projection_pushdown_opt
-                .optimize(alp, lp_arena, expr_arena)
-                .expect("projection pushdown failed");
+            let alp = projection_pushdown_opt.optimize(alp, lp_arena, expr_arena)?;
             lp_arena.replace(lp_top, alp);
         }
 
         if predicate_pushdown {
             let predicate_pushdown_opt = PredicatePushDown::default();
             let alp = lp_arena.take(lp_top);
-            let alp = predicate_pushdown_opt
-                .optimize(alp, lp_arena, expr_arena)
-                .expect("predicate pushdown failed");
+            let alp = predicate_pushdown_opt.optimize(alp, lp_arena, expr_arena)?;
             lp_arena.replace(lp_top, alp);
         }
 
@@ -625,9 +621,7 @@ impl LazyFrame {
         if slice_pushdown {
             let slice_pushdown_opt = SlicePushDown {};
             let alp = lp_arena.take(lp_top);
-            let alp = slice_pushdown_opt
-                .optimize(alp, lp_arena, expr_arena)
-                .expect("slice pushdown failed");
+            let alp = slice_pushdown_opt.optimize(alp, lp_arena, expr_arena)?;
 
             lp_arena.replace(lp_top, alp);
 

diff --git a/polars/polars-lazy/src/logical_plan/optimizer/projection_pushdown.rs b/polars/polars-lazy/src/logical_plan/optimizer/projection_pushdown.rs
@@ -98,12 +98,16 @@ fn update_scan_schema(
     // sorting parsers: csv,
     // non-sorting: parquet, ipc
     sort_projections: bool,
-) -> Schema {
+) -> Result<Schema> {
     let mut new_schema = Schema::with_capacity(acc_projections.len());
     let mut new_cols = Vec::with_capacity(acc_projections.len());
     for node in acc_projections.iter() {
         for name in aexpr_to_root_names(*node, expr_arena) {
-            let item = schema.get_full(&*name).unwrap();
+            let item = schema.get_full(&*name).ok_or_else(|| {
+                PolarsError::ComputeError(
+                    format!("column {} not available in schema {:?}", name, schema).into(),
+                )
+            })?;
             new_cols.push(item);
         }
     }
@@ -114,7 +118,7 @@ fn update_scan_schema(
     for item in new_cols {
         new_schema.with_column(item.1.clone(), item.2.clone())
     }
-    new_schema
+    Ok(new_schema)
 }
 
 pub(crate) struct ProjectionPushDown {}
@@ -356,7 +360,7 @@ impl ProjectionPushDown {
                             expr_arena,
                             &*schema,
                             true,
-                        )))
+                        )?))
                     };
                     options.output_schema = output_schema.clone();
 
@@ -394,7 +398,7 @@ impl ProjectionPushDown {
                         expr_arena,
                         &*schema,
                         false,
-                    ));
+                    )?);
                     projection = Some(acc_projections);
                 }
                 let lp = DataFrameScan {
@@ -423,7 +427,7 @@ impl ProjectionPushDown {
                         expr_arena,
                         &*schema,
                         false,
-                    )))
+                    )?))
                 };
                 options.with_columns = with_columns;
 
@@ -456,7 +460,7 @@ impl ProjectionPushDown {
                         expr_arena,
                         &*schema,
                         false,
-                    )))
+                    )?))
                 };
                 options.with_columns = with_columns;
 
@@ -482,7 +486,7 @@ impl ProjectionPushDown {
                         expr_arena,
                         &*options.schema,
                         true,
-                    )))
+                    )?))
                 };
                 Ok(PythonScan { options })
             }
@@ -505,7 +509,7 @@ impl ProjectionPushDown {
                         expr_arena,
                         &*schema,
                         true,
-                    )))
+                    )?))
                 };
 
                 let lp = CsvScan {

diff --git a/py-polars/tests/test_errors.py b/py-polars/tests/test_errors.py
@@ -83,3 +83,17 @@ def test_join_lazy_on_df() -> None:
         match="Expected a `LazyFrame` as join table, got <class 'polars.internals.frame.DataFrame'>",
     ):
         df_left.lazy().join_asof(df_right, on="Id")
+
+
+def test_projection_update_schema_missing_column() -> None:
+    with pytest.raises(
+        pl.ComputeError, match="column colC not available in schema Schema:*"
+    ):
+        (
+            pl.DataFrame({"colA": ["a", "b", "c"], "colB": [1, 2, 3]})
+            .lazy()
+            .filter(~pl.col("colC").is_null())
+            .groupby(["colA"])
+            .agg([pl.col("colB").sum().alias("result")])
+            .collect()
+        )