laminar_sql/datafusion/format_bridge_udf/
mod.rs

1//! Format bridge scalar UDFs (F-SCHEMA-014).
2//!
3//! SQL-callable functions for inline format conversion within queries:
4//!
5//! - [`ParseEpochUdf`] — `parse_epoch(number, unit) -> timestamp`
6//! - [`ParseTimestampUdf`] — `parse_timestamp(string, format, timezone) -> timestamp`
7//! - [`ToJsonUdf`] — `to_json(value) -> text`
8//! - [`FromJsonUdf`] — `from_json(string) -> jsonb`
9
10use std::hash::{Hash, Hasher};
11use std::sync::Arc;
12
13use arrow::datatypes::{DataType, TimeUnit};
14use arrow_array::{
15    builder::{LargeBinaryBuilder, StringBuilder, TimestampMicrosecondBuilder},
16    Array, ArrayRef, LargeBinaryArray, StringArray, TimestampMicrosecondArray,
17};
18use datafusion_common::Result;
19use datafusion_expr::{
20    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
21};
22
23use super::json_types;
24
25// ── Helpers ──────────────────────────────────────────────────────
26
27/// Expand all args to arrays of the same length.
28fn expand_args(args: &[ColumnarValue]) -> Result<Vec<ArrayRef>> {
29    let len = args
30        .iter()
31        .find_map(|a| match a {
32            ColumnarValue::Array(arr) => Some(arr.len()),
33            ColumnarValue::Scalar(_) => None,
34        })
35        .unwrap_or(1);
36
37    args.iter()
38        .map(|a| match a {
39            ColumnarValue::Array(arr) => Ok(Arc::clone(arr)),
40            ColumnarValue::Scalar(s) => s.to_array_of_size(len),
41        })
42        .collect()
43}
44
45// ══════════════════════════════════════════════════════════════════
46// parse_epoch(number, unit) -> timestamp
47// ══════════════════════════════════════════════════════════════════
48
49/// `parse_epoch(number, unit) -> timestamp`
50///
51/// Converts an epoch number to a timestamp. The unit parameter specifies
52/// the time unit: 'seconds', 'milliseconds', 'microseconds', 'nanoseconds'.
53///
54/// Returns `TimestampMicrosecond` (microsecond precision, no timezone).
55/// Executes in Ring 0 — pure arithmetic, no allocation.
56///
57/// # Examples
58///
59/// ```sql
60/// SELECT parse_epoch(1708528800, 'seconds') AS ts;
61/// SELECT parse_epoch(1708528800000, 'milliseconds') AS ts;
62/// ```
63#[derive(Debug)]
64pub struct ParseEpochUdf {
65    signature: Signature,
66}
67
68impl ParseEpochUdf {
69    /// Creates a new `parse_epoch` UDF.
70    #[must_use]
71    pub fn new() -> Self {
72        Self {
73            signature: Signature::new(
74                TypeSignature::Exact(vec![DataType::Int64, DataType::Utf8]),
75                Volatility::Immutable,
76            ),
77        }
78    }
79}
80
81impl Default for ParseEpochUdf {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl PartialEq for ParseEpochUdf {
88    fn eq(&self, _other: &Self) -> bool {
89        true
90    }
91}
92
93impl Eq for ParseEpochUdf {}
94
95impl Hash for ParseEpochUdf {
96    fn hash<H: Hasher>(&self, state: &mut H) {
97        "parse_epoch".hash(state);
98    }
99}
100
101impl ScalarUDFImpl for ParseEpochUdf {
102    fn as_any(&self) -> &dyn std::any::Any {
103        self
104    }
105
106    fn name(&self) -> &'static str {
107        "parse_epoch"
108    }
109
110    fn signature(&self) -> &Signature {
111        &self.signature
112    }
113
114    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
115        Ok(DataType::Timestamp(TimeUnit::Microsecond, None))
116    }
117
118    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
119        let expanded = expand_args(&args.args)?;
120        let val_arr = expanded[0]
121            .as_any()
122            .downcast_ref::<arrow_array::Int64Array>()
123            .ok_or_else(|| {
124                datafusion_common::DataFusionError::Internal(
125                    "parse_epoch: first arg must be Int64".into(),
126                )
127            })?;
128        let unit_arr = expanded[1]
129            .as_any()
130            .downcast_ref::<StringArray>()
131            .ok_or_else(|| {
132                datafusion_common::DataFusionError::Internal(
133                    "parse_epoch: second arg must be Utf8".into(),
134                )
135            })?;
136
137        let mut builder = TimestampMicrosecondBuilder::with_capacity(val_arr.len());
138        for i in 0..val_arr.len() {
139            if val_arr.is_null(i) || unit_arr.is_null(i) {
140                builder.append_null();
141            } else {
142                let value = val_arr.value(i);
143                let unit = unit_arr.value(i);
144                let micros = epoch_to_micros(value, unit)?;
145                builder.append_value(micros);
146            }
147        }
148        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
149    }
150}
151
152/// Convert an epoch value with the given unit string to microseconds.
153fn epoch_to_micros(value: i64, unit: &str) -> Result<i64> {
154    match unit.to_ascii_lowercase().as_str() {
155        "seconds" | "s" => Ok(value.saturating_mul(1_000_000)),
156        "milliseconds" | "ms" => Ok(value.saturating_mul(1_000)),
157        "microseconds" | "us" => Ok(value),
158        "nanoseconds" | "ns" => Ok(value / 1_000),
159        _ => Err(datafusion_common::DataFusionError::Execution(format!(
160            "parse_epoch: invalid unit '{unit}'. \
161             Expected: seconds, milliseconds, microseconds, nanoseconds"
162        ))),
163    }
164}
165
166// ══════════════════════════════════════════════════════════════════
167// parse_timestamp(string, format, timezone) -> timestamp
168// ══════════════════════════════════════════════════════════════════
169
170/// `parse_timestamp(string, format) -> timestamp`
171///
172/// Parses a timestamp string using the given chrono format string.
173/// Use `'iso8601'` as a shortcut for ISO 8601 parsing.
174///
175/// Returns `TimestampMicrosecond`.
176///
177/// # Examples
178///
179/// ```sql
180/// SELECT parse_timestamp('2026-02-21 14:30:00', '%Y-%m-%d %H:%M:%S');
181/// SELECT parse_timestamp('2026-02-21T14:30:00Z', 'iso8601');
182/// ```
183#[derive(Debug)]
184pub struct ParseTimestampUdf {
185    signature: Signature,
186}
187
188impl ParseTimestampUdf {
189    /// Creates a new `parse_timestamp` UDF.
190    #[must_use]
191    pub fn new() -> Self {
192        Self {
193            signature: Signature::new(
194                TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]),
195                Volatility::Immutable,
196            ),
197        }
198    }
199}
200
201impl Default for ParseTimestampUdf {
202    fn default() -> Self {
203        Self::new()
204    }
205}
206
207impl PartialEq for ParseTimestampUdf {
208    fn eq(&self, _other: &Self) -> bool {
209        true
210    }
211}
212
213impl Eq for ParseTimestampUdf {}
214
215impl Hash for ParseTimestampUdf {
216    fn hash<H: Hasher>(&self, state: &mut H) {
217        "parse_timestamp".hash(state);
218    }
219}
220
221impl ScalarUDFImpl for ParseTimestampUdf {
222    fn as_any(&self) -> &dyn std::any::Any {
223        self
224    }
225
226    fn name(&self) -> &'static str {
227        "parse_timestamp"
228    }
229
230    fn signature(&self) -> &Signature {
231        &self.signature
232    }
233
234    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
235        Ok(DataType::Timestamp(TimeUnit::Microsecond, None))
236    }
237
238    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
239        let expanded = expand_args(&args.args)?;
240        let str_arr = expanded[0]
241            .as_any()
242            .downcast_ref::<StringArray>()
243            .ok_or_else(|| {
244                datafusion_common::DataFusionError::Internal(
245                    "parse_timestamp: first arg must be Utf8".into(),
246                )
247            })?;
248        let fmt_arr = expanded[1]
249            .as_any()
250            .downcast_ref::<StringArray>()
251            .ok_or_else(|| {
252                datafusion_common::DataFusionError::Internal(
253                    "parse_timestamp: second arg must be Utf8".into(),
254                )
255            })?;
256
257        let mut builder = TimestampMicrosecondBuilder::with_capacity(str_arr.len());
258        for i in 0..str_arr.len() {
259            if str_arr.is_null(i) || fmt_arr.is_null(i) {
260                builder.append_null();
261            } else {
262                let ts_str = str_arr.value(i);
263                let fmt = fmt_arr.value(i);
264                match parse_ts_string(ts_str, fmt) {
265                    Ok(micros) => builder.append_value(micros),
266                    Err(_) => builder.append_null(),
267                }
268            }
269        }
270        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
271    }
272}
273
274/// Parse a timestamp string with the given format to microseconds since epoch.
275fn parse_ts_string(ts_str: &str, fmt: &str) -> std::result::Result<i64, String> {
276    use chrono::NaiveDateTime;
277
278    if fmt.eq_ignore_ascii_case("iso8601") {
279        // Try RFC 3339 / ISO 8601
280        let dt = ts_str
281            .parse::<chrono::DateTime<chrono::Utc>>()
282            .or_else(|_| {
283                NaiveDateTime::parse_from_str(ts_str, "%Y-%m-%dT%H:%M:%S%.f")
284                    .map(|ndt| ndt.and_utc())
285            })
286            .or_else(|_| {
287                NaiveDateTime::parse_from_str(ts_str, "%Y-%m-%dT%H:%M:%S").map(|ndt| ndt.and_utc())
288            })
289            .map_err(|e| e.to_string())?;
290        return Ok(dt.timestamp_micros());
291    }
292
293    let ndt = NaiveDateTime::parse_from_str(ts_str, fmt).map_err(|e| e.to_string())?;
294    Ok(ndt.and_utc().timestamp_micros())
295}
296
297// ══════════════════════════════════════════════════════════════════
298// to_json(value) -> text
299// ══════════════════════════════════════════════════════════════════
300
301/// `to_json(value) -> text`
302///
303/// Converts any SQL value to a JSON text string.
304/// Works with all Arrow data types — strings, numbers, booleans, structs.
305///
306/// # Examples
307///
308/// ```sql
309/// SELECT to_json(42);           -- Returns: '42'
310/// SELECT to_json('hello');      -- Returns: '"hello"'
311/// ```
312#[derive(Debug)]
313pub struct ToJsonUdf {
314    signature: Signature,
315}
316
317impl ToJsonUdf {
318    /// Creates a new `to_json` UDF.
319    #[must_use]
320    pub fn new() -> Self {
321        Self {
322            signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
323        }
324    }
325}
326
327impl Default for ToJsonUdf {
328    fn default() -> Self {
329        Self::new()
330    }
331}
332
333impl PartialEq for ToJsonUdf {
334    fn eq(&self, _other: &Self) -> bool {
335        true
336    }
337}
338
339impl Eq for ToJsonUdf {}
340
341impl Hash for ToJsonUdf {
342    fn hash<H: Hasher>(&self, state: &mut H) {
343        "to_json".hash(state);
344    }
345}
346
347impl ScalarUDFImpl for ToJsonUdf {
348    fn as_any(&self) -> &dyn std::any::Any {
349        self
350    }
351
352    fn name(&self) -> &'static str {
353        "to_json"
354    }
355
356    fn signature(&self) -> &Signature {
357        &self.signature
358    }
359
360    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
361        Ok(DataType::Utf8)
362    }
363
364    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
365        let expanded = expand_args(&args.args)?;
366        let arr = &expanded[0];
367        let len = arr.len();
368
369        let mut builder = StringBuilder::with_capacity(len, 256);
370        for row in 0..len {
371            if arr.is_null(row) {
372                builder.append_value("null");
373            } else {
374                let val = arrow_value_to_json(arr, row);
375                builder.append_value(val.to_string());
376            }
377        }
378        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
379    }
380}
381
382/// Convert an Arrow array value at a given row to a `serde_json::Value`.
383fn arrow_value_to_json(arr: &ArrayRef, row: usize) -> serde_json::Value {
384    if arr.is_null(row) {
385        return serde_json::Value::Null;
386    }
387    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
388        return serde_json::Value::String(a.value(row).to_owned());
389    }
390    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Int64Array>() {
391        return serde_json::Value::Number(a.value(row).into());
392    }
393    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Int32Array>() {
394        return serde_json::Value::Number(i64::from(a.value(row)).into());
395    }
396    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::Float64Array>() {
397        if let Some(n) = serde_json::Number::from_f64(a.value(row)) {
398            return serde_json::Value::Number(n);
399        }
400        return serde_json::Value::Null;
401    }
402    if let Some(a) = arr.as_any().downcast_ref::<arrow_array::BooleanArray>() {
403        return serde_json::Value::Bool(a.value(row));
404    }
405    // JSONB passthrough
406    if let Some(a) = arr.as_any().downcast_ref::<LargeBinaryArray>() {
407        if let Some(text) = json_types::jsonb_to_text(a.value(row)) {
408            if let Ok(val) = serde_json::from_str::<serde_json::Value>(&text) {
409                return val;
410            }
411            return serde_json::Value::String(text);
412        }
413        return serde_json::Value::Null;
414    }
415    // Timestamp types
416    if let Some(a) = arr.as_any().downcast_ref::<TimestampMicrosecondArray>() {
417        let micros = a.value(row);
418        let secs = micros / 1_000_000;
419        let sub_micros = (micros % 1_000_000).unsigned_abs() * 1_000;
420        #[allow(clippy::cast_possible_truncation)]
421        let nsecs = sub_micros as u32;
422        if let Some(dt) = chrono::DateTime::from_timestamp(secs, nsecs) {
423            return serde_json::Value::String(dt.to_rfc3339());
424        }
425        return serde_json::Value::Number(micros.into());
426    }
427    // Fallback via ScalarValue display
428    let sv = datafusion_common::ScalarValue::try_from_array(arr, row).ok();
429    match sv {
430        Some(s) => serde_json::Value::String(s.to_string()),
431        None => serde_json::Value::Null,
432    }
433}
434
435// ══════════════════════════════════════════════════════════════════
436// from_json(string) -> jsonb
437// ══════════════════════════════════════════════════════════════════
438
439/// `from_json(string) -> jsonb`
440///
441/// Parses a JSON string and returns the JSONB binary representation.
442/// Returns NULL for invalid JSON.
443///
444/// # Examples
445///
446/// ```sql
447/// SELECT from_json('{"name": "Alice", "age": 30}');
448/// SELECT json_typeof(from_json('42'));  -- Returns: 'number'
449/// ```
450#[derive(Debug)]
451pub struct FromJsonUdf {
452    signature: Signature,
453}
454
455impl FromJsonUdf {
456    /// Creates a new `from_json` UDF.
457    #[must_use]
458    pub fn new() -> Self {
459        Self {
460            signature: Signature::new(
461                TypeSignature::Exact(vec![DataType::Utf8]),
462                Volatility::Immutable,
463            ),
464        }
465    }
466}
467
468impl Default for FromJsonUdf {
469    fn default() -> Self {
470        Self::new()
471    }
472}
473
474impl PartialEq for FromJsonUdf {
475    fn eq(&self, _other: &Self) -> bool {
476        true
477    }
478}
479
480impl Eq for FromJsonUdf {}
481
482impl Hash for FromJsonUdf {
483    fn hash<H: Hasher>(&self, state: &mut H) {
484        "from_json".hash(state);
485    }
486}
487
488impl ScalarUDFImpl for FromJsonUdf {
489    fn as_any(&self) -> &dyn std::any::Any {
490        self
491    }
492
493    fn name(&self) -> &'static str {
494        "from_json"
495    }
496
497    fn signature(&self) -> &Signature {
498        &self.signature
499    }
500
501    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
502        Ok(DataType::LargeBinary) // Returns JSONB
503    }
504
505    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
506        let expanded = expand_args(&args.args)?;
507        let str_arr = expanded[0]
508            .as_any()
509            .downcast_ref::<StringArray>()
510            .ok_or_else(|| {
511                datafusion_common::DataFusionError::Internal("from_json: arg must be Utf8".into())
512            })?;
513
514        let mut builder = LargeBinaryBuilder::with_capacity(str_arr.len(), 256);
515        for i in 0..str_arr.len() {
516            if str_arr.is_null(i) {
517                builder.append_null();
518            } else {
519                let json_str = str_arr.value(i);
520                match serde_json::from_str::<serde_json::Value>(json_str) {
521                    Ok(val) => {
522                        let jsonb = json_types::encode_jsonb(&val);
523                        builder.append_value(&jsonb);
524                    }
525                    Err(_) => builder.append_null(),
526                }
527            }
528        }
529        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
530    }
531}
532
533// ══════════════════════════════════════════════════════════════════
534// Tests
535// ══════════════════════════════════════════════════════════════════
536
537#[cfg(test)]
538mod tests;
laminar_sql/datafusion/format_bridge_udf/mod.rs

laminar_sql/datafusion/format_bridge_udf/
mod.rs
