refactor: remove 'fast-projection' node (#15253)

pola-rs · Mar 23, 2024 · 07538dd · 07538dd
1 parent 6793954
commit 07538dd
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 99 deletions.
diff --git a/crates/polars-pipe/src/pipeline/convert.rs b/crates/polars-pipe/src/pipeline/convert.rs
@@ -643,17 +643,6 @@ where
             let op = operators::FilterOperator { predicate };
             Box::new(op) as Box<dyn Operator>
         },
-        MapFunction {
-            function: FunctionNode::FastProjection { columns, .. },
-            input,
-        } => {
-            let input_schema = lp_arena.get(*input).schema(lp_arena);
-            let op = operators::SimpleProjectionOperator::new(
-                columns.clone(),
-                input_schema.into_owned(),
-            );
-            Box::new(op) as Box<dyn Operator>
-        },
         MapFunction { function, .. } => {
             let op = operators::FunctionOperator::new(function.clone());
             Box::new(op) as Box<dyn Operator>

diff --git a/crates/polars-plan/src/logical_plan/functions/mod.rs b/crates/polars-plan/src/logical_plan/functions/mod.rs
@@ -61,10 +61,6 @@ pub enum FunctionNode {
     Unnest {
         columns: Arc<[Arc<str>]>,
     },
-    FastProjection {
-        columns: Arc<[SmartString]>,
-        duplicate_check: bool,
-    },
     DropNulls {
         subset: Arc<[Arc<str>]>,
     },
@@ -103,16 +99,6 @@ impl PartialEq for FunctionNode {
     fn eq(&self, other: &Self) -> bool {
         use FunctionNode::*;
         match (self, other) {
-            (
-                FastProjection {
-                    columns: l,
-                    duplicate_check: dl,
-                },
-                FastProjection {
-                    columns: r,
-                    duplicate_check: dr,
-                },
-            ) => l == r && dl == dr,
             (DropNulls { subset: l }, DropNulls { subset: r }) => l == r,
             (Rechunk, Rechunk) => true,
             (Count { paths: paths_l, .. }, Count { paths: paths_r, .. }) => paths_l == paths_r,
@@ -144,12 +130,9 @@ impl FunctionNode {
             Rechunk | Pipeline { .. } => false,
             #[cfg(feature = "merge_sorted")]
             MergeSorted { .. } => false,
-            DropNulls { .. }
-            | FastProjection { .. }
-            | Count { .. }
-            | Unnest { .. }
-            | Rename { .. }
-            | Explode { .. } => true,
+            DropNulls { .. } | Count { .. } | Unnest { .. } | Rename { .. } | Explode { .. } => {
+                true
+            },
             Melt { args, .. } => args.streamable,
             Opaque { streamable, .. } => *streamable,
             #[cfg(feature = "python")]
@@ -188,16 +171,6 @@ impl FunctionNode {
                 .map(|schema| Cow::Owned(schema.clone()))
                 .unwrap_or_else(|| Cow::Borrowed(input_schema))),
             Pipeline { schema, .. } => Ok(Cow::Owned(schema.clone())),
-            FastProjection { columns, .. } => {
-                let schema = columns
-                    .iter()
-                    .map(|name| {
-                        let name = name.as_ref();
-                        input_schema.try_get_field(name)
-                    })
-                    .collect::<PolarsResult<Schema>>()?;
-                Ok(Cow::Owned(Arc::new(schema)))
-            },
             DropNulls { .. } => Ok(Cow::Borrowed(input_schema)),
             Count { alias, .. } => {
                 let mut schema: Schema = Schema::with_capacity(1);
@@ -262,8 +235,7 @@ impl FunctionNode {
             Opaque { predicate_pd, .. } => *predicate_pd,
             #[cfg(feature = "python")]
             OpaquePython { predicate_pd, .. } => *predicate_pd,
-            FastProjection { .. }
-            | DropNulls { .. }
+            DropNulls { .. }
             | Rechunk
             | Unnest { .. }
             | Rename { .. }
@@ -282,8 +254,7 @@ impl FunctionNode {
             Opaque { projection_pd, .. } => *projection_pd,
             #[cfg(feature = "python")]
             OpaquePython { projection_pd, .. } => *projection_pd,
-            FastProjection { .. }
-            | DropNulls { .. }
+            DropNulls { .. }
             | Rechunk
             | Count { .. }
             | Unnest { .. }
@@ -319,16 +290,6 @@ impl FunctionNode {
                 schema,
                 ..
             } => python_udf::call_python_udf(function, df, *validate_output, schema.as_deref()),
-            FastProjection {
-                columns,
-                duplicate_check,
-            } => {
-                if *duplicate_check {
-                    df._select_impl(columns.as_ref())
-                } else {
-                    df._select_impl_unchecked(columns.as_ref())
-                }
-            },
             DropNulls { subset } => df.drop_nulls(Some(subset.as_ref())),
             Count {
                 paths, scan_type, ..
@@ -386,11 +347,6 @@ impl Display for FunctionNode {
             Opaque { fmt_str, .. } => write!(f, "{fmt_str}"),
             #[cfg(feature = "python")]
             OpaquePython { .. } => write!(f, "python dataframe udf"),
-            FastProjection { columns, .. } => {
-                write!(f, "FAST_PROJECT: ")?;
-                let columns = columns.as_ref();
-                fmt_column_delimited(f, columns, "[", "]")
-            },
             DropNulls { subset } => {
                 write!(f, "DROP_NULLS by: ")?;
                 let subset = subset.as_ref();

diff --git a/crates/polars-plan/src/logical_plan/optimizer/mod.rs b/crates/polars-plan/src/logical_plan/optimizer/mod.rs
@@ -12,7 +12,6 @@ mod collect_members;
 mod count_star;
 #[cfg(feature = "cse")]
 mod cse_expr;
-mod fast_projection;
 #[cfg(any(
     feature = "ipc",
     feature = "parquet",
@@ -26,6 +25,7 @@ mod flatten_union;
 mod fused;
 mod predicate_pushdown;
 mod projection_pushdown;
+mod simple_projection;
 mod simplify_expr;
 mod simplify_functions;
 mod slice_pushdown_expr;
@@ -35,10 +35,10 @@ mod type_coercion;
 
 use delay_rechunk::DelayRechunk;
 use drop_nulls::ReplaceDropNulls;
-use fast_projection::FastProjectionAndCollapse;
 use polars_io::predicates::PhysicalIoExpr;
 pub use predicate_pushdown::PredicatePushDown;
 pub use projection_pushdown::ProjectionPushDown;
+use simple_projection::SimpleProjectionAndCollapse;
 pub use simplify_expr::{SimplifyBooleanRule, SimplifyExprRule};
 use slice_pushdown_lp::SlicePushDown;
 pub use stack_opt::{OptimizationRule, StackOptimizer};
@@ -156,7 +156,7 @@ pub fn optimize(
 
     // make sure its before slice pushdown.
     if fast_projection {
-        rules.push(Box::new(FastProjectionAndCollapse::new(eager)));
+        rules.push(Box::new(SimpleProjectionAndCollapse::new(eager)));
     }
 
     if !eager {

diff --git a/...logical_plan/optimizer/fast_projection.rs → ...gical_plan/optimizer/simple_projection.rs b/...logical_plan/optimizer/fast_projection.rs → ...gical_plan/optimizer/simple_projection.rs
@@ -11,14 +11,14 @@ use super::*;
 /// It is important that this optimization is ran after projection pushdown.
 ///
 /// The schema reported after this optimization is also
-pub(super) struct FastProjectionAndCollapse {
+pub(super) struct SimpleProjectionAndCollapse {
     /// keep track of nodes that are already processed when they
     /// can be expensive. Schema materialization can be for instance.
     processed: BTreeSet<Node>,
     eager: bool,
 }
 
-impl FastProjectionAndCollapse {
+impl SimpleProjectionAndCollapse {
     pub(super) fn new(eager: bool) -> Self {
         Self {
             processed: Default::default(),
@@ -27,7 +27,7 @@ impl FastProjectionAndCollapse {
     }
 }
 
-impl OptimizationRule for FastProjectionAndCollapse {
+impl OptimizationRule for SimpleProjectionAndCollapse {
     fn optimize_plan(
         &mut self,
         lp_arena: &mut Arena<ALogicalPlan>,
@@ -93,39 +93,6 @@ impl OptimizationRule for FastProjectionAndCollapse {
                     _ => None,
                 }
             },
-            MapFunction {
-                input,
-                function: FunctionNode::FastProjection { columns, .. },
-            } if !self.eager => {
-                // if there are 2 subsequent fast projections, flatten them and only take the last
-                match lp_arena.get(*input) {
-                    MapFunction {
-                        function: FunctionNode::FastProjection { .. },
-                        input: prev_input,
-                    } => Some(MapFunction {
-                        input: *prev_input,
-                        function: FunctionNode::FastProjection {
-                            columns: columns.clone(),
-                            duplicate_check: true,
-                        },
-                    }),
-                    // cleanup projections set in projection pushdown just above caches
-                    // they are not needed.
-                    cache_lp @ Cache { .. } if self.processed.insert(node) => {
-                        let cache_schema = cache_lp.schema(lp_arena);
-                        if cache_schema.len() == columns.len()
-                            && cache_schema.iter_names().zip(columns.iter()).all(
-                                |(left_name, right_name)| left_name.as_str() == right_name.as_str(),
-                            )
-                        {
-                            Some(cache_lp.clone())
-                        } else {
-                            None
-                        }
-                    },
-                    _ => None,
-                }
-            },
             // if there are 2 subsequent caches, flatten them and only take the inner
             Cache {
                 input,