llm_optimizer_processor/pipeline/
operator.rs

1//! Stream operators for transforming and processing events
2//!
3//! This module provides the core operators for stream processing pipelines:
4//! - Map: Transform events one-to-one
5//! - Filter: Select events based on predicates
6//! - FlatMap: Transform events one-to-many
7//! - KeyBy: Partition events by key
8//! - Window: Assign events to windows
9//! - Aggregate: Compute aggregations over windows
10
11use crate::aggregation::Aggregator;
12use crate::core::{EventKey, EventTimeExtractor, KeyExtractor};
13use crate::error::Result;
14use crate::watermark::Watermark;
15use crate::window::{Window, WindowAssigner};
16use async_trait::async_trait;
17use serde::Serialize;
18use std::fmt::Debug;
19use std::sync::Arc;
20
21/// Context provided to operators during execution
22#[derive(Debug, Clone)]
23pub struct OperatorContext {
24    /// Current watermark
25    pub watermark: Watermark,
26
27    /// Partition ID
28    pub partition: u32,
29
30    /// Processing timestamp
31    pub processing_time: i64,
32}
33
34impl OperatorContext {
35    /// Create a new operator context
36    pub fn new(watermark: Watermark, partition: u32) -> Self {
37        Self {
38            watermark,
39            partition,
40            processing_time: chrono::Utc::now().timestamp_millis(),
41        }
42    }
43}
44
45/// Trait for stream operators
46///
47/// All stream processing operators implement this trait to provide a uniform
48/// interface for processing events in the pipeline.
49#[async_trait]
50pub trait StreamOperator: Send + Sync + Debug {
51    /// Process an event and return transformed events
52    ///
53    /// # Arguments
54    /// * `input` - The input event data
55    /// * `ctx` - Operator context with watermark and partition info
56    ///
57    /// # Returns
58    /// A vector of output events (may be empty, one, or many)
59    async fn process(
60        &self,
61        input: Vec<u8>,
62        ctx: &OperatorContext,
63    ) -> Result<Vec<Vec<u8>>>;
64
65    /// Get the operator name
66    fn name(&self) -> &str;
67
68    /// Called when the operator is initialized
69    async fn open(&mut self) -> Result<()> {
70        Ok(())
71    }
72
73    /// Called when the operator is shut down
74    async fn close(&mut self) -> Result<()> {
75        Ok(())
76    }
77}
78
79/// Map operator - transforms events one-to-one
80///
81/// # Example
82///
83/// ```rust
84/// use processor::pipeline::operator::MapOperator;
85///
86/// // Create a map operator that doubles values
87/// let operator = MapOperator::new("double", |x: i32| x * 2);
88/// ```
89#[derive(Clone)]
90pub struct MapOperator<F, T, U>
91where
92    F: Fn(T) -> U + Send + Sync + Clone,
93    T: Send + Sync + Clone,
94    U: Send + Sync + Clone,
95{
96    name: String,
97    func: F,
98    _phantom: std::marker::PhantomData<(T, U)>,
99}
100
101impl<F, T, U> MapOperator<F, T, U>
102where
103    F: Fn(T) -> U + Send + Sync + Clone,
104    T: Send + Sync + Clone + serde::de::DeserializeOwned,
105    U: Send + Sync + Clone + Serialize,
106{
107    /// Create a new map operator
108    pub fn new<S: Into<String>>(name: S, func: F) -> Self {
109        Self {
110            name: name.into(),
111            func,
112            _phantom: std::marker::PhantomData,
113        }
114    }
115}
116
117impl<F, T, U> Debug for MapOperator<F, T, U>
118where
119    F: Fn(T) -> U + Send + Sync + Clone,
120    T: Send + Sync + Clone,
121    U: Send + Sync + Clone,
122{
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("MapOperator")
125            .field("name", &self.name)
126            .finish()
127    }
128}
129
130#[async_trait]
131impl<F, T, U> StreamOperator for MapOperator<F, T, U>
132where
133    F: Fn(T) -> U + Send + Sync + Clone,
134    T: Send + Sync + Clone + serde::de::DeserializeOwned,
135    U: Send + Sync + Clone + Serialize,
136{
137    async fn process(
138        &self,
139        input: Vec<u8>,
140        _ctx: &OperatorContext,
141    ) -> Result<Vec<Vec<u8>>> {
142        let value: T = bincode::deserialize(&input)?;
143        let result = (self.func)(value);
144        let serialized = bincode::serialize(&result)?;
145        Ok(vec![serialized])
146    }
147
148    fn name(&self) -> &str {
149        &self.name
150    }
151}
152
153/// Filter operator - selects events based on a predicate
154///
155/// # Example
156///
157/// ```rust
158/// use processor::pipeline::operator::FilterOperator;
159///
160/// // Create a filter that only passes even numbers
161/// let operator = FilterOperator::new("even", |x: &i32| x % 2 == 0);
162/// ```
163#[derive(Clone)]
164pub struct FilterOperator<F, T>
165where
166    F: Fn(&T) -> bool + Send + Sync + Clone,
167    T: Send + Sync + Clone,
168{
169    name: String,
170    predicate: F,
171    _phantom: std::marker::PhantomData<T>,
172}
173
174impl<F, T> FilterOperator<F, T>
175where
176    F: Fn(&T) -> bool + Send + Sync + Clone,
177    T: Send + Sync + Clone + serde::de::DeserializeOwned,
178{
179    /// Create a new filter operator
180    pub fn new<S: Into<String>>(name: S, predicate: F) -> Self {
181        Self {
182            name: name.into(),
183            predicate,
184            _phantom: std::marker::PhantomData,
185        }
186    }
187}
188
189impl<F, T> Debug for FilterOperator<F, T>
190where
191    F: Fn(&T) -> bool + Send + Sync + Clone,
192    T: Send + Sync + Clone,
193{
194    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
195        f.debug_struct("FilterOperator")
196            .field("name", &self.name)
197            .finish()
198    }
199}
200
201#[async_trait]
202impl<F, T> StreamOperator for FilterOperator<F, T>
203where
204    F: Fn(&T) -> bool + Send + Sync + Clone,
205    T: Send + Sync + Clone + serde::de::DeserializeOwned + Serialize,
206{
207    async fn process(
208        &self,
209        input: Vec<u8>,
210        _ctx: &OperatorContext,
211    ) -> Result<Vec<Vec<u8>>> {
212        let value: T = bincode::deserialize(&input)?;
213        if (self.predicate)(&value) {
214            Ok(vec![input])
215        } else {
216            Ok(vec![])
217        }
218    }
219
220    fn name(&self) -> &str {
221        &self.name
222    }
223}
224
225/// FlatMap operator - transforms events one-to-many
226///
227/// # Example
228///
229/// ```rust
230/// use processor::pipeline::operator::FlatMapOperator;
231///
232/// // Create a flatmap that splits strings into words
233/// let operator = FlatMapOperator::new("split_words", |s: String| {
234///     s.split_whitespace().map(|w| w.to_string()).collect::<Vec<_>>()
235/// });
236/// ```
237#[derive(Clone)]
238pub struct FlatMapOperator<F, T, U>
239where
240    F: Fn(T) -> Vec<U> + Send + Sync + Clone,
241    T: Send + Sync + Clone,
242    U: Send + Sync + Clone,
243{
244    name: String,
245    func: F,
246    _phantom: std::marker::PhantomData<(T, U)>,
247}
248
249impl<F, T, U> FlatMapOperator<F, T, U>
250where
251    F: Fn(T) -> Vec<U> + Send + Sync + Clone,
252    T: Send + Sync + Clone + serde::de::DeserializeOwned,
253    U: Send + Sync + Clone + Serialize,
254{
255    /// Create a new flatmap operator
256    pub fn new<S: Into<String>>(name: S, func: F) -> Self {
257        Self {
258            name: name.into(),
259            func,
260            _phantom: std::marker::PhantomData,
261        }
262    }
263}
264
265impl<F, T, U> Debug for FlatMapOperator<F, T, U>
266where
267    F: Fn(T) -> Vec<U> + Send + Sync + Clone,
268    T: Send + Sync + Clone,
269    U: Send + Sync + Clone,
270{
271    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        f.debug_struct("FlatMapOperator")
273            .field("name", &self.name)
274            .finish()
275    }
276}
277
278#[async_trait]
279impl<F, T, U> StreamOperator for FlatMapOperator<F, T, U>
280where
281    F: Fn(T) -> Vec<U> + Send + Sync + Clone,
282    T: Send + Sync + Clone + serde::de::DeserializeOwned,
283    U: Send + Sync + Clone + Serialize,
284{
285    async fn process(
286        &self,
287        input: Vec<u8>,
288        _ctx: &OperatorContext,
289    ) -> Result<Vec<Vec<u8>>> {
290        let value: T = bincode::deserialize(&input)?;
291        let results = (self.func)(value);
292        let mut serialized = Vec::new();
293        for result in results {
294            serialized.push(bincode::serialize(&result)?);
295        }
296        Ok(serialized)
297    }
298
299    fn name(&self) -> &str {
300        &self.name
301    }
302}
303
304/// KeyBy operator - partitions events by key
305///
306/// # Example
307///
308/// ```rust
309/// use processor::pipeline::operator::KeyByOperator;
310/// use processor::core::{EventKey, KeyExtractor};
311///
312/// // Define a key extractor
313/// struct MyKeyExtractor;
314/// impl KeyExtractor<i32> for MyKeyExtractor {
315///     fn extract_key(&self, event: &i32) -> EventKey {
316///         EventKey::Session(event.to_string())
317///     }
318/// }
319///
320/// let operator = KeyByOperator::new("partition", MyKeyExtractor);
321/// ```
322pub struct KeyByOperator<T, E>
323where
324    T: Send + Sync + Clone,
325    E: KeyExtractor<T> + Send + Sync,
326{
327    name: String,
328    extractor: E,
329    _phantom: std::marker::PhantomData<T>,
330}
331
332impl<T, E> KeyByOperator<T, E>
333where
334    T: Send + Sync + Clone + serde::de::DeserializeOwned,
335    E: KeyExtractor<T> + Send + Sync,
336{
337    /// Create a new key-by operator
338    pub fn new<S: Into<String>>(name: S, extractor: E) -> Self {
339        Self {
340            name: name.into(),
341            extractor,
342            _phantom: std::marker::PhantomData,
343        }
344    }
345
346    /// Extract the key from an event
347    pub fn extract_key(&self, event: &T) -> EventKey {
348        self.extractor.extract_key(event)
349    }
350}
351
352impl<T, E> Debug for KeyByOperator<T, E>
353where
354    T: Send + Sync + Clone,
355    E: KeyExtractor<T> + Send + Sync,
356{
357    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
358        f.debug_struct("KeyByOperator")
359            .field("name", &self.name)
360            .finish()
361    }
362}
363
364#[async_trait]
365impl<T, E> StreamOperator for KeyByOperator<T, E>
366where
367    T: Send + Sync + Clone + serde::de::DeserializeOwned + Serialize,
368    E: KeyExtractor<T> + Send + Sync,
369{
370    async fn process(
371        &self,
372        input: Vec<u8>,
373        _ctx: &OperatorContext,
374    ) -> Result<Vec<Vec<u8>>> {
375        // KeyBy doesn't transform data, just extracts keys for partitioning
376        // The actual partitioning is handled by the executor
377        Ok(vec![input])
378    }
379
380    fn name(&self) -> &str {
381        &self.name
382    }
383}
384
385/// Window operator - assigns events to windows
386///
387/// # Example
388///
389/// ```rust
390/// use processor::pipeline::operator::WindowOperator;
391/// use processor::window::{TumblingWindowAssigner, WindowAssigner};
392/// use processor::core::EventTimeExtractor;
393/// use std::time::Duration;
394///
395/// // Define a time extractor
396/// struct MyTimeExtractor;
397/// impl EventTimeExtractor<i64> for MyTimeExtractor {
398///     fn extract_event_time(&self, event: &i64) -> chrono::DateTime<chrono::Utc> {
399///         chrono::DateTime::from_timestamp(*event, 0).unwrap()
400///     }
401/// }
402///
403/// let assigner = TumblingWindowAssigner::new(Duration::from_secs(60));
404/// let operator = WindowOperator::new("windows", assigner, MyTimeExtractor);
405/// ```
406pub struct WindowOperator<T, A, E>
407where
408    T: Send + Sync + Clone,
409    A: WindowAssigner + Send + Sync,
410    E: EventTimeExtractor<T> + Send + Sync,
411{
412    name: String,
413    assigner: A,
414    time_extractor: E,
415    _phantom: std::marker::PhantomData<T>,
416}
417
418impl<T, A, E> WindowOperator<T, A, E>
419where
420    T: Send + Sync + Clone + serde::de::DeserializeOwned,
421    A: WindowAssigner + Send + Sync,
422    E: EventTimeExtractor<T> + Send + Sync,
423{
424    /// Create a new window operator
425    pub fn new<S: Into<String>>(name: S, assigner: A, time_extractor: E) -> Self {
426        Self {
427            name: name.into(),
428            assigner,
429            time_extractor,
430            _phantom: std::marker::PhantomData,
431        }
432    }
433
434    /// Assign an event to windows
435    pub fn assign_windows(&self, event: &T, _key: &str) -> Vec<Window> {
436        let timestamp = self.time_extractor.extract_event_time(event);
437        self.assigner.assign_windows(timestamp)
438    }
439}
440
441impl<T, A, E> Debug for WindowOperator<T, A, E>
442where
443    T: Send + Sync + Clone,
444    A: WindowAssigner + Send + Sync,
445    E: EventTimeExtractor<T> + Send + Sync,
446{
447    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
448        f.debug_struct("WindowOperator")
449            .field("name", &self.name)
450            .finish()
451    }
452}
453
454#[async_trait]
455impl<T, A, E> StreamOperator for WindowOperator<T, A, E>
456where
457    T: Send + Sync + Clone + serde::de::DeserializeOwned + Serialize,
458    A: WindowAssigner + Send + Sync,
459    E: EventTimeExtractor<T> + Send + Sync,
460{
461    async fn process(
462        &self,
463        input: Vec<u8>,
464        _ctx: &OperatorContext,
465    ) -> Result<Vec<Vec<u8>>> {
466        // Window operator doesn't transform data, just assigns windows
467        // The actual windowing logic is handled by the executor
468        Ok(vec![input])
469    }
470
471    fn name(&self) -> &str {
472        &self.name
473    }
474}
475
476/// Aggregate operator - computes aggregations over windows
477///
478/// # Example
479///
480/// ```rust
481/// use processor::pipeline::operator::AggregateOperator;
482/// use processor::aggregation::{Aggregator, SumAggregator};
483///
484/// let aggregator = SumAggregator::new();
485/// let operator = AggregateOperator::new("sum", aggregator);
486/// ```
487pub struct AggregateOperator<T, A>
488where
489    T: Send + Sync + Clone,
490    A: Aggregator<Input = T> + Send + Sync,
491{
492    name: String,
493    aggregator: Arc<std::sync::Mutex<A>>,
494    _phantom: std::marker::PhantomData<T>,
495}
496
497impl<T, A> AggregateOperator<T, A>
498where
499    T: Send + Sync + Clone,
500    A: Aggregator<Input = T> + Send + Sync,
501{
502    /// Create a new aggregate operator
503    pub fn new<S: Into<String>>(name: S, aggregator: A) -> Self {
504        Self {
505            name: name.into(),
506            aggregator: Arc::new(std::sync::Mutex::new(aggregator)),
507            _phantom: std::marker::PhantomData,
508        }
509    }
510
511    /// Get a reference to the aggregator
512    pub fn aggregator(&self) -> &Arc<std::sync::Mutex<A>> {
513        &self.aggregator
514    }
515}
516
517impl<T, A> Debug for AggregateOperator<T, A>
518where
519    T: Send + Sync + Clone,
520    A: Aggregator<Input = T> + Send + Sync,
521{
522    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
523        f.debug_struct("AggregateOperator")
524            .field("name", &self.name)
525            .finish()
526    }
527}
528
529impl<T, A> Clone for AggregateOperator<T, A>
530where
531    T: Send + Sync + Clone,
532    A: Aggregator<Input = T> + Send + Sync,
533{
534    fn clone(&self) -> Self {
535        Self {
536            name: self.name.clone(),
537            aggregator: Arc::clone(&self.aggregator),
538            _phantom: std::marker::PhantomData,
539        }
540    }
541}
542
543#[async_trait]
544impl<T, A> StreamOperator for AggregateOperator<T, A>
545where
546    T: Send + Sync + Clone + serde::de::DeserializeOwned + Serialize,
547    A: Aggregator<Input = T, Output = T> + Send + Sync,
548{
549    async fn process(
550        &self,
551        input: Vec<u8>,
552        _ctx: &OperatorContext,
553    ) -> Result<Vec<Vec<u8>>> {
554        let value: T = bincode::deserialize(&input)?;
555        let mut agg = self.aggregator.lock().unwrap();
556        agg.update(value)?;
557        let result = agg.finalize()?;
558        let serialized = bincode::serialize(&result)?;
559        Ok(vec![serialized])
560    }
561
562    fn name(&self) -> &str {
563        &self.name
564    }
565}
566
567/// Deduplication operator - filters duplicate events
568///
569/// This operator filters out duplicate events based on a deduplication strategy.
570/// It maintains a set of seen event signatures and drops events that match
571/// previously processed signatures within the configured TTL window.
572///
573/// # Example
574///
575/// ```rust
576/// use processor::pipeline::operator::DeduplicationOperator;
577/// use processor::config::{DeduplicationConfig, DeduplicationStrategy};
578/// use std::time::Duration;
579///
580/// # async fn example() -> anyhow::Result<()> {
581/// let config = DeduplicationConfig::new()
582///     .enabled()
583///     .with_ttl(Duration::from_secs(3600))
584///     .with_strategy(DeduplicationStrategy::ContentHash);
585///
586/// let operator = DeduplicationOperator::new("dedup", config);
587/// # Ok(())
588/// # }
589/// ```
590pub struct DeduplicationOperator {
591    name: String,
592    config: crate::config::DeduplicationConfig,
593    // In-memory cache for recently seen event hashes
594    // In production, this would be backed by Redis
595    seen_hashes: Arc<std::sync::RwLock<std::collections::HashMap<String, std::time::Instant>>>,
596}
597
598impl DeduplicationOperator {
599    /// Create a new deduplication operator
600    pub fn new<S: Into<String>>(name: S, config: crate::config::DeduplicationConfig) -> Self {
601        Self {
602            name: name.into(),
603            config,
604            seen_hashes: Arc::new(std::sync::RwLock::new(std::collections::HashMap::new())),
605        }
606    }
607
608    /// Compute the signature for an event based on the configured strategy
609    fn compute_signature(&self, data: &[u8], _strategy: crate::config::DeduplicationStrategy) -> String {
610        // For now, use simple content hashing
611        // In production, this would handle different strategies
612        use sha2::{Digest, Sha256};
613        let mut hasher = Sha256::new();
614        hasher.update(data);
615        format!("{:x}", hasher.finalize())
616    }
617
618    /// Check if an event is a duplicate
619    fn is_duplicate(&self, signature: &str) -> bool {
620        let seen = self.seen_hashes.read().unwrap();
621        if let Some(timestamp) = seen.get(signature) {
622            // Check if the signature is still within TTL
623            timestamp.elapsed() < self.config.ttl
624        } else {
625            false
626        }
627    }
628
629    /// Mark an event signature as seen
630    fn mark_seen(&self, signature: String) {
631        let mut seen = self.seen_hashes.write().unwrap();
632        seen.insert(signature, std::time::Instant::now());
633
634        // Periodic cleanup of expired entries (simple implementation)
635        // In production, this would be done in a background task
636        if seen.len() > 10000 {
637            let ttl = self.config.ttl;
638            seen.retain(|_, timestamp| timestamp.elapsed() < ttl);
639        }
640    }
641
642    /// Get deduplication statistics
643    pub fn stats(&self) -> DeduplicationStats {
644        let seen = self.seen_hashes.read().unwrap();
645        DeduplicationStats {
646            cached_signatures: seen.len(),
647            enabled: self.config.enabled,
648        }
649    }
650}
651
652impl Debug for DeduplicationOperator {
653    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
654        f.debug_struct("DeduplicationOperator")
655            .field("name", &self.name)
656            .field("enabled", &self.config.enabled)
657            .field("strategy", &self.config.strategy)
658            .finish()
659    }
660}
661
662impl Clone for DeduplicationOperator {
663    fn clone(&self) -> Self {
664        Self {
665            name: self.name.clone(),
666            config: self.config.clone(),
667            seen_hashes: Arc::clone(&self.seen_hashes),
668        }
669    }
670}
671
672#[async_trait]
673impl StreamOperator for DeduplicationOperator {
674    async fn process(
675        &self,
676        input: Vec<u8>,
677        _ctx: &OperatorContext,
678    ) -> Result<Vec<Vec<u8>>> {
679        // If deduplication is disabled, pass through
680        if !self.config.enabled {
681            return Ok(vec![input]);
682        }
683
684        // Compute event signature
685        let signature = self.compute_signature(&input, self.config.strategy);
686
687        // Check for duplicates
688        if self.is_duplicate(&signature) {
689            // Drop duplicate event
690            Ok(vec![])
691        } else {
692            // Mark as seen and pass through
693            self.mark_seen(signature);
694            Ok(vec![input])
695        }
696    }
697
698    fn name(&self) -> &str {
699        &self.name
700    }
701}
702
703/// Statistics for deduplication operator
704#[derive(Debug, Clone)]
705pub struct DeduplicationStats {
706    /// Number of cached event signatures
707    pub cached_signatures: usize,
708    /// Whether deduplication is enabled
709    pub enabled: bool,
710}
711
712/// Chain of operators that can be applied sequentially
713pub struct OperatorChain {
714    operators: Vec<Arc<dyn StreamOperator>>,
715}
716
717impl OperatorChain {
718    /// Create a new operator chain
719    pub fn new() -> Self {
720        Self {
721            operators: Vec::new(),
722        }
723    }
724
725    /// Add an operator to the chain
726    pub fn add<O: StreamOperator + 'static>(&mut self, operator: O) {
727        self.operators.push(Arc::new(operator));
728    }
729
730    /// Get all operators in the chain
731    pub fn operators(&self) -> &[Arc<dyn StreamOperator>] {
732        &self.operators
733    }
734
735    /// Process an event through the entire operator chain
736    pub async fn process(
737        &self,
738        mut input: Vec<Vec<u8>>,
739        ctx: &OperatorContext,
740    ) -> Result<Vec<Vec<u8>>> {
741        for operator in &self.operators {
742            let mut output = Vec::new();
743            for item in input {
744                let results = operator.process(item, ctx).await?;
745                output.extend(results);
746            }
747            input = output;
748        }
749        Ok(input)
750    }
751
752    /// Get the number of operators in the chain
753    pub fn len(&self) -> usize {
754        self.operators.len()
755    }
756
757    /// Check if the chain is empty
758    pub fn is_empty(&self) -> bool {
759        self.operators.is_empty()
760    }
761}
762
763impl Default for OperatorChain {
764    fn default() -> Self {
765        Self::new()
766    }
767}
768
769impl Debug for OperatorChain {
770    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
771        f.debug_struct("OperatorChain")
772            .field("operator_count", &self.operators.len())
773            .finish()
774    }
775}
776
777#[cfg(test)]
778mod tests {
779    use super::*;
780    use crate::aggregation::SumAggregator;
781    use crate::core::{EventKey, FeedbackEventKeyExtractor, FeedbackEventTimeExtractor};
782    use crate::window::TumblingWindowAssigner;
783    use llm_optimizer_types::events::{EventSource, EventType, FeedbackEvent};
784    use std::time::Duration;
785
786    #[tokio::test]
787    async fn test_map_operator() {
788        let operator = MapOperator::new("double", |x: i32| x * 2);
789        let ctx = OperatorContext::new(Watermark::min(), 0);
790
791        let input = bincode::serialize(&5).unwrap();
792        let output = operator.process(input, &ctx).await.unwrap();
793
794        assert_eq!(output.len(), 1);
795        let result: i32 = bincode::deserialize(&output[0]).unwrap();
796        assert_eq!(result, 10);
797    }
798
799    #[tokio::test]
800    async fn test_filter_operator() {
801        let operator = FilterOperator::new("even", |x: &i32| x % 2 == 0);
802        let ctx = OperatorContext::new(Watermark::min(), 0);
803
804        // Test with even number (should pass)
805        let input = bincode::serialize(&4).unwrap();
806        let output = operator.process(input, &ctx).await.unwrap();
807        assert_eq!(output.len(), 1);
808
809        // Test with odd number (should filter out)
810        let input = bincode::serialize(&5).unwrap();
811        let output = operator.process(input, &ctx).await.unwrap();
812        assert_eq!(output.len(), 0);
813    }
814
815    #[tokio::test]
816    async fn test_flatmap_operator() {
817        let operator = FlatMapOperator::new("range", |n: i32| (0..n).collect::<Vec<_>>());
818        let ctx = OperatorContext::new(Watermark::min(), 0);
819
820        let input = bincode::serialize(&3).unwrap();
821        let output = operator.process(input, &ctx).await.unwrap();
822
823        assert_eq!(output.len(), 3);
824        let v0: i32 = bincode::deserialize(&output[0]).unwrap();
825        let v1: i32 = bincode::deserialize(&output[1]).unwrap();
826        let v2: i32 = bincode::deserialize(&output[2]).unwrap();
827        assert_eq!(v0, 0);
828        assert_eq!(v1, 1);
829        assert_eq!(v2, 2);
830    }
831
832    #[tokio::test]
833    async fn test_keyby_operator() {
834        let operator = KeyByOperator::new("key-by", FeedbackEventKeyExtractor);
835        let ctx = OperatorContext::new(Watermark::min(), 0);
836
837        let event = FeedbackEvent::new(
838            EventSource::Observatory,
839            EventType::Metric,
840            serde_json::json!({}),
841        );
842
843        let input = bincode::serialize(&event).unwrap();
844        let output = operator.process(input.clone(), &ctx).await.unwrap();
845
846        assert_eq!(output.len(), 1);
847        assert_eq!(output[0], input); // KeyBy doesn't transform
848    }
849
850    #[tokio::test]
851    async fn test_window_operator() {
852        let assigner = TumblingWindowAssigner::new(chrono::Duration::seconds(60));
853        let operator = WindowOperator::new("window", assigner, FeedbackEventTimeExtractor);
854        let ctx = OperatorContext::new(Watermark::min(), 0);
855
856        let event = FeedbackEvent::new(
857            EventSource::Observatory,
858            EventType::Metric,
859            serde_json::json!({}),
860        );
861
862        let input = bincode::serialize(&event).unwrap();
863        let output = operator.process(input.clone(), &ctx).await.unwrap();
864
865        assert_eq!(output.len(), 1);
866        assert_eq!(output[0], input); // Window doesn't transform
867    }
868
869    #[tokio::test]
870    async fn test_aggregate_operator() {
871        let aggregator: SumAggregator<f64> = SumAggregator::new();
872        let operator = AggregateOperator::new("sum", aggregator);
873        let ctx = OperatorContext::new(Watermark::min(), 0);
874
875        // Process first value
876        let input1 = bincode::serialize(&5.0).unwrap();
877        operator.process(input1, &ctx).await.unwrap();
878
879        // Process second value
880        let input2 = bincode::serialize(&3.0).unwrap();
881        let output = operator.process(input2, &ctx).await.unwrap();
882
883        // Result should be the sum
884        assert_eq!(output.len(), 1);
885        let result: f64 = bincode::deserialize(&output[0]).unwrap();
886        assert_eq!(result, 8.0);
887    }
888
889    #[tokio::test]
890    async fn test_operator_chain() {
891        let mut chain = OperatorChain::new();
892
893        // Add operators to the chain
894        chain.add(MapOperator::new("add-one", |x: i32| x + 1));
895        chain.add(FilterOperator::new("even", |x: &i32| x % 2 == 0));
896        chain.add(MapOperator::new("double", |x: i32| x * 2));
897
898        assert_eq!(chain.len(), 3);
899        assert!(!chain.is_empty());
900
901        let ctx = OperatorContext::new(Watermark::min(), 0);
902
903        // Test with 3: 3 -> 4 (add-one) -> 4 (filter pass) -> 8 (double)
904        let input = bincode::serialize(&3).unwrap();
905        let output = chain.process(vec![input], &ctx).await.unwrap();
906        assert_eq!(output.len(), 1);
907        let result: i32 = bincode::deserialize(&output[0]).unwrap();
908        assert_eq!(result, 8);
909
910        // Test with 4: 4 -> 5 (add-one) -> filtered out
911        let input = bincode::serialize(&4).unwrap();
912        let output = chain.process(vec![input], &ctx).await.unwrap();
913        assert_eq!(output.len(), 0);
914    }
915
916    #[tokio::test]
917    async fn test_operator_context() {
918        let watermark = Watermark::new(12345);
919        let ctx = OperatorContext::new(watermark, 7);
920
921        assert_eq!(ctx.watermark, watermark);
922        assert_eq!(ctx.partition, 7);
923        assert!(ctx.processing_time > 0);
924    }
925
926    #[test]
927    fn test_map_operator_name() {
928        let operator = MapOperator::new("my-map", |x: i32| x + 1);
929        assert_eq!(operator.name(), "my-map");
930    }
931
932    #[test]
933    fn test_filter_operator_name() {
934        let operator = FilterOperator::new("my-filter", |x: &i32| *x > 0);
935        assert_eq!(operator.name(), "my-filter");
936    }
937
938    #[test]
939    fn test_flatmap_operator_name() {
940        let operator = FlatMapOperator::new("my-flatmap", |x: i32| vec![x]);
941        assert_eq!(operator.name(), "my-flatmap");
942    }
943
944    #[test]
945    fn test_keyby_operator_extract_key() {
946        let operator = KeyByOperator::new("my-keyby", FeedbackEventKeyExtractor);
947
948        let mut event = FeedbackEvent::new(
949            EventSource::Observatory,
950            EventType::Metric,
951            serde_json::json!({}),
952        );
953        event.metadata.insert("model_id".to_string(), "test-model".to_string());
954
955        let key = operator.extract_key(&event);
956        match key {
957            EventKey::Model(id) => assert_eq!(id, "test-model"),
958            _ => panic!("Expected Model key"),
959        }
960    }
961
962    #[test]
963    fn test_window_operator_assign_windows() {
964        let assigner = TumblingWindowAssigner::new(chrono::Duration::seconds(60));
965        let operator = WindowOperator::new("my-window", assigner, FeedbackEventTimeExtractor);
966
967        let event = FeedbackEvent::new(
968            EventSource::Observatory,
969            EventType::Metric,
970            serde_json::json!({}),
971        );
972
973        let windows = operator.assign_windows(&event, "test-key");
974        assert!(!windows.is_empty());
975    }
976
977    #[test]
978    fn test_aggregate_operator_clone() {
979        let aggregator: SumAggregator<f64> = SumAggregator::new();
980        let operator = AggregateOperator::new("sum", aggregator);
981        let _cloned = operator.clone();
982    }
983
984    #[test]
985    fn test_operator_chain_default() {
986        let chain = OperatorChain::default();
987        assert!(chain.is_empty());
988        assert_eq!(chain.len(), 0);
989    }
990
991    #[tokio::test]
992    async fn test_deduplication_operator_disabled() {
993        use crate::config::DeduplicationConfig;
994
995        let config = DeduplicationConfig::default(); // disabled by default
996        let operator = DeduplicationOperator::new("dedup", config);
997        let ctx = OperatorContext::new(Watermark::min(), 0);
998
999        let input = bincode::serialize(&42).unwrap();
1000        let output1 = operator.process(input.clone(), &ctx).await.unwrap();
1001        let output2 = operator.process(input.clone(), &ctx).await.unwrap();
1002
1003        // Both should pass through when disabled
1004        assert_eq!(output1.len(), 1);
1005        assert_eq!(output2.len(), 1);
1006    }
1007
1008    #[tokio::test]
1009    async fn test_deduplication_operator_enabled() {
1010        use crate::config::DeduplicationConfig;
1011
1012        let config = DeduplicationConfig::new().enabled();
1013        let operator = DeduplicationOperator::new("dedup", config);
1014        let ctx = OperatorContext::new(Watermark::min(), 0);
1015
1016        let input = bincode::serialize(&42).unwrap();
1017
1018        // First event should pass through
1019        let output1 = operator.process(input.clone(), &ctx).await.unwrap();
1020        assert_eq!(output1.len(), 1);
1021
1022        // Duplicate event should be filtered
1023        let output2 = operator.process(input.clone(), &ctx).await.unwrap();
1024        assert_eq!(output2.len(), 0);
1025    }
1026
1027    #[tokio::test]
1028    async fn test_deduplication_operator_different_events() {
1029        use crate::config::DeduplicationConfig;
1030
1031        let config = DeduplicationConfig::new().enabled();
1032        let operator = DeduplicationOperator::new("dedup", config);
1033        let ctx = OperatorContext::new(Watermark::min(), 0);
1034
1035        let input1 = bincode::serialize(&42).unwrap();
1036        let input2 = bincode::serialize(&43).unwrap();
1037
1038        // Both should pass through as they're different
1039        let output1 = operator.process(input1, &ctx).await.unwrap();
1040        let output2 = operator.process(input2, &ctx).await.unwrap();
1041
1042        assert_eq!(output1.len(), 1);
1043        assert_eq!(output2.len(), 1);
1044    }
1045
1046    #[tokio::test]
1047    async fn test_deduplication_operator_ttl() {
1048        use crate::config::DeduplicationConfig;
1049        use std::time::Duration;
1050
1051        let config = DeduplicationConfig::new()
1052            .enabled()
1053            .with_ttl(Duration::from_millis(100));
1054
1055        let operator = DeduplicationOperator::new("dedup", config);
1056        let ctx = OperatorContext::new(Watermark::min(), 0);
1057
1058        let input = bincode::serialize(&42).unwrap();
1059
1060        // First event should pass
1061        let output1 = operator.process(input.clone(), &ctx).await.unwrap();
1062        assert_eq!(output1.len(), 1);
1063
1064        // Wait for TTL to expire
1065        tokio::time::sleep(tokio::time::Duration::from_millis(150)).await;
1066
1067        // Should pass again after TTL expires
1068        let output2 = operator.process(input.clone(), &ctx).await.unwrap();
1069        assert_eq!(output2.len(), 1);
1070    }
1071
1072    #[test]
1073    fn test_deduplication_operator_stats() {
1074        use crate::config::DeduplicationConfig;
1075
1076        let config = DeduplicationConfig::new().enabled();
1077        let operator = DeduplicationOperator::new("dedup", config);
1078
1079        let stats = operator.stats();
1080        assert_eq!(stats.cached_signatures, 0);
1081        assert!(stats.enabled);
1082    }
1083
1084    #[test]
1085    fn test_deduplication_operator_clone() {
1086        use crate::config::DeduplicationConfig;
1087
1088        let config = DeduplicationConfig::new().enabled();
1089        let operator = DeduplicationOperator::new("dedup", config);
1090        let _cloned = operator.clone();
1091    }
1092}
llm_optimizer_processor/pipeline/operator.rs

llm_optimizer_processor/pipeline/
operator.rs
