update schema when projections are pushed down to the dataframe level (…

…#1660)

polars/polars-lazy/src/logical_plan/optimizer/projection_pushdown.rs

@@ -86,6 +86,21 @@ fn add_str_to_accumulated(
     }
 }
 
+fn update_scan_schema(
+    acc_projections: &[Node],
+    expr_arena: &Arena<AExpr>,
+    schema: &Schema,
+) -> Result<Schema> {
+    let mut new_fields = Vec::with_capacity(acc_projections.len());
+    for node in acc_projections.iter() {
+        for name in aexpr_to_root_names(*node, expr_arena) {
+            let field = schema.field_with_name(&*name)?;
+            new_fields.push(field.clone())
+        }
+    }
+    Ok(Schema::new(new_fields))
+}
+
 pub(crate) struct ProjectionPushDown {}
 
 impl ProjectionPushDown {
@@ -278,13 +293,14 @@ impl ProjectionPushDown {
             }
             DataFrameScan {
                 df,
-                schema,
+                mut schema,
                 selection,
                 ..
             } => {
                 let mut projection = None;
                 if !acc_projections.is_empty() {
-                    projection = Some(acc_projections)
+                    schema = Arc::new(update_scan_schema(&acc_projections, expr_arena, &*schema)?);
+                    projection = Some(acc_projections);
                 }
                 let lp = DataFrameScan {
                     df,
@@ -326,6 +342,7 @@ impl ProjectionPushDown {
                 ..
             } => {
                 options.with_columns = get_scan_columns(&mut acc_projections, expr_arena);
+
                 let lp = CsvScan {
                     path,
                     schema,

polars/polars-lazy/src/test.rs

@@ -2024,3 +2024,18 @@ fn test_binary_expr() -> Result<()> {
     assert_eq!(out.dtypes(), &[DataType::Float64]);
     Ok(())
 }
+
+#[test]
+fn test_drop_and_select() -> Result<()> {
+    let df = fruits_cars();
+
+    let out = df
+        .lazy()
+        .drop_columns(["A", "B"])
+        .select([col("fruits")])
+        .collect()?;
+
+    assert_eq!(out.get_column_names(), &["fruits"]);
+
+    Ok(())
+}