tinyflow_framework/stream/
data_stream.rs

1//! Linear operator chain DSL (`map`, `filter`, `flat_map`, `sink`, `name`).
2//!
3//! Build-time errors on the chain **panic** (duplicate source, steps after terminal `sink`).
4//! Runtime errors use [`crate::api::error::StreamResult`] in [`crate::stream::chain_execute::execute`].
5//!
6//! Terminate with [`DataStream::sink`], [`crate::connectors::print::PrintDataStreamExt::print`], or
7//! [`crate::connectors::iceberg::data_stream_ext::IcebergDataStreamExt::add_iceberg_sink`],
8//! then [`crate::env::stream_env::StreamEnv::execute`].
9
10use std::any::TypeId;
11use std::marker::PhantomData;
12
13use crate::env::stream_env::StreamEnv;
14use crate::runtime::chain_job::{ChainOperatorKind, ChainStep};
15use crate::runtime::chain_segment::ChainSegmentFactory;
16use tinyflow_api::functions::*;
17
18pub struct DataStream<T> {
19    pub(crate) env: StreamEnv,
20    pub(crate) tail_step_index: usize,
21    pub(crate) _marker: PhantomData<T>,
22}
23
24impl<T: 'static> DataStream<T> {
25    pub(crate) fn new(env: StreamEnv, tail_step_index: usize) -> Self {
26        Self {
27            env,
28            tail_step_index,
29            _marker: PhantomData,
30        }
31    }
32
33    /// Sets the user-visible name for the current tail operator.
34    pub fn name(self, n: impl Into<String>) -> Self {
35        let name = n.into();
36        let tail = self.tail_step_index;
37        {
38            let mut job = self.env.chain_job.lock().unwrap();
39            if let Some(job) = job.as_mut()
40                && tail < job.steps.len()
41            {
42                job.steps[tail].user_name = Some(name);
43            }
44        }
45        self
46    }
47
48    fn push_step<OUT: Send + 'static>(
49        self,
50        kind: ChainOperatorKind,
51        factory: ChainSegmentFactory,
52    ) -> DataStream<OUT> {
53        let mut job_guard = self.env.chain_job.lock().unwrap();
54        let job = job_guard
55            .as_mut()
56            .expect("ChainJob missing; call add_source first");
57        job.push_step_with_types(
58            ChainStep {
59                kind,
60                user_name: None,
61                factory,
62            },
63            Some(TypeId::of::<T>()),
64            TypeId::of::<OUT>(),
65        );
66        let idx = job.steps.len() - 1;
67        DataStream::new(self.env.clone(), idx)
68    }
69}
70
71impl<T: Send + Clone + 'static> DataStream<T> {
72    pub fn map<O: Send + Clone + 'static, F: MapFunction<T, O> + Clone + Send + 'static>(
73        self,
74        f: F,
75    ) -> DataStream<O> {
76        let factory = crate::runtime::chain_segment::map_chain_factory(f);
77        self.push_step(ChainOperatorKind::Map, factory)
78    }
79
80    pub fn filter<F: FilterFunction<T> + Clone + Send + 'static>(self, f: F) -> DataStream<T> {
81        let factory = crate::runtime::chain_segment::filter_chain_factory(f);
82        self.push_step(ChainOperatorKind::Filter, factory)
83    }
84
85    pub fn flat_map<
86        O: Send + Clone + 'static,
87        F: FlatMapFunction<T, O> + Clone + Send + 'static,
88    >(
89        self,
90        f: F,
91    ) -> DataStream<O> {
92        let factory = crate::runtime::chain_segment::flat_map_chain_factory(f);
93        self.push_step(ChainOperatorKind::FlatMap, factory)
94    }
95
96    /// Terminal sink; freezes the chain (no further operators).
97    pub fn sink<F>(self, f: F)
98    where
99        F: SinkFunction<T> + Clone + Send + 'static,
100    {
101        self.sink_named_impl(None, f);
102    }
103
104    /// Terminal sink with a user-visible operator name (for logs and `RuntimeContext::task_name`).
105    pub fn sink_named<F>(self, name: impl Into<String>, f: F)
106    where
107        F: SinkFunction<T> + Clone + Send + 'static,
108    {
109        self.sink_named_impl(Some(name.into()), f);
110    }
111
112    fn sink_named_impl<F>(self, user_name: Option<String>, f: F)
113    where
114        F: SinkFunction<T> + Clone + Send + 'static,
115    {
116        let factory = crate::runtime::chain_segment::sink_chain_factory(f);
117        let mut job_guard = self.env.chain_job.lock().unwrap();
118        let job = job_guard
119            .as_mut()
120            .expect("ChainJob missing; call add_source first");
121        job.push_step_with_types(
122            ChainStep {
123                kind: ChainOperatorKind::Sink,
124                user_name,
125                factory,
126            },
127            Some(TypeId::of::<T>()),
128            TypeId::of::<T>(),
129        );
130        drop(job_guard);
131    }
132}
tinyflow_framework/stream/data_stream.rs

tinyflow_framework/stream/
data_stream.rs
