Allow to ignore missing columns when calling "SDF.drop()" (#434)

Author: daniil-quix

docs/processing.md

@@ -150,6 +150,9 @@ It accepts either one column name as a string or a list of names.
 
 The `.drop()` method updates the existing `StreamingDataFrame` object and returns the same `StreamingDataFrame` instance so that you can chain other methods after the `drop()` call, too.
 
+By default, the `.drop()` fails if one of the columns is missing in the value.  
+If you expect some columns to be missing, you can pass `errors="ignore"` to the `.drop()` method to suppress the errors. 
+
 Internally, it mutates the record's value and deletes the keys in place.
 
 **Example**:
@@ -184,6 +187,9 @@ sdf.drop("metadata")
 
 # You may also drop multiple keys by providing a list of names:
 sdf.drop(["metadata", "timestamp"])
+
+# You may suppress KeyErrors if some columns are missing in the value dictionary
+sdf.drop(["missing_column"], errors='ignore')
 ```
 
 > **_NOTE:_**  The `StreamingDataFrame.drop()` method works only with mapping-like values like dictionaries.

quixstreams/dataframe/dataframe.py

@@ -966,7 +966,11 @@ def hopping_window(
             name=name,
         )
 
-    def drop(self, columns: Union[str, List[str]]) -> Self:
+    def drop(
+        self,
+        columns: Union[str, List[str]],
+        errors: Literal["ignore", "raise"] = "raise",
+    ) -> Self:
         """
         Drop column(s) from the message value (value must support `del`, like a dict).
 
@@ -985,6 +989,9 @@ def drop(self, columns: Union[str, List[str]]) -> Self:
         ```
 
         :param columns: a single column name or a list of names, where names are `str`
+        :param errors: If "ignore", suppress error and only existing labels are dropped.
+            Default - `"raise"`.
+
         :return: a new StreamingDataFrame instance
         """
         if isinstance(columns, list):
@@ -998,7 +1005,10 @@ def drop(self, columns: Union[str, List[str]]) -> Self:
             raise TypeError(
                 f"Expected a string or a list of strings, not {type(columns)}"
             )
-        return self._add_update(lambda value: _drop(value, columns), metadata=False)
+        return self._add_update(
+            lambda value: _drop(value, columns, ignore_missing=errors == "ignore"),
+            metadata=False,
+        )
 
     def _produce(
         self,
@@ -1145,14 +1155,19 @@ def __bool__(self):
         )
 
 
-def _drop(value: Dict, columns: List[str]):
+def _drop(value: Dict, columns: List[str], ignore_missing: bool = False):
     """
     remove columns from the value, inplace
     :param value: a dict or something that supports `del`
     :param columns: a list of column names
+    :param ignore_missing: if True, ignore missing columns
     """
     for column in columns:
-        del value[column]
+        try:
+            del value[column]
+        except KeyError:
+            if not ignore_missing:
+                raise
 
 
 def _as_metadata_func(

tests/test_quixstreams/test_dataframe/test_dataframe.py

@@ -1,4 +1,3 @@
-import json
 import operator
 import uuid
 from collections import namedtuple
@@ -423,6 +422,24 @@ def test_drop_empty_list(self, dataframe_factory):
         post_drop_stream = sdf.stream.tree()
         assert pre_drop_stream == post_drop_stream
 
+    def test_drop_missing_columns_errors_raise(self, dataframe_factory):
+        value = {"col_a": 1}
+        key, timestamp, headers = b"key", 0, []
+        sdf = dataframe_factory()
+        sdf.drop(["col_b", "col_c"], errors="raise")
+        with pytest.raises(KeyError):
+            assert sdf.test(value=value, key=key, timestamp=timestamp, headers=headers)
+
+    def test_drop_missing_columns_errors_ignore(self, dataframe_factory):
+        value = {"col_a": 1}
+        expected = {"col_a": 1}
+        key, timestamp, headers = b"key", 0, []
+        sdf = dataframe_factory()
+        sdf.drop(["col_b", "col_c"], errors="ignore")
+        assert sdf.test(value=value, key=key, timestamp=timestamp, headers=headers)[
+            0
+        ] == (expected, key, timestamp, headers)
+
 
 class TestStreamingDataFrameApplyExpand:
     def test_apply_expand(self, dataframe_factory):