Skip to content

Commit

Permalink
change drop to inplace after rebase
Browse files Browse the repository at this point in the history
  • Loading branch information
tim-quix committed Jul 17, 2024
1 parent 1823ece commit b89c82e
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 4 deletions.
9 changes: 6 additions & 3 deletions quixstreams/dataframe/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -970,14 +970,18 @@ def drop(self, columns: Union[str, List[str]]) -> Self:
"""
Drop column(s) from the message value (value must support `del`, like a dict).
This operation occurs in-place, meaning reassignment is entirely OPTIONAL: the
original `StreamingDataFrame` is returned for chaining (`sdf.update().print()`).
Example Snippet:
```python
# Remove columns "x" and "y" from the value.
# This would transform {"x": 1, "y": 2, "z": 3} to {"z": 3}
sdf = StreamingDataframe()
sdf = sdf.drop(["x", "y"])
sdf.drop(["x", "y"])
```
:param columns: a single column name or a list of names, where names are `str`
Expand All @@ -994,8 +998,7 @@ def drop(self, columns: Union[str, List[str]]) -> Self:
raise TypeError(
f"Expected a string or a list of strings, not {type(columns)}"
)
stream = self.stream.add_update(lambda value: _drop(value, columns))
return self.__dataframe_clone__(stream)
return self._add_update(lambda value: _drop(value, columns), metadata=False)

def _produce(
self,
Expand Down
2 changes: 1 addition & 1 deletion tests/test_quixstreams/test_dataframe/test_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,7 +402,7 @@ def test_drop(self, dataframe_factory, columns, expected):
value = {"col_a": 1, "col_b": 2, "col_c": 3}
key, timestamp, headers = b"key", 0, []
sdf = dataframe_factory()
sdf = sdf.drop(columns)
sdf.drop(columns)
assert sdf.test(value=value, key=key, timestamp=timestamp, headers=headers)[
0
] == (expected, key, timestamp, headers)
Expand Down

0 comments on commit b89c82e

Please sign in to comment.