laminar_sql/datafusion/
proctime_udf.rs

1//! Processing-time UDF for `DataFusion` integration
2//!
3//! Provides a `proctime()` scalar function that returns the current
4//! wall-clock timestamp. Used to declare processing-time watermarks:
5//!
6//! ```sql
7//! CREATE SOURCE events (
8//!     data VARCHAR,
9//!     ts TIMESTAMP,
10//!     WATERMARK FOR ts AS PROCTIME()
11//! );
12//! ```
13
14use std::hash::{Hash, Hasher};
15use std::time::{SystemTime, UNIX_EPOCH};
16
17use arrow::datatypes::{DataType, TimeUnit};
18use datafusion_common::{Result, ScalarValue};
19use datafusion_expr::{
20    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
21};
22
23/// Scalar UDF that returns the current processing time.
24///
25/// `proctime()` takes no arguments and returns a `TimestampMillisecond`
26/// representing the current wall-clock time. Unlike `watermark()`, this
27/// function is volatile — it returns a different value on each invocation.
28#[derive(Debug)]
29pub struct ProcTimeUdf {
30    signature: Signature,
31}
32
33impl ProcTimeUdf {
34    /// Creates a new processing-time UDF.
35    #[must_use]
36    pub fn new() -> Self {
37        Self {
38            signature: Signature::new(TypeSignature::Nullary, Volatility::Volatile),
39        }
40    }
41}
42
43impl Default for ProcTimeUdf {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl PartialEq for ProcTimeUdf {
50    fn eq(&self, _other: &Self) -> bool {
51        true
52    }
53}
54
55impl Eq for ProcTimeUdf {}
56
57impl Hash for ProcTimeUdf {
58    fn hash<H: Hasher>(&self, state: &mut H) {
59        "proctime".hash(state);
60    }
61}
62
63impl ScalarUDFImpl for ProcTimeUdf {
64    fn as_any(&self) -> &dyn std::any::Any {
65        self
66    }
67
68    fn name(&self) -> &'static str {
69        "proctime"
70    }
71
72    fn signature(&self) -> &Signature {
73        &self.signature
74    }
75
76    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
77        Ok(DataType::Timestamp(TimeUnit::Millisecond, None))
78    }
79
80    #[allow(clippy::cast_possible_truncation)]
81    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
82        let now_ms = SystemTime::now()
83            .duration_since(UNIX_EPOCH)
84            .unwrap_or_default()
85            .as_millis() as i64;
86        Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
87            Some(now_ms),
88            None,
89        )))
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use std::sync::Arc;
97
98    use arrow_schema::Field;
99    use datafusion_common::config::ConfigOptions;
100    use datafusion_expr::ScalarUDF;
101
102    fn make_args() -> ScalarFunctionArgs {
103        ScalarFunctionArgs {
104            args: vec![],
105            arg_fields: vec![],
106            number_rows: 1,
107            return_field: Arc::new(Field::new(
108                "output",
109                DataType::Timestamp(TimeUnit::Millisecond, None),
110                true,
111            )),
112            config_options: Arc::new(ConfigOptions::default()),
113        }
114    }
115
116    #[test]
117    fn test_proctime_returns_timestamp() {
118        let udf = ProcTimeUdf::new();
119        let result = udf.invoke_with_args(make_args()).unwrap();
120        match result {
121            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(Some(v), _)) => {
122                // Should be a reasonable timestamp (after 2020-01-01)
123                assert!(v > 1_577_836_800_000, "timestamp too old: {v}");
124            }
125            other => panic!("Expected TimestampMillisecond, got: {other:?}"),
126        }
127    }
128
129    #[test]
130    fn test_proctime_registration() {
131        let udf = ScalarUDF::new_from_impl(ProcTimeUdf::new());
132        assert_eq!(udf.name(), "proctime");
133    }
134
135    #[test]
136    fn test_proctime_volatile() {
137        let udf = ProcTimeUdf::new();
138        assert_eq!(udf.signature().volatility, Volatility::Volatile);
139    }
140
141    #[test]
142    fn test_proctime_return_type() {
143        let udf = ProcTimeUdf::new();
144        let rt = udf.return_type(&[]).unwrap();
145        assert_eq!(rt, DataType::Timestamp(TimeUnit::Millisecond, None));
146    }
147}
laminar_sql/datafusion/proctime_udf.rs

laminar_sql/datafusion/
proctime_udf.rs
