ignore empty lists when dropping

quixio · Jul 16, 2024 · 8d42733 · 8d42733
1 parent eb5e1ac
commit 8d42733
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/quixstreams/dataframe/dataframe.py b/quixstreams/dataframe/dataframe.py
@@ -933,6 +933,8 @@ def drop(self, columns: Union[str, List[str]]) -> Self:
         :return: a new StreamingDataFrame instance
         """
         if isinstance(columns, list):
+            if not columns:
+                return self
             if not all(isinstance(s, str) for s in columns):
                 raise TypeError(f"column list must contain strings only")
         elif isinstance(columns, str):

diff --git a/tests/test_quixstreams/test_dataframe/test_dataframe.py b/tests/test_quixstreams/test_dataframe/test_dataframe.py
@@ -392,6 +392,16 @@ def test_drop_invalid_columns(self, dataframe_factory, columns):
         with pytest.raises(TypeError):
             sdf.drop(columns)
 
+    def test_drop_empty_list(self, dataframe_factory):
+        """
+        Dropping an empty list is ignored entirely.
+        """
+        sdf = dataframe_factory()
+        pre_drop_stream = sdf.stream.tree()
+        sdf = sdf.drop([])
+        post_drop_stream = sdf.stream.tree()
+        assert pre_drop_stream == post_drop_stream
+
 
 class TestStreamingDataFrameApplyExpand:
     def test_apply_expand(self, dataframe_factory):