Add partition_unique node to core

Author: bdewilde

streamz/core.py

@@ -8,6 +8,7 @@
 import sys
 import threading
 from time import time
+from typing import Any, Callable, Hashable, Union
 import weakref
 
 import toolz
@@ -1070,6 +1071,107 @@ def update(self, x, who=None, metadata=None):
             )
 
 
+@Stream.register_api()
+class partition_unique(Stream):
+    """
+    Partition stream elements into groups of equal size with unique keys only.
+
+    Parameters
+    ----------
+    n: int
+        Number of (unique) elements to pass through as a group.
+    key: Union[Hashable, Callable[[Any], Hashable]]
+        Callable that accepts a stream element and returns a unique, hashable
+        representation of the incoming data (``key(x)``), or a hashable that gets
+        the corresponding value of a stream element (``x[key]``). For example,
+        ``key=lambda x: x["a"]`` would allow only elements with unique ``"a"`` values
+        to pass through.
+
+        .. note:: By default, we simply use the element object itself as the key,
+            so that object must be hashable. If that's not the case, a non-default
+            key must be provided.
+
+    keep: str
+        Which element to keep in the case that a unique key is already found
+        in the group. If "first", keep element from the first occurrence of a given
+        key; if "last", keep element from the most recent occurrence. Note that
+        relative ordering of *elements* is preserved in the data passed through,
+        and not ordering of *keys*.
+    **kwargs
+
+    Examples
+    --------
+    >>> source = Stream()
+    >>> stream = source.partition_unique(n=3, keep="first").sink(print)
+    >>> eles = [1, 2, 1, 3, 1, 3, 3, 2]
+    >>> for ele in eles:
+    ...     source.emit(ele)
+    (1, 2, 3)
+    (1, 3, 2)
+
+    >>> source = Stream()
+    >>> stream = source.partition_unique(n=3, keep="last").sink(print)
+    >>> eles = [1, 2, 1, 3, 1, 3, 3, 2]
+    >>> for ele in eles:
+    ...     source.emit(ele)
+    (2, 1, 3)
+    (1, 3, 2)
+
+    >>> source = Stream()
+    >>> stream = source.partition_unique(n=3, key=lambda x: len(x), keep="last").sink(print)
+    >>> eles = ["f", "fo", "f", "foo", "f", "foo", "foo", "fo"]
+    >>> for ele in eles:
+    ...     source.emit(ele)
+    ('fo', 'f', 'foo')
+    ('f', 'foo', 'fo')
+    """
+    _graphviz_shape = "diamond"
+
+    def __init__(
+        self,
+        upstream,
+        n: int,
+        key: Union[Hashable, Callable[[Any], Hashable]] = identity,
+        keep: str = "first",  # Literal["first", "last"]
+        **kwargs
+    ):
+        self.n = n
+        self.key = key
+        self.keep = keep
+        self._buffer = {}
+        self._metadata_buffer = {}
+        Stream.__init__(self, upstream, **kwargs)
+
+    def _get_key(self, x):
+        if callable(self.key):
+            return self.key(x)
+        else:
+            return x[self.key]
+
+    def update(self, x, who=None, metadata=None):
+        self._retain_refs(metadata)
+        y = self._get_key(x)
+        if self.keep == "last":
+            # remove key if already present so that emitted value
+            # will reflect elements' actual relative ordering
+            self._buffer.pop(y, None)
+            self._metadata_buffer.pop(y, None)
+            self._buffer[y] = x
+            self._metadata_buffer[y] = metadata
+        else:  # self.keep == "first"
+            if y not in self._buffer:
+                self._buffer[y] = x
+                self._metadata_buffer[y] = metadata
+        if len(self._buffer) == self.n:
+            result, self._buffer = tuple(self._buffer.values()), {}
+            metadata_result, self._metadata_buffer = list(self._metadata_buffer.values()), {}
+            ret = self._emit(result, metadata_result)
+            self._release_refs(metadata_result)
+            return ret
+        else:
+            return []
+
+
 @Stream.register_api()
 class sliding_window(Stream):
     """ Produce overlapping tuples of size n