feat: allow creating drop message for accumulator

packages/pynumaflow-lite/Cargo.toml

@@ -48,6 +48,10 @@ path = "tests/bin/session_reduce.rs"
 name = "test_accumulator"
 path = "tests/bin/accumulator.rs"
 
+[[bin]]
+name = "test_accumulator_blackhole"
+path = "tests/bin/accumulator_blackhole.rs"
+
 [[bin]]
 name = "test_reducestream"
 path = "tests/bin/reducestream.rs"

packages/pynumaflow-lite/pynumaflow_lite/_accumulator_dtypes.py

@@ -22,7 +22,7 @@ async def handler(self, datums: AsyncIterator[Datum]) -> AsyncIterator[Message]:
         """
         Accumulate can read unordered from the input stream and emit the ordered data to the output stream.
         Once the watermark (WM) of the output stream progresses, the data in WAL until that WM will be garbage collected.
-        NOTE: A message can be silently dropped if need be, and it will be cleared from the WAL when the WM progresses.
+        NOTE: To drop a datum, yield `Message.message_to_drop(datum)` so the WM can progress.
 
         Args:
             datums: An async iterator of Datum objects

packages/pynumaflow-lite/pynumaflow_lite/accumulator.pyi

@@ -25,9 +25,14 @@ class Message:
         watermark: datetime | None = None,
     ) -> None: ...
     @staticmethod
-    def message_to_drop() -> Message:
+    def message_to_drop(datum: Datum) -> Message:
         """
-        Drop a Message, do not forward to the next vertex.
+        Builds a Message from the given Datum with drop tags set, so the message is not
+        forwarded to the next vertex but still allows the accumulator to advance the watermark
+        and release tracked state.
+
+        Args:
+            datum: The input Datum to drop the results for
         """
         ...
 

packages/pynumaflow-lite/src/accumulate/mod.rs

@@ -68,18 +68,20 @@ impl Message {
         }
     }
 
-    /// Drop a Message, do not forward to the next vertex.
-    #[pyo3(signature = ())]
+    /// Builds a Message from the given Datum with drop tags set, so the message is not
+    /// forwarded to the next vertex but still allows the accumulator to advance the watermark
+    /// and release tracked state.
+    #[pyo3(signature = (datum))]
     #[staticmethod]
-    fn message_to_drop() -> Self {
+    fn message_to_drop(datum: &Datum) -> Self {
         Self {
-            keys: None,
+            keys: Some(datum.keys.clone()),
             value: vec![],
             tags: Some(vec![numaflow::shared::DROP.to_string()]),
-            id: String::new(),
-            headers: HashMap::new(),
-            event_time: chrono::Utc::now(),
-            watermark: chrono::Utc::now(),
+            id: datum.id.clone(),
+            headers: datum.headers.clone(),
+            event_time: datum.event_time,
+            watermark: datum.watermark,
         }
     }
 

packages/pynumaflow-lite/tests/bin/accumulator_blackhole.rs

@@ -0,0 +1,135 @@
+use std::collections::HashMap;
+use std::env;
+use std::path::PathBuf;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use tokio::net::UnixStream;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::{Request, transport::Uri};
+use tower::service_fn;
+
+use numaflow::proto::accumulator as acc_proto;
+use numaflow::proto::accumulator::accumulator_request::window_operation::Event;
+
+fn ts_from_secs(secs: i64) -> prost_types::Timestamp {
+    prost_types::Timestamp {
+        seconds: secs,
+        nanos: 0,
+    }
+}
+
+fn keyed_window(base_time: i64) -> acc_proto::KeyedWindow {
+    acc_proto::KeyedWindow {
+        start: Some(ts_from_secs(base_time)),
+        end: Some(ts_from_secs(base_time + 60)),
+        slot: "slot-0".to_string(),
+        keys: vec!["key1".into()],
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let sock_path = env::args()
+        .nth(1)
+        .or_else(|| env::var("NUMAFLOW_ACCUMULATOR_SOCK").ok())
+        .unwrap_or_else(|| "/tmp/var/run/numaflow/accumulator.sock".to_string());
+
+    let channel = tonic::transport::Endpoint::try_from("http://[::]:50051")?
+        .connect_with_connector(service_fn(move |_: Uri| {
+            let sock = PathBuf::from(sock_path.clone());
+            async move {
+                Ok::<_, std::io::Error>(hyper_util::rt::TokioIo::new(
+                    UnixStream::connect(sock).await?,
+                ))
+            }
+        }))
+        .await?;
+
+    let mut client =
+        numaflow::proto::accumulator::accumulator_client::AccumulatorClient::new(channel);
+
+    let (tx, rx) = mpsc::channel(16);
+
+    let base_time = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() as i64;
+
+    // (id, event_time offset, watermark offset)
+    let inputs = [("msg1", 30, 5), ("msg2", 10, 15), ("msg3", 20, 25)];
+    let headers = HashMap::from([("h1".to_string(), "v1".to_string())]);
+
+    for (i, (id, et, wm)) in inputs.iter().enumerate() {
+        let event = if i == 0 { Event::Open } else { Event::Append };
+        tx.send(acc_proto::AccumulatorRequest {
+            payload: Some(acc_proto::Payload {
+                keys: vec!["key1".into()],
+                value: format!("value-{id}").into_bytes(),
+                watermark: Some(ts_from_secs(base_time + wm)),
+                event_time: Some(ts_from_secs(base_time + et)),
+                headers: headers.clone(),
+                id: id.to_string(),
+            }),
+            operation: Some(acc_proto::accumulator_request::WindowOperation {
+                event: event as i32,
+                keyed_window: Some(keyed_window(base_time)),
+            }),
+        })
+        .await?;
+    }
+
+    tx.send(acc_proto::AccumulatorRequest {
+        payload: None,
+        operation: Some(acc_proto::accumulator_request::WindowOperation {
+            event: Event::Close as i32,
+            keyed_window: Some(keyed_window(base_time)),
+        }),
+    })
+    .await?;
+    drop(tx);
+
+    let request = Request::new(ReceiverStream::new(rx));
+    let mut resp = client.accumulate_fn(request).await?.into_inner();
+
+    let mut dropped = Vec::new();
+    let mut found_eof = false;
+
+    while let Some(r) = resp.message().await? {
+        if r.eof {
+            found_eof = true;
+            continue;
+        }
+        assert_eq!(
+            r.tags,
+            vec![numaflow::shared::DROP.to_string()],
+            "Every response should carry the DROP tag"
+        );
+        let payload = r.payload.expect("Drop response should carry a payload");
+        assert!(
+            payload.value.is_empty(),
+            "Drop message value should be empty"
+        );
+        assert_eq!(payload.keys, vec!["key1".to_string()]);
+        assert_eq!(payload.headers, headers);
+        dropped.push(payload);
+    }
+
+    assert!(found_eof, "Should have received EOF");
+    assert_eq!(dropped.len(), inputs.len(), "Expected one drop per datum");
+
+    for (payload, (id, et, wm)) in dropped.iter().zip(inputs.iter()) {
+        assert_eq!(payload.id, *id, "Drop message id should match the datum");
+        assert_eq!(
+            payload.event_time,
+            Some(ts_from_secs(base_time + et)),
+            "Drop message event_time should match the datum"
+        );
+        assert_eq!(
+            payload.watermark,
+            Some(ts_from_secs(base_time + wm)),
+            "Drop message watermark should match the datum"
+        );
+    }
+
+    println!("All datums were dropped with their metadata preserved!");
+
+    Ok(())
+}

packages/pynumaflow-lite/tests/examples/accumulator_blackhole.py

@@ -0,0 +1,49 @@
+"""
+Blackhole accumulator example.
+
+This accumulator discards every datum it receives. Instead of yielding nothing, it yields a
+drop message for each datum so the watermark can progress and the tracked state can be released.
+"""
+
+import asyncio
+import signal
+from collections.abc import AsyncIterator
+
+from pynumaflow_lite.accumulator import (
+    Accumulator,
+    AccumulatorAsyncServer,
+    Datum,
+    Message,
+)
+
+
+class Blackhole(Accumulator):
+    async def handler(self, datums: AsyncIterator[Datum]) -> AsyncIterator[Message]:
+        async for datum in datums:
+            print(f"Dropping datum: id={datum.id}, event_time={datum.event_time}")
+            yield Message.message_to_drop(datum)
+
+
+async def main():
+    sock_file = "/tmp/var/run/numaflow/accumulator.sock"
+    server_info_file = "/tmp/var/run/numaflow/accumulator-server-info"
+    server = AccumulatorAsyncServer(sock_file, server_info_file)
+
+    loop = asyncio.get_running_loop()
+    try:
+        loop.add_signal_handler(signal.SIGINT, lambda: server.stop())
+        loop.add_signal_handler(signal.SIGTERM, lambda: server.stop())
+    except (NotImplementedError, RuntimeError):
+        pass
+
+    try:
+        print("Starting Blackhole Accumulator Server...")
+        await server.start(Blackhole)
+        print("Shutting down gracefully...")
+    except asyncio.CancelledError:
+        server.stop()
+        return
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

packages/pynumaflow-lite/tests/test_accumulator.py

@@ -21,3 +21,13 @@ def test_python_accumulator_server_and_rust_client(script: str, tmp_path: Path):
         rust_bin_name="test_accumulator",
         rust_bin_args=[str(SOCK_PATH)],
     )
+
+
+def test_python_accumulator_drop_messages(tmp_path: Path):
+    run_python_server_with_rust_client(
+        script="accumulator_blackhole.py",
+        sock_path=SOCK_PATH,
+        server_info_path=SERVER_INFO,
+        rust_bin_name="test_accumulator_blackhole",
+        rust_bin_args=[str(SOCK_PATH)],
+    )