Skip to main content

fraiseql_wire/stream/filter/
mod.rs

1//! Filtered JSON stream
2
3use crate::Result;
4use futures::stream::Stream;
5use serde_json::Value;
6use std::pin::Pin;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::task::{Context, Poll};
9
10/// Predicate function type
11pub type Predicate = Box<dyn Fn(&Value) -> bool + Send>;
12
13/// Filtered JSON stream
14pub struct FilteredStream<S> {
15    inner: S,
16    predicate: Predicate,
17    // Sampling counter for metrics recording
18    eval_count: AtomicU64,
19}
20
21impl<S> FilteredStream<S> {
22    /// Create new filtered stream
23    pub fn new(inner: S, predicate: Predicate) -> Self {
24        Self {
25            inner,
26            predicate,
27            eval_count: AtomicU64::new(0),
28        }
29    }
30}
31
32impl<S> Stream for FilteredStream<S>
33where
34    S: Stream<Item = Result<Value>> + Unpin,
35{
36    type Item = Result<Value>;
37
38    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
39        loop {
40            match Pin::new(&mut self.inner).poll_next(cx) {
41                Poll::Ready(Some(Ok(value))) => {
42                    // Apply predicate
43                    // Sample timing: only record 1 in 1000 evaluations
44                    let eval_idx = self.eval_count.fetch_add(1, Ordering::Relaxed);
45                    let passed = if eval_idx.is_multiple_of(1000) {
46                        // Record timing for sampled evaluation
47                        let filter_start = std::time::Instant::now();
48                        let result = (self.predicate)(&value);
49                        let filter_duration = filter_start.elapsed().as_millis() as u64;
50                        crate::metrics::histograms::filter_duration("unknown", filter_duration);
51                        result
52                    } else {
53                        // No timing, just evaluate
54                        (self.predicate)(&value)
55                    };
56
57                    if passed {
58                        return Poll::Ready(Some(Ok(value)));
59                    }
60                    // Predicate failed, try next value (filter out this row)
61                    crate::metrics::counters::rows_filtered("unknown", 1);
62                    continue;
63                }
64                Poll::Ready(Some(Err(e))) => {
65                    // Propagate errors
66                    return Poll::Ready(Some(Err(e)));
67                }
68                Poll::Ready(None) => {
69                    // End of stream
70                    return Poll::Ready(None);
71                }
72                Poll::Pending => {
73                    return Poll::Pending;
74                }
75            }
76        }
77    }
78}
79
80#[cfg(test)]
81mod tests;