feat: implement drop as special case of select (#10885)

pola-rs · Sep 3, 2023 · f6fe1b9 · f6fe1b9
1 parent c67d056
commit f6fe1b9
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 64 deletions.
diff --git a/crates/polars-lazy/src/frame/mod.rs b/crates/polars-lazy/src/frame/mod.rs
@@ -431,22 +431,7 @@ impl LazyFrame {
         I: IntoIterator<Item = T>,
         T: AsRef<str>,
     {
-        let columns: Vec<SmartString> = columns
-            .into_iter()
-            .map(|name| name.as_ref().into())
-            .collect();
-        self.drop_columns_impl(columns)
-    }
-
-    #[allow(clippy::ptr_arg)]
-    fn drop_columns_impl(self, columns: Vec<SmartString>) -> Self {
-        if let Some(lp) = self.check_names(&columns, None) {
-            lp
-        } else {
-            self.map_private(FunctionNode::Drop {
-                names: columns.into(),
-            })
-        }
+        self.select(vec![col("*").exclude(columns)])
     }
 
     /// Shift the values by a given period and fill the parts that will be empty due to this operation

diff --git a/crates/polars-lazy/src/physical_plan/executors/projection_utils.rs b/crates/polars-lazy/src/physical_plan/executors/projection_utils.rs
@@ -185,7 +185,9 @@ pub(super) fn check_expand_literals(
     mut selected_columns: Vec<Series>,
     zero_length: bool,
 ) -> PolarsResult<DataFrame> {
-    let first_len = selected_columns[0].len();
+    let Some(first_len) = selected_columns.get(0).map(|s| s.len()) else {
+        return Ok(DataFrame::empty());
+    };
     let mut df_height = 0;
     let mut all_equal_len = true;
     {

diff --git a/crates/polars-plan/src/logical_plan/functions/drop.rs b/crates/polars-plan/src/logical_plan/functions/drop.rs
diff --git a/crates/polars-plan/src/logical_plan/functions/mod.rs b/crates/polars-plan/src/logical_plan/functions/mod.rs
@@ -1,4 +1,3 @@
-mod drop;
 #[cfg(feature = "merge_sorted")]
 mod merge_sorted;
 #[cfg(feature = "python")]
@@ -80,9 +79,6 @@ pub enum FunctionNode {
         // A column name gets swapped with an existing column
         swapping: bool,
     },
-    Drop {
-        names: Arc<[SmartString]>,
-    },
     Explode {
         columns: Arc<[Arc<str>]>,
         schema: SchemaRef,
@@ -117,7 +113,6 @@ impl PartialEq for FunctionNode {
                     ..
                 },
             ) => existing_l == existing_r && new_l == new_r,
-            (Drop { names: l }, Drop { names: r }) => l == r,
             (Explode { columns: l, .. }, Explode { columns: r, .. }) => l == r,
             (Melt { args: l, .. }, Melt { args: r, .. }) => l == r,
             (RowCount { name: l, .. }, RowCount { name: r, .. }) => l == r,
@@ -138,8 +133,7 @@ impl FunctionNode {
             | FastProjection { .. }
             | Unnest { .. }
             | Rename { .. }
-            | Explode { .. }
-            | Drop { .. } => true,
+            | Explode { .. } => true,
             Melt { args, .. } => args.streamable,
             Opaque { streamable, .. } => *streamable,
             #[cfg(feature = "python")]
@@ -229,7 +223,6 @@ impl FunctionNode {
             #[cfg(feature = "merge_sorted")]
             MergeSorted { .. } => Ok(Cow::Borrowed(input_schema)),
             Rename { existing, new, .. } => rename::rename_schema(input_schema, existing, new),
-            Drop { names } => drop::drop_schema(input_schema, names),
             Explode { schema, .. } | RowCount { schema, .. } | Melt { schema, .. } => {
                 Ok(Cow::Owned(schema.clone()))
             },
@@ -248,8 +241,7 @@ impl FunctionNode {
             | Unnest { .. }
             | Rename { .. }
             | Explode { .. }
-            | Melt { .. }
-            | Drop { .. } => true,
+            | Melt { .. } => true,
             #[cfg(feature = "merge_sorted")]
             MergeSorted { .. } => true,
             RowCount { .. } => false,
@@ -269,8 +261,7 @@ impl FunctionNode {
             | Unnest { .. }
             | Rename { .. }
             | Explode { .. }
-            | Melt { .. }
-            | Drop { .. } => true,
+            | Melt { .. } => true,
             #[cfg(feature = "merge_sorted")]
             MergeSorted { .. } => true,
             RowCount { .. } => true,
@@ -332,7 +323,6 @@ impl FunctionNode {
                 }
             },
             Rename { existing, new, .. } => rename::rename_impl(df, existing, new),
-            Drop { names } => drop::drop_impl(df, names),
             Explode { columns, .. } => df.explode(columns.as_ref()),
             Melt { args, .. } => {
                 let args = (**args).clone();
@@ -385,7 +375,6 @@ impl Display for FunctionNode {
                 }
             },
             Rename { .. } => write!(f, "RENAME"),
-            Drop { .. } => write!(f, "DROP"),
             Explode { .. } => write!(f, "EXPLODE"),
             Melt { .. } => write!(f, "MELT"),
             RowCount { .. } => write!(f, "WITH ROW COUNT"),