Initial fix for rename perf issue

pola-rs · Sep 24, 2024 · 8128ce8 · 8128ce8
1 parent 54218e7
commit 8128ce8
Showing 1 changed file with 15 additions and 4 deletions.
diff --git a/crates/polars-plan/src/plans/functions/rename.rs b/crates/polars-plan/src/plans/functions/rename.rs
@@ -5,10 +5,21 @@ pub(super) fn rename_impl(
     existing: &[PlSmallStr],
     new: &[PlSmallStr],
 ) -> PolarsResult<DataFrame> {
-    let positions = existing
-        .iter()
-        .map(|old| df.get_column_index(old))
-        .collect::<Vec<_>>();
+    let positions = if existing.len() > 1 && df.get_columns().len() > 10 {
+        let schema = df.schema();
+        existing
+            .iter()
+            .map(|old| match schema.get_full(old) {
+                Some((idx, _, _)) => Some(idx),
+                None => None,
+            })
+            .collect::<Vec<_>>()
+    } else {
+        existing
+            .iter()
+            .map(|old| df.get_column_index(old))
+            .collect::<Vec<_>>()
+    };
 
     for (pos, name) in positions.iter().zip(new.iter()) {
         // the column might be removed due to projection pushdown