laminar_sql/datafusion/
watermark_udf.rs

1//! Watermark UDF for `DataFusion` integration
2//!
3//! Provides a `watermark()` scalar function that returns the current
4//! watermark timestamp from Ring 0 via a shared `Arc<AtomicI64>`.
5//!
6//! The watermark represents the boundary below which all events are
7//! assumed to have arrived. Queries can use `WHERE event_time > watermark()`
8//! to filter stale data.
9
10use std::hash::{Hash, Hasher};
11use std::sync::atomic::{AtomicI64, Ordering};
12use std::sync::Arc;
13
14use arrow::datatypes::{DataType, TimeUnit};
15use datafusion_common::{Result, ScalarValue};
16use datafusion_expr::{
17    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
18};
19
20/// No-watermark sentinel value. When the atomic holds this value,
21/// `watermark()` returns `NULL`.
22pub const NO_WATERMARK: i64 = -1;
23
24/// Scalar UDF that returns the current watermark timestamp.
25///
26/// The watermark is stored in a shared `Arc<AtomicI64>` that Ring 0
27/// updates as events are processed. This UDF reads it with relaxed
28/// ordering (appropriate since watermarks are monotonically advancing
29/// and a slightly stale read is acceptable).
30///
31/// Returns `NULL` when no watermark has been set (value < 0).
32#[derive(Debug)]
33pub struct WatermarkUdf {
34    signature: Signature,
35    watermark_ms: Arc<AtomicI64>,
36}
37
38impl WatermarkUdf {
39    /// Creates a new watermark UDF backed by the given atomic value.
40    ///
41    /// # Arguments
42    ///
43    /// * `watermark_ms` - Shared atomic holding the current watermark
44    ///   in milliseconds since epoch. Values < 0 mean "no watermark".
45    pub fn new(watermark_ms: Arc<AtomicI64>) -> Self {
46        Self {
47            signature: Signature::new(TypeSignature::Nullary, Volatility::Volatile),
48            watermark_ms,
49        }
50    }
51
52    /// Creates a watermark UDF that always returns `NULL`.
53    #[must_use]
54    pub fn unset() -> Self {
55        Self::new(Arc::new(AtomicI64::new(NO_WATERMARK)))
56    }
57
58    /// Returns a reference to the underlying atomic watermark value.
59    #[must_use]
60    pub fn watermark_ref(&self) -> &Arc<AtomicI64> {
61        &self.watermark_ms
62    }
63}
64
65impl PartialEq for WatermarkUdf {
66    fn eq(&self, _other: &Self) -> bool {
67        true // All watermark UDF instances serve the same purpose
68    }
69}
70
71impl Eq for WatermarkUdf {}
72
73impl Hash for WatermarkUdf {
74    fn hash<H: Hasher>(&self, state: &mut H) {
75        "watermark".hash(state);
76    }
77}
78
79impl ScalarUDFImpl for WatermarkUdf {
80    fn as_any(&self) -> &dyn std::any::Any {
81        self
82    }
83
84    fn name(&self) -> &'static str {
85        "watermark"
86    }
87
88    fn signature(&self) -> &Signature {
89        &self.signature
90    }
91
92    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
93        Ok(DataType::Timestamp(TimeUnit::Millisecond, None))
94    }
95
96    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
97        let wm = self.watermark_ms.load(Ordering::Relaxed);
98        if wm < 0 {
99            Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
100                None, None,
101            )))
102        } else {
103            Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
104                Some(wm),
105                None,
106            )))
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use arrow::datatypes::TimeUnit;
115    use arrow_schema::Field;
116    use datafusion_common::config::ConfigOptions;
117    use datafusion_expr::ScalarUDF;
118
119    fn make_args() -> ScalarFunctionArgs {
120        ScalarFunctionArgs {
121            args: vec![],
122            arg_fields: vec![],
123            number_rows: 1,
124            return_field: Arc::new(Field::new(
125                "output",
126                DataType::Timestamp(TimeUnit::Millisecond, None),
127                true,
128            )),
129            config_options: Arc::new(ConfigOptions::default()),
130        }
131    }
132
133    #[test]
134    fn test_watermark_default_null() {
135        let udf = WatermarkUdf::unset();
136        let result = udf.invoke_with_args(make_args()).unwrap();
137        match result {
138            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(None, _)) => {}
139            other => panic!("Expected NULL when no watermark set, got: {other:?}"),
140        }
141    }
142
143    #[test]
144    fn test_watermark_fixed_value() {
145        let wm = Arc::new(AtomicI64::new(1_000_000));
146        let udf = WatermarkUdf::new(Arc::clone(&wm));
147
148        let result = udf.invoke_with_args(make_args()).unwrap();
149        match result {
150            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(Some(v), _)) => {
151                assert_eq!(v, 1_000_000);
152            }
153            other => panic!("Expected TimestampMillisecond(1_000_000), got: {other:?}"),
154        }
155    }
156
157    #[test]
158    fn test_watermark_atomic_update() {
159        let wm = Arc::new(AtomicI64::new(NO_WATERMARK));
160        let udf = WatermarkUdf::new(Arc::clone(&wm));
161
162        // Initially NULL
163        let result = udf.invoke_with_args(make_args()).unwrap();
164        assert!(matches!(
165            result,
166            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(None, _))
167        ));
168
169        // Update watermark
170        wm.store(500_000, Ordering::Relaxed);
171        let result = udf.invoke_with_args(make_args()).unwrap();
172        match result {
173            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(Some(v), _)) => {
174                assert_eq!(v, 500_000);
175            }
176            other => panic!("Expected updated watermark, got: {other:?}"),
177        }
178    }
179
180    #[test]
181    fn test_watermark_registration() {
182        let udf = ScalarUDF::new_from_impl(WatermarkUdf::unset());
183        assert_eq!(udf.name(), "watermark");
184    }
185}
laminar_sql/datafusion/watermark_udf.rs

laminar_sql/datafusion/
watermark_udf.rs
