From a88c714be70f111e454589c25c6933c103fef161 Mon Sep 17 00:00:00 2001
From: Matt Green <emgeee@users.noreply.github.com>
Date: Wed, 6 Nov 2024 11:24:29 -0800
Subject: [PATCH] add drop_columns

---
 crates/core/src/datastream.rs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/crates/core/src/datastream.rs b/crates/core/src/datastream.rs
index 8b83ec4..42a8f4f 100644
--- a/crates/core/src/datastream.rs
+++ b/crates/core/src/datastream.rs
@@ -113,6 +113,15 @@ impl DataStream {
         })
     }
 
+    pub fn drop_columns(self, columns: &[&str]) -> Result<Self> {
+        Ok(Self {
+            df: Arc::new(self.df.as_ref().clone().drop_columns(columns)?),
+            context: self.context.clone(),
+            shutdown_tx: self.shutdown_tx.clone(),
+            shutdown_rx: self.shutdown_rx.clone(),
+        })
+    }
+
     // Join two streams using the specified expression
     pub fn join_on(
         self,