From f6fe1b9f45b20a8eacfbfa07d01fbaa68977abe5 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Sun, 3 Sep 2023 19:58:49 +0200 Subject: [PATCH] feat: implement drop as special case of `select` (#10885) --- crates/polars-lazy/src/frame/mod.rs | 17 +--------- .../executors/projection_utils.rs | 4 ++- .../src/logical_plan/functions/drop.rs | 33 ------------------- .../src/logical_plan/functions/mod.rs | 17 ++-------- 4 files changed, 7 insertions(+), 64 deletions(-) delete mode 100644 crates/polars-plan/src/logical_plan/functions/drop.rs diff --git a/crates/polars-lazy/src/frame/mod.rs b/crates/polars-lazy/src/frame/mod.rs index dea0a304668b..a0ef4e18099e 100644 --- a/crates/polars-lazy/src/frame/mod.rs +++ b/crates/polars-lazy/src/frame/mod.rs @@ -431,22 +431,7 @@ impl LazyFrame { I: IntoIterator, T: AsRef, { - let columns: Vec = 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) -> 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 index 70e33fb986a0..40532904c0ac 100644 --- 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, zero_length: bool, ) -> PolarsResult { - 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 deleted file mode 100644 index 242289329ead..000000000000 --- a/crates/polars-plan/src/logical_plan/functions/drop.rs +++ /dev/null @@ -1,33 +0,0 @@ -use super::*; - -pub(super) fn drop_impl(mut df: DataFrame, names: &[SmartString]) -> PolarsResult { - for name in names { - // ignore names that are not in there - // they might already be removed by projection pushdown - if let Some(idx) = df.find_idx_by_name(name) { - let _ = unsafe { df.get_columns_mut().remove(idx) }; - } - } - - Ok(df) -} - -pub(super) fn drop_schema<'a>( - input_schema: &'a SchemaRef, - names: &[SmartString], -) -> PolarsResult> { - let to_drop = PlHashSet::from_iter(names); - - let new_schema = input_schema - .iter() - .flat_map(|(name, dtype)| { - if to_drop.contains(name) { - None - } else { - Some(Field::new(name, dtype.clone())) - } - }) - .collect::(); - - Ok(Cow::Owned(Arc::new(new_schema))) -} diff --git a/crates/polars-plan/src/logical_plan/functions/mod.rs b/crates/polars-plan/src/logical_plan/functions/mod.rs index 2be3ff57a1e0..34582e4e5a7b 100644 --- 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]>, 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"),